# coding=utf-8

#: ## Lab 9 ##
#: 
#: CS-2910 Network Protocols  
#: Dr. Yoder 
#: Fall quarter 2016-2017
#: 
#: | Team members (username) |
#: |:------------------------|
#: | Grace Hopper (hopperg)  |
#: | Donald Knuth (knuthd)   |
#: 
# If you are curious about the names above, look them up!

import random
import sys
import math

BITS_PER_HEX_DIGIT = 4 # binary digits per hex digit -- always 4
BIT_LENGTH = 28  # "B" in the lab handout. Length of n in bits
HEX_LENGTH = BIT_LENGTH/BITS_PER_HEX_DIGIT # Length of n in hexadecimal digits
MAX_PRIME = 0b11111111111111  # The maximum value a prime number can have
MIN_PRIME = 0b11110000000001  # The minimum value a prime number can have
PUBLIC_EXPONENT = 17  # The default public exponent


#
# Provide the user with a variety of encryption-related actions
#
def main():
    # Get chosen operation from the user.
    action = input("Select an option from the menu below:\n"
                   "(1-CK) create_keys\n"
                   "(2-CC) compute_checksum\n"
                   "(3-VC) verify_checksum\n"
                   "(4-EM) encrypt_message\n"
                   "(5-DM) decrypt_message\n"
                   "(6-BK) break_key\n "
                   "Please enter the option you want:\n")
    # Execute the chosen operation.
    if action in ['1', 'CK', 'ck', 'create_keys']:
        create_keys_interactive()
    elif action in ['2', 'CC', 'cc', 'compute_checksum']:
        compute_checksum_interactive()
    elif action in ['3', 'VC', 'vc', 'verify_checksum']:
        verify_checksum_interactive()
    elif action in ['4', 'EM', 'em', 'encrypt_message']:
        encrypt_message_interactive()
    elif action in ['5', 'DM', 'dm', 'decrypt_message']:
        decrypt_message_interactive()
    elif action in ['6', 'BK', 'bk', 'break_key']:
        break_key_interactive()
    else:
        print("Unknown action: '{0}'".format(action))


#
# Create new public keys
#
# Returns the private key for use by other interactive methods
#
def create_keys_interactive():
    key_pair = create_keys()
    pub = get_public_key(key_pair)
    priv = get_private_key(key_pair)
    print("Public key: ")
    print(pub)
    print("Private key: ")
    print(priv)
    return priv


#
# Compute the checksum for a message, and encrypt it
#  
def compute_checksum_interactive():
    priv = create_keys_interactive()

    message = input('Please enter the message to be checksummed: ')

    hash = compute_checksum(message)
    print('Hash:',as_hex(hash))
    cipher = apply_key(priv, hash)
    print('Encrypted Hash:', as_hex(cipher))


#
# Verify a message with its checksum, interactively
#
def verify_checksum_interactive():
    pub = enter_public_key_interactive()
    message = input('Please enter the message to be verified: ')
    recomputed_hash = compute_checksum(message)

    string_hash = input('Please enter the encrypted hash (in hexadecimal): ')
    encrypted_hash = int(string_hash, 16)
    decrypted_hash = apply_key(pub, encrypted_hash)
    print('Recomputed hash:', as_hex(format(recomputed_hash)))
    print('Decrypted hash: ', as_hex(format(decrypted_hash)))
    if recomputed_hash == decrypted_hash:
        print('Hashes match -- message is verified')
    else:
        print('Hashes do not match -- has tampering occured?')


#
# Encrypt a message
#
def encrypt_message_interactive():
    message = input('Please enter the message to be encrypted: ')
    pub = enter_public_key_interactive()
    encrypted = ''
    for c in message:
        encrypted += apply_key(pub,ord(c))
    print("Encrypted message:", encrypted)


