Clean up and reorganize keygen.py

This makes it easier to add new features
- Clearer variable names
- Split up large functions into smaller ones for each action
- Help make it easier to add new features

keygen.py

@@ -8,41 +8,63 @@ import sys
 import os
 
 
-# Check if encryption key already exists
+# Load secret
-enckey = Path(settings.ENCKEY_PATH)
+def load_secret():
-if enckey.is_file():
+    with open(settings.ENCKEY_PATH, "rb") as sf:
-    print("Encryption key found.")
+        secret = sf.read()
-else:
+    return secret
-    print("Encryption key not found.")
-    print("Generating key...")
-    key = Fernet.generate_key()
-    with open(settings.ENCKEY_PATH, "wb") as key_file:
-        key_file.write(key)
-    print("Encryption key generated and stored in secret.key.")
-
-
-# Load encryption key
-def load_key():
-    with open(settings.ENCKEY_PATH, "rb") as kf:
-        kdata = kf.read()
-    return kdata
 
 
 # Encrypting and storing of key
-def encrypt_key(message):
+def append_uploadkey(akey):
-    key = load_key()
+    with open('uploadkeys', 'a+') as uploadkeysf:
-    keyf = Fernet(key)
+        print(str(akey), file=uploadkeysf)
 
-    with open('uploadkeys', 'a+') as uploadkeys:
-        print(str(token), file=uploadkeys)
 
-    with open("uploadkeys", "rb") as keyfile:
+def decrypt_uploadkeys():
-        keyfile_data = keyfile.read()
+    with open("uploadkeys", "rb") as uploadkeysf:
+        uploadkeys_data = uploadkeysf.read()
 
-    encrypted_data = keyf.encrypt(keyfile_data)
+    try:
+        secret = load_secret()
+        secretf = Fernet(secret)
+        decrypted_data = secretf.decrypt(uploadkeys_data)  # decrypt data
+        with open("uploadkeys", "wb") as ukf:
+            ukf.write(decrypted_data)  # write the original file
+        print("Done decrypting")  # debug
+        return True
+    except InvalidToken:
+        print("InvalidToken")  # debug
+        print("The encrypted key data is invalid and cannot be read.")
+        print("It may be necessary to clear the file entirely, which will invalidate all tokens.")
+        proceed = ask_yn("Do you wish to proceed to clearing the uploadkeys file? [y/n] ")
 
-    with open("uploadkeys", "wb") as keyfile:
+        if proceed:
-        keyfile.write(encrypted_data)
+            print("Proceed1")
+            os.remove("uploadkeys")
+            print("Removed uploadkeys file.")
+            proceed2 = ask_yn("Would you like to continue and generate a new key? [y/n] ")
+            if not proceed2:
+                print("not proceed2")
+                return False
+            else:
+                print("proceed2")
+                return True
+        else:
+            print("not Proceed1")
+            return False
+
+
+def encrypt_uploadkeys():
+    with open("uploadkeys", "rb") as uploadkeysf:
+        uploadkeys_data = uploadkeysf.read()
+
+    secret = load_secret()
+    secretf = Fernet(secret)
+    encrypted_data = secretf.encrypt(uploadkeys_data)
+
+    with open("uploadkeys", "wb") as uploadkeysf:
+        uploadkeysf.write(encrypted_data)
 
 
 def ask_yn(msg):
@@ -58,20 +80,27 @@ def ask_yn(msg):
     return proceed
 
 
+# Check if encryption secret already exists
+if Path(settings.ENCKEY_PATH).is_file():
+    print("Encryption secret found.")
+else:
+    print("Encryption secret not found.")
+    print("Generating secret...")
+    newsecret = Fernet.generate_key()
+    with open(settings.ENCKEY_PATH, "wb") as secret_file:
+        secret_file.write(newsecret)
+    print("Encryption secret generated and stored in {0}".format(settings.ENCKEY_PATH))
+
+
+if __name__ == "__main__":
     start = ask_yn("Have you run this program as the correct user (for example, nginx uses www-data)? [y/n] ")
     if not start:
-    print("Please run this as the correct user with: sudo su [user] -s /bin/sh -c 'python3 keygen/py'")
+        print("Please run this as the correct user with: sudo su [user] -s /bin/sh -c 'python3 keygen.py'")
 
     else:
-
-    N = 64  # Size of token
-
-    # Generate key
-    token = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(N))
-
         # Decrypt the existing keyfile
-    key = load_key()
+        secret = load_secret()
-    keyf = Fernet(key)
+        keyf = Fernet(secret)
 
         genkey = True
         uploadkeysp = Path("uploadkeys")
@@ -82,22 +111,13 @@ else:
                 # read the encrypted data
                 encrypted_data = ukf.read()
 
-        try:
-            decrypted_data = keyf.decrypt(encrypted_data)  # decrypt data
-            with open("uploadkeys", "wb") as ukf:
-                ukf.write(decrypted_data)  # write the original file
-        except InvalidToken:
-            print("The encrypted key data is invalid and cannot be read.")
-            print("It may be necessary to clear the file entirely, which will invalidate all tokens.")
-            proceed = ask_yn("Do you wish to proceed to clearing the uploadkeys file? [y/n] ")
-
-            if proceed:
-                os.remove("uploadkeys")
-                print("Removed uploadkeys file.")
-                proceed2 = ask_yn("Would you like to continue and generate a new token? [y/n] ")
-                if not proceed2:
-                    genkey = False
 
         if genkey:
-        print("Your new token is: " + str(token))  # Print token
+            if decrypt_uploadkeys():  # Decrypt the file
-        encrypt_key(str(token))  # Encrypt the key and save
+                N = 64  # Size of key
+                key = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(N))
+                print("Your new key is: " + str(key))  # Print key
+                append_uploadkey(key)  # Save the new key to file unencrypted
+                encrypt_uploadkeys()  # Encrypt the uploadkeys file
+            else:
+                print("Exiting.")