"""
Purpose of the RSA cryptosystem is to have a secure way of transmitting data

Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
*   2019-08-26 Initial programming

"""

from math import gcd
from sympy import isprime
import random
from euclid_gcd import extended_euclidean


def generate_pq(bits):
    # Randomly generate two primes p,q
    p = random.getrandbits(bits)
    q = random.getrandbits(bits)

    # Check if p,q is prime
    p_isprime = isprime(p)
    q_isprime = isprime(q)

    # Keep generating until both are primes
    while not (p_isprime and q_isprime):
        if not p_isprime:
            p = random.getrandbits(bits)
        if not q_isprime:
            q = random.getrandbits(bits)

        p_isprime = isprime(p)
        q_isprime = isprime(q)

    return p, q


def generate_e(totient):
    # Generate e such that 1 < e < phi(n)
    # phi(n) in this case is totient

    while True:
        # Should be (2,totient) so if it is stuck in infinite loop then restart or replace 80000 -> totient
        # Reason why I want e to be a low value is to make encryption faster
        e = random.randint(2, 50000)

        if gcd(e, totient) == 1:
            return e


def generate_d(e, totient):
    _, e_inverse, _ = extended_euclidean(e, totient)
    d = e_inverse % totient

    return d


def generate_all_values():
    num_bits = 1024
    p, q = generate_pq(num_bits)
    totient = (p - 1) * (q - 1)
    e = generate_e(totient)
    d = generate_d(e, totient)

    print("Generated value n: " + str(p * q))
    print("Generated e and d: " + str(e) + " and " + str(d))

    return p * q, e, d


def encrypt(message, n, e):
    encrypted = ""
    for letter in message:
        pad = 3 - len(str(ord(letter)))

        if pad > 0:
            new_letter = "0" * pad + str(ord(letter))
        else:
            new_letter = ord(letter)

        encrypted += str(new_letter)

    encrypted = pow(int(encrypted), e, n)
    return encrypted


def decrypt(encrypted, n, d):
    decrypted_message = ""
    decrypted_code = str(pow(encrypted, d, n))

    if len(decrypted_code) % 3 != 0:
        decrypted_code = "0" + decrypted_code

    while len(decrypted_code):
        decrypted_message += chr(int(decrypted_code[0:3]))
        decrypted_code = decrypted_code[3:]

    return decrypted_message


def example():
    # An example of a test case where we generate all necessary primes, encrypt and then decrypt the message.
    # Only to show how all parts of the code is working. This is not how it's going to be used in practice.
    hidden_message = "i really love peanuts"
    n, e, d = generate_all_values()
    encrypted_message = encrypt(hidden_message, n, e)
    decrypted_message = decrypt(encrypted_message, n, d)

    print("\n")
    print("Original message: " + hidden_message)
    print("Encrypted message: " + str(encrypted_message))
    print("Decrypted message: " + decrypted_message)


def main():
    # Write the values of your RSA encryption (Note: Never give the 'd' to someone that doesn't want it)
    n = 354089397494626050014776605732143027269473328409397973403863001639624332101789181044818951483060155060788030618162673282176493895463414816015601230408140046833172059490430968956729878861381343553446553025440156523477822105773362480000716985478565013956749662865189691539813391686696182702224364834273144673717742246537383454469146642154754778836797926780437490677663302034284308892191362266103193070200405420180296005388479418941723827243187899338980201782128797489464650981164232057548015010630986959083998487019465357524040595865260220030689502065850060761344148196291328192760801074939658292752592564874822996765430361631210613041006858858506787439506504448316606509551260553919757840169593791152166515571202450662850988377002989153080277915454500432640601643512909764636398157415600050468972065216354878984114648007494687081718749734915103155014825420081658864982423629447913147575382146725524407739875786801876011026010419782863232303861065841801863420557617962438178979549855959377311548527613240676904989886382444381261628076009466895878852398923237601309285642954207693266358989851324012643315688180744573155217352955083176785543099571257338683938756048920161738393295253775030232399282686809347027784971441882883201432807953
    e = 9361
    # d =

    enc_or_dec = input(
        "Would you like to encrypt or decrypt a message (input: 'enc' or 'dec'): "
    )

    if enc_or_dec.lower() == "enc":
        hidden_message = input("What is your hidden message?: ")
        print("Encrypted message: " + str(encrypt(hidden_message, n, e)))

    elif enc_or_dec.lower() == "dec":
        encrypted = input("What is your encrypted message?: ")
        print(encrypted)
        print("Decrypted message: " + str(decrypt(int(encrypted), n, d)))

    else:
        print("Not sure what you typed")


main()
