ADDED PRIVATE KEY ENCRYPTION

.gitignore

@@ -1 +1,3 @@
-*.pyc
+*.pyc
+venv/
+.vscode/

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):