imgupload-legacy/keygen.py

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.")