# This method uses the 'pycrypto' package
# Import all the packages needed for this notebook
from Crypto.PublicKey import RSA
from Crypto import Random
import hashlib

# Generate an RSA of 2048 bits
random_generator = Random.new().read
private_key = RSA.generate(2048, random_generator)
public_key = private_key.publickey()

# Test the encryption and decrytion process
message = b'Super secret message'

# Encrypt with the public key
crypt_msg = public_key.encrypt(message, K=None)
print(crypt_msg)

# Now decrypt using the private key
decrypt_msg = private_key.decrypt(crypt_msg)
print(decrypt_msg)

"""
Explanation: Generate public and private keys using RSA of 2048 bits for the key size
Possible python libraries to use:
- pycrypto
- rsa
- ...
End of explanation
"""


# Print the public and private keys to console in PEM format  

# Print private key
print(private_key.exportKey(format='PEM'))

print('\n')

# Print public key
print(public_key.exportKey(format='PEM'))

"""
Explanation: Read and write to PEM formatted text
End of explanation
"""


"""
Similate a basic transaction
1.  Create digital wallet objects for Alice and Bob.
2.  Create private and digital keys for these wallets
3.  Send Bob's public key in 'PEM' format to Alice
4.  Alice then encrypt's a message using Bob's public key
5.  Alice sends encrypted message to Bob
6.  Bob receives message and decrypts message using his private key
"""

class DigitalWallet(object):
    def __init__(self, name):
                
        # Owner of the wallet
        self._name = name
        
        # Generate RSA keys on object creation
        private_key, public_key = self.__generate_RSA_keys(2048)
        
        # Objects RSA keys
        self._private_key = private_key # This should never be sent out
        self._public_keys = {self._name: public_key} 
        
    def __generate_RSA_keys(self, num_bits):
        # Keys that are generated when the digital wallet is created
        random_generator = Random.new().read
        private_key = RSA.generate(2048, random_generator)
        public_key = private_key.publickey()
        return private_key, public_key
    
    def send_public_keys(self):
        # Public key can be sent to others for message encryption
        return {self._name: self._public_keys[self._name].exportKey(format='PEM')}
    
    def receive_public_key(self, key_dict):
        # Recieve a public key from another person
        for key, val in key_dict.items():
            public_key = RSA.importKey(val)
            self._public_keys[key] = public_key
        
    def send_encrypted_message(self, to_user, message):
        # Encrypt with the public key and return to user
        return self._public_keys[to_user].encrypt(message, K=None)
    
    def receive_encrypted_message(self, encrypted_message):
        decoded_message = self._private_key.decrypt(encrypted_message)
        print(decoded_message)
        
    
# Build wallets for Alice and Bob
alice_wallet = DigitalWallet(name='Alice')
bob_wallet = DigitalWallet(name='Bob')

# Receive a public key from Bob
bob_public_key = bob_wallet.send_public_keys()
alice_wallet.receive_public_key(bob_public_key)

# Send an encrypted message to Bob that only Bob can decrypt
bob_encrypted_message = alice_wallet.send_encrypted_message('Bob', b'Secret message for bob')

# Receive the encrypted message
bob_wallet.receive_encrypted_message(bob_encrypted_message)

"""
Explanation: Build typical transaction object
A transaction should have the following properties:
1.  The owners public key and private key (Alice's public key, for example.
2.  Another users public key, so that they can send encrypted messages
End of explanation
"""


import hashlib

### An example of creating a password, storing this hashed value and validating future logins

# Create a new password, hash it and store it in a database 
new_password = b'Super_secret_password'
hashed_new_password = hashlib.sha256(new_password) 
hex_new_password = hashed_new_password.hexdigest() # Export this as a hexadecimal value (store this in a database)
print('New password hex: \t{0}'.format(hex_new_password))

# Now a user wants to log in using this password
# We validate whether the password is correct by hashing it and comparing it against the value in the database
# First try an incorrect password
bad_password = b'Wr0ngPassword'
hashed_bad_password = hashlib.sha256(bad_password)
hex_bad_password = hashed_bad_password.hexdigest() 
print('Bad password hex: \t{0}'.format(hex_bad_password))

# Now attempt to compare
if hex_new_password == hex_bad_password:
    print("Successfully logged in!\n")
else:
    print("Login failed!\n")

# Now do correct password validation    
correct_password = b'Super_secret_password'
hashed_correct_password = hashlib.sha256(correct_password)
hex_correct_password = hashed_correct_password.hexdigest() 
print('Correct password hex: \t{0}'.format(hex_correct_password))

# Now attempt to compare
if hex_new_password == hex_correct_password:
    print("Successfully logged in!\n")
else:
    print("Login failed!\n")

### This doesn't quite work just yet.

# You can also hash objects
hashed_alice_wallet = hashlib.sha256(bytes(alice_wallet))
print('Hashed version of alice_wallet: \t{0}'.format(hashed_alice_wallet.hexdigest()))
hashed_bob_wallet = hashlib.sha256(bytes(bob_wallet))
print('Hashed version of bob_wallet: \t\t{0}'.format(hashed_bob_wallet.hexdigest()))

# Create a new version of bob_wallet
new_bob_wallet = DigitalWallet(name='Bob')
hashed_new_bob_wallet = hashlib.sha256(bytes(new_bob_wallet))
print('Hashed version of new_bob_wallet: \t\t{0}'.format(hashed_new_bob_wallet.hexdigest()))

"""
Explanation: Examples of hashing
Use hashlib at: https://docs.python.org/3/library/hashlib.html
Thought I would add in some use cases of where you would use hashing
End of explanation
"""


