Bugfixes in keygen.py

- Handle if uploadkeys becomes corrupted
- Disambiguate variable names
- Handle case where the uploadkeys file doesn't already exist
This commit is contained in:
BBaoVanC 2020-08-31 20:32:00 -05:00
parent 797bebb1a1
commit 7ccaafc6c6
Signed by: bbaovanc
GPG Key ID: 18089E4E3CCF1D3A

View File

@ -1,4 +1,5 @@
from cryptography.fernet import Fernet from cryptography.fernet import Fernet
from cryptography.fernet import InvalidToken
from pathlib import Path from pathlib import Path
import settings import settings
import string import string
@ -6,11 +7,13 @@ import secrets
import sys import sys
import os import os
# Check if the script is ran as root
# Check if the script was run as root
if os.geteuid() != 0: if os.geteuid() != 0:
exit("Root privileges are necessary to run this script.\nPlease try again as root or using `sudo`.") exit("Root privileges are necessary to run this script.\nPlease try again as root or using `sudo`.")
# Check if encryption key exists
# Check if encryption key already exists
enckey = Path(settings.ENCKEY_PATH) enckey = Path(settings.ENCKEY_PATH)
if enckey.is_file(): if enckey.is_file():
print("Encryption key found.") print("Encryption key found.")
@ -22,32 +25,18 @@ else:
key_file.write(key) key_file.write(key)
print("Encryption key generated and stored in secret.key.") print("Encryption key generated and stored in secret.key.")
# Load encryption key # Load encryption key
def load_key(): def load_key():
return open(settings.ENCKEY_PATH, "rb").read() with open(settings.ENCKEY_PATH, "rb") as kf:
kdata = kf.read()
return kdata
# Set size of string
N = 64
# Generating of key
token = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(64))
# Decrypt the existing keyfile
key = load_key()
f = Fernet(key)
with open("uploadkeys", "rb") as file:
# read the encrypted data
encrypted_data = file.read()
# decrypt data
decrypted_data = f.decrypt(encrypted_data)
# write the original file
with open("uploadkeys", "wb") as file:
file.write(decrypted_data)
# Encrypting and storing of key # Encrypting and storing of key
def encrypt_key(message): def encrypt_key(message):
key = load_key() key = load_key()
f = Fernet(key) keyf = Fernet(key)
with open('uploadkeys', 'a+') as uploadkeys: with open('uploadkeys', 'a+') as uploadkeys:
print(str(token), file=uploadkeys) print(str(token), file=uploadkeys)
@ -55,11 +44,59 @@ def encrypt_key(message):
with open("uploadkeys", "rb") as keyfile: with open("uploadkeys", "rb") as keyfile:
keyfile_data = keyfile.read() keyfile_data = keyfile.read()
encrypted_data = f.encrypt(keyfile_data) encrypted_data = keyf.encrypt(keyfile_data)
with open("uploadkeys", "wb") as keyfile: with open("uploadkeys", "wb") as keyfile:
keyfile.write(encrypted_data) keyfile.write(encrypted_data)
# Print result on display and call encrypt_key
print("Your new token is: " + str(token)) def ask_yn(msg):
encrypt_key(str(token)) 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
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