import numpy as np
import math
import gmpy2
from gmpy2 import powmod,mpz,isqrt,invert
from qiskit.aqua.algorithms import Shor
from qiskit.aqua import QuantumInstance
from qiskit import Aer,execute,QuantumCircuit
from qiskit.tools.visualization import plot_histogram
from qiskit.providers.ibmq import least_busy
from qiskit import IBMQ, execute

def generate_keys():
    # prime number of 3 digits i.e 7 bits
    random1 = np.random.randint(3,40)
    random2 = np.random.randint(3,40)
    p = int(gmpy2.next_prime(random1))
    q = int(gmpy2.next_prime(random2))
    n = p*q
    while (n<100 or n>127):
        random1 = np.random.randint(3,40)
        random2 = np.random.randint(3,40)
        p = int(gmpy2.next_prime(random1))
        q = int(gmpy2.next_prime(random2))
        n = p*q
    phi = (p-1)*(q-1)
    e = 2
    while True:
        if gmpy2.gcd(phi,e) != 1:
            e = e + 1
        else :
            break
    d = gmpy2.invert(e,phi)
    return n,e,d

def encrypt(plain_text_blocks,public_keys):
    cipher_text_blocks = []
    n,e = public_keys
    for plain_text in plain_text_blocks:
        cipher_text = (gmpy2.powmod(plain_text,e,n))
        cipher_text_blocks.append(cipher_text)
    return cipher_text_blocks

def decrypt(cipher_text_blocks,secret_key,public_keys):
    n,e = public_keys
    d = secret_key
    decypted_plain_text_blocks = []
    for cipher_text in cipher_text_blocks:
        plain_text = (gmpy2.powmod(cipher_text,d,n))
        decypted_plain_text_blocks.append(plain_text)
    return decypted_plain_text_blocks

def get_factors(public_keys):
    n,e = public_keys
    # backend = Aer.get_backend('qasm_simulator')
    provider = IBMQ.load_account()
    backend = provider.get_backend('ibmq_qasm_simulator')
    quantum_instance = QuantumInstance(backend,shots=2500)
    find_factors = Shor(n,a=2,quantum_instance=quantum_instance)
    factors = Shor.run(find_factors)
    p = ((factors['factors'])[0])[0]
    q = ((factors['factors'])[0])[1]
    print('Factors of',n,'are :',p,q)
    
    return p,q

# taken in 'Hello World!!!' returns ['Hello World!','!!']
def get_blocks(PT,block_size):
    blocks = []
    i = 0
    while i<len(PT):
        temp_str=''
        if i+block_size-1 < len(PT):
            temp_str=temp_str+PT[i:i+block_size]
        else :
            temp_str=temp_str+PT[i::]
        blocks.append(temp_str)
        i=i+block_size
    return blocks
        
# covert plain_text block from characters to the numbers
def format_plain_text(PT):
    plain_text_blocks = []
    for block in PT:
        plain_text = 0
        for i in range(len(block)):
            # for 'd'
            if ord(block[i]) == 100:
                plain_text = plain_text*100 + 28
            # between (101,127)
            elif ord(block[i])>100:
                plain_text = plain_text*100 + (ord(block[i])-100)
            else :
                plain_text = plain_text*100 + (ord(block[i]))
        plain_text_blocks.append(plain_text)
    return plain_text_blocks

# convert numeric decypted_plain_text_blocks into a single plain text of characters
def format_decrypted_plain_text(decypted_plain_text_blocks):
    plain_text_blocks = []
    for dc_pt in decypted_plain_text_blocks:
        plain_text = ''
        temp = dc_pt
        # for 'd' temp = 28
        while temp > 0:
            if temp%100 == 28:
                plain_text = plain_text + 'd'
            elif (temp%100) in range(0,27):
                plain_text = plain_text + chr((temp%100)+100)
            else :
                plain_text = plain_text + chr((temp%100))
            temp = temp//100
        plain_text = plain_text[::-1] 
        plain_text_blocks.append(plain_text)
    final_plain_text = ''
    for plain_text_block in plain_text_blocks:
        final_plain_text = final_plain_text + plain_text_block
    return final_plain_text

n,e,d = generate_keys()

public_keys = (n,e)
secret_key = d
print("\nPublic Key :")
print('n :',n)
print('e :',e)
print("Secret Key :\nd :",d)

PT = input("\nEnter Plain Text to encrypt : ")

original_plain_text = PT

block_size = 1
PT = get_blocks(PT,block_size)
print('\nPlain Text after converting to blocks',PT)

plain_text_blocks = format_plain_text(PT)
print('\nPlain text blocks after formatting to numbers:',plain_text_blocks)

cipher_text_blocks = encrypt(plain_text_blocks,public_keys)
print("\nCipher Text Blocks After RSA encryption :",cipher_text_blocks)

p,q = get_factors(public_keys)
phi = (p-1)*(q-1)
broken_d = gmpy2.invert(e,phi)
compromised_PT = decrypt(cipher_text_blocks,broken_d,public_keys)
compromised_PT = format_decrypted_plain_text(compromised_PT)
compromised_PT = '!!!Your message has been attacked!!! ' + compromised_PT
compromised_PT = get_blocks(compromised_PT,block_size) 
compromised_PT = format_plain_text(compromised_PT)
compromised_CT = encrypt(compromised_PT,public_keys)

cipher_text_blocks = compromised_CT

decypted_plain_text_blocks = decrypt(cipher_text_blocks,secret_key,public_keys)
print("\nPlain Text blocks after decryption of Cipher Text blocks :",decypted_plain_text_blocks)

plain_text_after_decryption = format_decrypted_plain_text(decypted_plain_text_blocks)
print("\nAfter decryption Plain Text :",plain_text_after_decryption)

if (original_plain_text == plain_text_after_decryption):
    print("\nHurrayyy!!!\n\nDecrypted plain_text is same as original plain_text! :) ")
else :
    print('RSA was attacked!!! :(')

