Remove a few leftover lines

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

.gitignore

@@ -134,3 +134,4 @@ savelog.log
 uwsgi.log
 settings.py
 functions.py
+secret.key

configtest.py

@@ -1,9 +1,3 @@
-#!/usr/bin/env python3
-"""
-configtest.py
-
-Tests the validity of your configuration in settings.py.
-"""
 import os
 import settings as settings
 
@@ -16,6 +10,7 @@ defaults = {
     "SAVELOG": "savelog.log",
     "SAVELOG_CHMOD": "0o644",
     "SAVELOG_KEYPREFIX": 4,
+    "ENCKEY_PATH": "secret.key"
 }
 
 deftypes = {
@@ -25,6 +20,7 @@ deftypes = {
     "SAVELOG": str,
     "SAVELOG_CHMOD": int,
     "SAVELOG_KEYPREFIX": int,
+    "ENCKEY_PATH": str,
 }
 
 
@@ -98,6 +94,16 @@ if "ROOTURL" in checksettings:
         print("[" + u"\u2713" + "] ROOTURL is good!")
 
 
+# Check if ENCKEY_PATH exists
+enckey_exists = True
+if "ENCKEY_PATH" in checksettings:
+    if not os.path.isfile(settings.ENCKEY_PATH):
+        enckey_exists = False
+        print("[!] The path set in ENCKEY_PATH ('{0}') doesn't exist!".format(settings.ENCKEY_PATH))
+    else:
+        print("[" + u"\u2713" + "] ENCKEY_PATH exists!")
+
+
 # Ask the user if SAVELOG is the intended filename
 if "SAVELOG" in checksettings:
     print("[*] SAVELOG was interpreted to be {0}".format(settings.SAVELOG))
@@ -130,6 +136,10 @@ if not uploadfolder_exists:
     summarygood = False
     print("UPLOAD_FOLDER ({0}) does not exist!".format(settings.UPLOAD_FOLDER))
 
+if not enckey_exists:
+    summarygood = False
+    print("ENCKEY_PATH ({0}) does not exist!".format(settings.ENCKEY_PATH))
+
 if not rooturl_good:
     summarygood = False
     print("ROOTURL may cause issues!")

functions.py.default

@@ -1,9 +1,3 @@
-#!/usr/bin/env python3
-"""
-functions.py
-
-Functions used by imgupload which can be easily customized.
-"""
 import string
 import random
 

imgupload.py

@@ -1,12 +1,8 @@
-#!/usr/bin/env python3
+from flask import Flask, request, jsonify, abort, Response
-"""
+from cryptography.fernet import Fernet
-imgupload.py
-
-Flask application for processing images uploaded through POST requests.
-"""
-from flask import Flask, request, jsonify, Response
 from flask_api import status
 from pathlib import Path
+import random
 import os
 import datetime
 
@@ -17,7 +13,7 @@ app = Flask(__name__)  # app is the app
 
 
 def allowed_extension(testext):
-    if testext.lower() in settings.ALLOWED_EXTENSIONS:
+    if testext in settings.ALLOWED_EXTENSIONS:
         return True
     else:
         return False
@@ -38,22 +34,24 @@ def upload():
     if request.method == "POST":  # sanity check: make sure it's a POST request
         print("Request method was POST!")
 
-        with open("uploadkeys", "r") as keyfile:  # load valid keys
+        with open(settings.ENCKEY_PATH,"rb") as enckey: # load encryption key
-            validkeys = keyfile.readlines()
+            key = enckey.read()
-        validkeys = [x.strip("\n") for x in validkeys]
+        f = Fernet(key)
+
+        with open("uploadkeys", "rb") as keyfile:
+            encrypted_data = keyfile.read()
+        decrypted_data = str(f.decrypt(encrypted_data).decode('utf-8'))
+        decrypted_data = decrypted_data.splitlines()
+
+        validkeys = [x.strip("\n") for x in decrypted_data]
         while "" in validkeys:
             validkeys.remove("")
+            print("Removed blank key(s)")
         print("Loaded validkeys")
-
         if "uploadKey" in request.form:  # if an uploadKey was provided
             if request.form["uploadKey"] in validkeys:  # check if uploadKey is valid
                 print("Key is valid!")
 
-                if "verify" in request.form.keys():
-                    if request.form["verify"] == "true":
-                        print("Request is asking if key is valid (it is)")
-                        return jsonify({'status': 'key_valid'})
-
                 if "imageUpload" in request.files:  # check if image to upload was provided
                     f = request.files["imageUpload"]  # f is the image to upload
                 else:
@@ -86,16 +84,21 @@ def upload():
 
                 else:  # if the extension was invalid
                     print("Uploaded extension is invalid!")
-                    return jsonify({'status': 'error', 'error': 'INVALID_EXTENSION'}), status.HTTP_415_UNSUPPORTED_MEDIA_TYPE
+                    abort(415)
 
             else:  # if the key was not valid
                 print("Key is invalid!")
                 print("Request key: {0}".format(request.form["uploadKey"]))
-                return jsonify({'status': 'error', 'error': 'UNAUTHORIZED'}), status.HTTP_401_UNAUTHORIZED
+                abort(401)
 
         else:  # if uploadKey was not found in request body
             print("No uploadKey found in request!")
-            return jsonify({'status': 'error', 'error': 'UNAUTHORIZED'}), status.HTTP_401_UNAUTHORIZED
+            abort(401)
+
+
+    else:  # if the request method wasn't post
+        print("Request method was not POST!")
+        abort(405)
 
 if __name__ == "__main__":
     print("Run with `flask` or a WSGI server!")

keyctl.py

@@ -1,181 +0,0 @@
-#!/usr/bin/env python3
-"""
-keyctl.py
-
-Command-line utility for easy management of the uploadkeys file.
-"""
-
-from pathlib import Path
-import argparse
-import logging
-import secrets
-import string
-
-
-def read_keyfile():
-    with open("uploadkeys", "r") as keyfile:  # open uploadkeys
-        keys = keyfile.readlines()  # read all the keys
-        logging.debug("Read uploadkeys")
-    keys = [x.strip("\n") for x in keys]  # strip newlines from keys
-    logging.debug("Stripped newlines from keys")
-    return keys
-
-
-def genkey(length):
-    key = ''.join(secrets.choice(string.ascii_letters + string.digits) for x in range(length))
-    return key
-
-
-def savekey(key):
-    if not Path("uploadkeys").is_file():  # if uploadkeys doesn't exist, log an info message
-        logging.info("uploadkeys file doesn't exist, it will be created.")
-    with open("uploadkeys", "a+") as keyfile:
-        keyfile.write(str(key) + "\n")  # add the key
-    logging.debug("Saved a key to uploadkeys: {0}".format(key))
-
-
-def rmkey(delkey):
-    removedkey = False
-    allkeys = read_keyfile()
-    if delkey in allkeys:  # if the key to remove exists
-        allkeys.remove(delkey)  # remove the first instance of the key
-        removedkey = True
-        logging.debug("Removed one instance of the key")
-
-    with open("uploadkeys", "w") as keyfile:
-        for k in allkeys:
-            keyfile.write(k + "\n")  # write the remaining keys
-
-    if removedkey:
-        return True
-    else:
-        return False
-
-
-def find_duplicates():
-    allkeys = read_keyfile()
-
-    seen = set()
-    ukeys = []
-    dupkeys = []
-    for x in allkeys:
-        if x not in seen:
-            ukeys.append(x)
-            seen.add(x)
-        else:
-            dupkeys.append(x)
-    return dupkeys
-
-
-def get_keys():
-    validkeys = read_keyfile()
-    while "" in validkeys:
-        validkeys.remove("")
-    logging.debug("Removed blank keys")
-    return validkeys
-
-
-
-def cmd_list(args):
-    validkeys = get_keys()
-
-    print("List of upload keys:")
-    for i in range(len(validkeys)):
-        showkey = validkeys[i][:6]
-        if len(validkeys[i]) > 6:
-            showkey += "..."  # add ellipses since the key was shortened in list
-
-        print("    [{0}] {1}".format(i+1, showkey))
-
-
-def cmd_generate(args):
-    k = genkey(args.length)
-    logging.debug("Generated a new key: {0}".format(k))
-    savekey(k)
-    print("Your new key is: {0}".format(k))
-
-
-def cmd_add(args):
-    print("Please type/paste the key you would like to add.")
-    akr = input("> ")
-    ak = akr.strip()
-    print()
-    logging.debug("Ran strip() on key")
-    print(ak)
-    if input("Is the above key correct? [y/N] ").lower() == "y":
-        logging.debug("Interpreted as yes")
-        ask_for_key = False
-        savekey(ak)
-        logging.info("Added.")
-    else:
-        logging.debug("Interpreted as no")
-        print("No key has been saved.")
-
-
-def cmd_remove(args):
-    if rmkey(args.key):
-        logging.debug("Successfully removed the requested key")
-    else:
-        logging.info("No key was removed.")
-
-def cmd_dedupe(args):
-    dupes = find_duplicates()
-    if len(dupes) > 0:
-        for d in dupes:
-            r = rmkey(d)
-            logging.debug(r)
-            logging.info("Removed duplicate key: {0}".format(d))
-    else:
-        logging.info("[" + u"\u2713" + "] No duplicate keys found!")
-
-def cmd_show(args):
-    for k in get_keys():
-        if k[:6] == args.prefix:
-            print("Key: {0}".format(k))
-            break
-
-
-parser = argparse.ArgumentParser()  # create instance of argument parser class
-
-parlog = parser.add_mutually_exclusive_group()
-parlog.add_argument("-v", "--verbose", help="show debugging messages", action="store_true")
-parlog.add_argument("-q", "--quiet", help="show only warning messages and up", action="store_true")
-
-subparsers = parser.add_subparsers(help="sub-commands")
-parser_list = subparsers.add_parser("list", help="list the beginning of each key")
-parser_list.set_defaults(func=cmd_list)
-
-parser_gen = subparsers.add_parser("generate", help="generate a key and save it to uploadkeys")
-parser_gen.add_argument("length", help="length of key to generate", default=64, type=int, nargs="?")
-parser_gen.set_defaults(func=cmd_generate)
-
-parser_add = subparsers.add_parser("add", help="prompts for a key to add to uploadkeys")
-parser_add.set_defaults(func=cmd_add)
-
-parser_remove = subparsers.add_parser("remove", help="remove (one instance of) a key from uploadkeys")
-parser_remove.add_argument("key", help="key to remove")
-parser_remove.set_defaults(func=cmd_remove)
-
-parser_dedupe = subparsers.add_parser("dedupe", help="remove duplicate keys")
-parser_dedupe.set_defaults(func=cmd_dedupe)
-
-parser_show = subparsers.add_parser("show", help="show the full key based on the first 6 characters")
-parser_show.add_argument("prefix", help="first 6 characters of key (shown by `python3 keyctl.py list`)")
-parser_show.set_defaults(func=cmd_show)
-
-
-args = parser.parse_args()  # parse the arguments
-
-if args.verbose:
-    loglevel = logging.DEBUG
-elif args.quiet:
-    loglevel = logging.WARNING
-else:
-    loglevel = logging.INFO
-logging.basicConfig(level=loglevel, format="%(levelname)s: %(message)s")
-
-try:
-    args.func(args)
-except AttributeError:
-    logging.error("AttributeError")
-    parser.print_help()

keygen.py

@@ -0,0 +1,108 @@
+from cryptography.fernet import Fernet
+from cryptography.fernet import InvalidToken
+from pathlib import Path
+import settings
+import string
+import secrets
+import sys
+import os
+
+
+# Load secret
+def load_secret():
+    with open(settings.ENCKEY_PATH, "rb") as sf:
+        secret = sf.read()
+    return secret
+
+
+# Encrypting and storing of key
+def append_uploadkey(akey):
+    with open('uploadkeys', 'a+') as uploadkeysf:
+        print(str(akey), file=uploadkeysf)
+
+
+def decrypt_uploadkeys():
+    with open("uploadkeys", "rb") as uploadkeysf:
+        uploadkeys_data = uploadkeysf.read()
+
+    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] ")
+
+        if proceed:
+            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:
+                return False
+            else:
+                return True
+        else:
+            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):
+    resps = {"y": True, "n": False}
+    ask = True
+    while ask:
+        proceedraw = input(msg)
+        if proceedraw.lower() in resps.keys():
+            proceed = resps[proceedraw]
+            ask = False
+        else:
+            print("Invalid response.")
+    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'")
+
+    else:
+        uploadkeysp = Path("uploadkeys")
+        if not uploadkeysp.is_file():
+            uploadkeysp.touch()
+
+        if decrypt_uploadkeys():  # Decrypt the file
+            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.")

requirements.txt

@@ -1,2 +1,3 @@
 Flask_API==2.0
+cryptography==3.1
 Flask==1.1.2

settings.py.default

@@ -1,13 +1,7 @@
-#!/usr/bin/env python3
-"""
-settings.py
-
-User-defined settings used by imgupload.py.
-"""
-
 UPLOAD_FOLDER = "/path/to/images"
 ALLOWED_EXTENSIONS = [".png", ".jpg", ".jpeg", ".svg", ".bmp", ".gif", ".ico",  ".webp"]
 ROOTURL = "https://example.com/"
 SAVELOG = "savelog.log"
 SAVELOG_CHMOD = 0o644
 SAVELOG_KEYPREFIX = 4
+ENCKEY_PATH = "secret.key"