#
# Decrypt a message
#
def decrypt_message_interactive(priv = None):
    encrypted = input('Please enter the message to be decrypted: ')
    if priv is None:
        priv = enter_key_interactive('private')
    message = ''
    for i in range(0,len(encrypted),HEX_LENGTH):
        enc_string = encrypted[i:i+HEX_LENGTH]
        enc = int(enc_string,16)
        dec = apply_key(priv,enc)
        if dec >= 0 and dec < 256:
            message += chr(dec)
        else:
            print('Warning: Could not decode encrypted entity: ' + enc_string)
            print('         decrypted as: ' + str(dec) + ' which is out of range.')
            print('         inserting _ at position of this character')
            message += '_'
    print("Decrypted message:", message)


#
# Break key, interactively
#
def break_key_interactive():
    pub = enter_public_key_interactive()
    priv = break_key(pub)
    print("Private key:")
    print(priv)
    decrypt_message_interactive(priv)


#
# Prompt user to enter the public modulus.
#
# returns the tuple (e,n)
#    
def enter_public_key_interactive():
    print('(Using public exponent = ' + str(PUBLIC_EXPONENT) + ')')
    string_modulus = input('Please enter the modulus (decimal): ')
    modulus = int(string_modulus)
    return (PUBLIC_EXPONENT, modulus)


#
# Prompt user to enter the exponent and modulus of a key
#
# key_type - either the string 'public' or 'private' -- used to prompt the user on how
#            this key is interpretted by the program.
#
# returns the tuple (e,n)
#    
def enter_key_interactive(key_type):
    string_exponent = input('Please enter the ' + key_type + ' exponent (decimal): ')
    exponent = int(string_exponent)
    string_modulus = input('Please enter the modulus (decimal): ')
    modulus = int(string_modulus)
    return (exponent, modulus)


#
# Compute simple hash
#
# Given a string, compute a simple hash as the sum of characters
# in the string.
#
# (If the sum goes over sixteen bits, the numbers should "wrap around"
# back into a sixteen bit number.  e.g. 0x3E6A7 should "wrap around" to
# 0xE6A7)
#
# This checksum is similar to the internet checksum used in UDP and TCP
# packets, but it is a two's complement sum rather than a one's
# complement sum.
#
# Returns the checksum as an integer
#
def compute_checksum(string):
    total = 0
    for c in string:
        total += ord(c)
    total %= 0x8000  # Guarantees checksum is only 4 hex digits
    # How many bytes is that?
    #
    # Also guarantees that that the checksum will
    # always be less than the modulus.
    return total


# ---------------------------------------
# Do not modify code above this line
# ---------------------------------------

#
# Create the public and private keys.
#
# Returns the keys as a three-tuple:
#
#  (e,d,n)
#
def create_keys():
    pass  # Delete this line and complete this method


#
# Apply the key, given as a tuple (e,n) or (d,n) to the message.
#
# This can be used both for encryption and decription.
#
# Returns the message with the key applied. For example,
# if given the public key and a message, encrypts the message
# and returns the ciphertext.
#
def apply_key(key, m):
    pass  # Delete this line and complete this method


#
# Break a key.  Given the public key, find the private key.
# Factorizes the modulus n to find the prime numbers p and q.
#
# You can follow the steps in the "optional" part of the in-class
# exercise.
#
# pub - a tuple containing the public key (e,n)
#
# returns a tuple containing the private key (d,n)
#
def break_key(pub):
    pass  # Delete this line and complete this method


# Your code and additional functions go here. (Replace this line.)

# ** Do not modify code below this line.

#
# Pulls the public key out of the tuple structure created by
# create_keys()
#
def get_public_key(key_pair):
    return (key_pair[0], key_pair[2])


#
# Pulls the private key out of the tuple structure created by
# create_keys()
#
def get_private_key(key_pair):
    return (key_pair[1], key_pair[2])


#
# Convert integer to a zero-padded hex string with the required number
# of characters to represent n, d, or and encrypted message.
# 
# Returns: The formatted string
#
def as_hex(number):
    return "{0:0"+str(HEX_LENGTH)+"x}".format(number)

main()