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

@@ -133,4 +133,5 @@ uploadkeys
 savelog.log
 uwsgi.log
 settings.py
+functions.py
 secret.key

README.md

@@ -4,6 +4,23 @@
 ### What is imgupload?
 imgupload is a Flask + uWSGI application to serve as an all-purpose image/file uploader over POST requests.
 
-### Usage
-Make sure you install the dependencies first. To do this, run `sudo python3 -m pip install -r requirements.txt`.
-To deploy imgupload, run `flask run`.
+### Installation
+1. Clone the repository: `git clone https://github.com/BBaoVanC/imgupload.git`
+2. Enter the imgupload directory: `cd imgupload`
+3. Create a virtualenv: `python3 -m venv env`
+4. Enter the virtualenv: `source env/bin/activate`
+5. Install dependencies: `python3 -m pip install -r requirements.txt`
+6. Run the Flask app
+
+## Running the Flask app
+### Using uWSGI
+[https://uwsgi-docs.readthedocs.io/en/latest/Configuration.html](https://uwsgi-docs.readthedocs.io/en/latest/Configuration.html)
+
+Instructions specific to imgupload are coming soon
+
+### Using Flask development server
+```shell
+$ source env/bin/activate  # if you haven't already entered the virtualenv
+$ export FLASK_APP=imgupload.py
+$ flask run
+```

functions.py.default

@@ -0,0 +1,8 @@
+import string
+import random
+
+
+def generate_name():
+    chars = string.ascii_letters + string.digits  # uppercase, lowercase, and numbers
+    name = ''.join((random.choice(chars) for i in range(8)))  # generate name
+    return name

imgupload.py

@@ -2,15 +2,12 @@ from flask import Flask, request, jsonify, abort, Response
 from cryptography.fernet import Fernet
 from flask_api import status
 from pathlib import Path
-import string
 import random
 import os
 import datetime
 
 import settings  # app settings (such as allowed extensions)
-
-
-ALPHANUMERIC = string.ascii_letters + string.digits  # uppercase, lowercase, and numbers
+import functions  # custom functions
 
 app = Flask(__name__)  # app is the app
 
@@ -22,15 +19,6 @@ def allowed_extension(testext):
         return False
 
 
-def generate_name(extension):
-    namefound = False
-    while not namefound:
-        fname = ''.join((random.choice(ALPHANUMERIC) for i in range(8))) + str(extension)
-        if not Path(fname).is_file():
-            namefound = True
-    return fname
-
-
 def log_savelog(key, ip, savedname):
     if settings.SAVELOG_KEYPREFIX > 0:
         with open(settings.SAVELOG, "a+") as slogf:
@@ -77,7 +65,7 @@ def upload():
                 fext = Path(f.filename).suffix  # get the uploaded extension
                 if allowed_extension(fext):  # if the extension is allowed
                     print("Generating file with extension {0}".format(fext))
-                    fname = generate_name(fext)  # generate file name
+                    fname = functions.generate_name() + fext  # generate file name
                     print("Generated name: {0}".format(fname))
 
                     if f:  # if the uploaded image exists

keygen.py

@@ -8,46 +8,59 @@ import sys
 import os
 
 
-# Check if the script was run as root
-if os.geteuid() != 0:
-    exit("Root privileges are necessary to run this script.\nPlease try again as root or using `sudo`.")
-
-
-# Check if encryption key already exists
-enckey = Path(settings.ENCKEY_PATH)
-if enckey.is_file():
-    print("Encryption key found.")
-else:
-    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
+# Load secret
+def load_secret():
+    with open(settings.ENCKEY_PATH, "rb") as sf:
+        secret = sf.read()
+    return secret
 
 
 # Encrypting and storing of key
-def encrypt_key(message):
-    key = load_key()
-    keyf = Fernet(key)
+def append_uploadkey(akey):
+    with open('uploadkeys', 'a+') as uploadkeysf:
+        print(str(akey), file=uploadkeysf)
 
-    with open('uploadkeys', 'a+') as uploadkeys:
-        print(str(token), file=uploadkeys)
 
-    with open("uploadkeys", "rb") as keyfile:
-        keyfile_data = keyfile.read()
+def decrypt_uploadkeys():
+    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:
-        keyfile.write(encrypted_data)
+        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):
@@ -63,40 +76,33 @@ def ask_yn(msg):
     return proceed
 
 
-N = 64  # Size of token
+# 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))
 
-# Generate key
-token = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(N))
 
-# Decrypt the existing keyfile
-key = load_key()
-keyf = Fernet(key)
+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'")
 
-genkey = True
+    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:
-    with open("uploadkeys", "rb") as ukf:
-        # 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
-    encrypt_key(str(token))  # Encrypt the key and save
+            print("Exiting.")