diff --git a/.gitignore b/.gitignore index 7e99e36..ed33af8 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ -*.pyc \ No newline at end of file +*.pyc +venv/ +.vscode/ \ No newline at end of file diff --git a/simpleCoin/wallet.py b/simpleCoin/wallet.py index 7eeb992..10b9778 100644 --- a/simpleCoin/wallet.py +++ b/simpleCoin/wallet.py @@ -22,6 +22,10 @@ import time import base64 import ecdsa +from cryptography.fernet import Fernet, InvalidToken +from hashlib import md5 +from json import loads, dumps +import os def wallet(): @@ -38,8 +42,33 @@ def wallet(): =========================================\n""") generate_ECDSA_keys() elif response == "2": - addr_from = input("From: introduce your wallet address (public key)\n") - private_key = input("Introduce your private key\n") + files_in_directory = []# All files with .wallet extension in the current directory + for file in os.listdir("./"): + if file.endswith(".wallet"): + files_in_directory.append(file)# Add all wallet files found + + if len(files_in_directory) == 0: + # If there's not any wallet files, the 'manual' process is used + addr_from = input("From: introduce your wallet address (public key)\n") + private_key = input("Introduce your private key\n") + elif len(files_in_directory) == 1: + wallet = get_data_from_file(files_in_directory[0]) + addr_from, private_key = wallet["public_key"], wallet["private_key"] + elif len(files_in_directory) > 1: + print("Various wallets detected:") + for i in files_in_directory: + print(i) + while True: + wallet_file = input("Which one do you want to use?: ") + try: + wallet = get_data_from_file(wallet_file) + break + except InvalidToken: + print("There was an error decrypting the private key.\nEither the password is not correct or the file is corrupted.") + except FileNotFoundError: + print("There was an error finding the file.") + + addr_from, private_key = wallet["public_key"], wallet["private_key"] addr_to = input("To: introduce destination wallet address\n") amount = input("Amount: number stating how much do you want to send\n") print("=========================================\n\n") @@ -51,6 +80,25 @@ def wallet(): else: # Will always occur when response == 3. check_transactions() +def get_data_from_file(wallet_file): + """Gets from a .wallet file the following wallet object: + + { + "public_key":"publicKey", + "private_key":"privateKey", + "encryption":True or False + }""" + with open(wallet_file, "r") as f: + content = loads(f.read())# Gets the content of the .wallet file + f.close() + if content["encryption"]:# If the private key is encrypted + password = input("The wallet file detected has been encrypted. Please, input the password: ") + key = gen_encryption_key(password)# Get the password key + f = Fernet(key)# Create the decryption object with the key + content["private_key"] = f.decrypt( content["private_key"].encode() ).decode() + else:# If not, just loads the json object + print("Not-encrypted wallet file detected.") + return content def send_transaction(addr_from, private_key, addr_to, amount): """Sends your transaction to different nodes. Once any of the nodes manage @@ -88,6 +136,14 @@ def check_transactions(): res = requests.get('http://localhost:5000/blocks') print(res.text) +def gen_encryption_key(password): + """A plain text password can't be used for encrpting or decrypting with Fernet. + This function returns a 32 bit url-safe base64 encoded key version of the password. + This key is valid for encrypting and decrypting.""" + h = md5() + h.update(password.encode()) + return base64.urlsafe_b64encode(h.hexdigest().encode()) + # Now it's ready to use. def generate_ECDSA_keys(): """This function takes care of creating your private and public (your address) keys. @@ -103,10 +159,32 @@ def generate_ECDSA_keys(): public_key = vk.to_string().hex() #we are going to encode the public key to make it shorter public_key = base64.b64encode(bytes.fromhex(public_key)) + + # The private and public key are assembled into a dictionary. + # Dumping and loading this object with json format will make + # them easier to manage. + content = {"public_key":public_key.decode(), "private_key":private_key, "encryption":False} + + filename = input("Write the name of your new address: ") + ".wallet" + password = input(""" +Write a password to encript the new wallet file. +Protecting your file with a password prevents your private key for being +stolen by someone that could have access to the wallet file. +Leave in blank to not use any password. (NOT RECOMMENDED) + +Password: """) + + if password != "": + key = gen_encryption_key(password)# Get the key version of the password + f = Fernet(key)# Create encryption object with the key version of the password + content["private_key"] = f.encrypt(content["private_key"].encode()).decode()# Use it to encrypt the private key + content["encryption"] = True# Make the program know that this file requires a password. + with open(filename, "w") as f: + f.write(dumps(content)) + else: + with open(filename, "w") as f: + f.write(dumps(content)) - filename = input("Write the name of your new address: ") + ".txt" - with open(filename, "w") as f: - f.write("Private key: {0}\nWallet address / Public key: {1}".format(private_key, public_key.decode())) print("Your new address and private key are now in the file {0}".format(filename)) def sign_ECDSA_msg(private_key):