Rename functions.py to functions.py.default

Since this is default settings and the user might want to customize
them, functions.py has been renamed. This also will prevent conflicts if
the user has updated their functions.py and then tries to pull.

.gitignore

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

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)
-
+import functions  # custom functions
-
-ALPHANUMERIC = string.ascii_letters + string.digits  # uppercase, lowercase, and numbers
 
 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,11 +8,6 @@ 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():
@@ -63,20 +58,26 @@ def ask_yn(msg):
     return proceed
 
 
-N = 64  # Size of token
+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'")
 
-# 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)
-
-genkey = True
-uploadkeysp = Path("uploadkeys")
-if not uploadkeysp.is_file():
-    uploadkeysp.touch()
 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()
+    keyf = Fernet(key)
+
+    genkey = True
+    uploadkeysp = Path("uploadkeys")
+    if not uploadkeysp.is_file():
+        uploadkeysp.touch()
+    else:
         with open("uploadkeys", "rb") as ukf:
             # read the encrypted data
             encrypted_data = ukf.read()
@@ -97,6 +98,6 @@ else:
                 if not proceed2:
                     genkey = False
 
-if genkey:
+    if genkey:
         print("Your new token is: " + str(token))  # Print token
         encrypt_key(str(token))  # Encrypt the key and save