104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
from cryptography.fernet import Fernet
|
|
from cryptography.fernet import InvalidToken
|
|
from pathlib import Path
|
|
import settings
|
|
import string
|
|
import secrets
|
|
import sys
|
|
import os
|
|
|
|
|
|
# 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
|
|
|
|
|
|
# Encrypting and storing of key
|
|
def encrypt_key(message):
|
|
key = load_key()
|
|
keyf = Fernet(key)
|
|
|
|
with open('uploadkeys', 'a+') as uploadkeys:
|
|
print(str(token), file=uploadkeys)
|
|
|
|
with open("uploadkeys", "rb") as keyfile:
|
|
keyfile_data = keyfile.read()
|
|
|
|
encrypted_data = keyf.encrypt(keyfile_data)
|
|
|
|
with open("uploadkeys", "wb") as keyfile:
|
|
keyfile.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
|
|
|
|
|
|
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:
|
|
|
|
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()
|
|
|
|
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
|