Anon/MiraiNikkiBot
Archived
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Simplified encryption logic

Browse Source
This commit is contained in:
Anon 2022-10-08 17:35:02 -07:00
parent 9975eaab63
commit 08206db1d9
1 changed files with 61 additions and 88 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

149
src/encryption.py
View File

@ -28,13 +28,11 @@ from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from collections import OrderedDict from collections import OrderedDict
class PasswordMismatch(Exception): class EncryptionFail(Exception):
pass pass
class PasswordMismatch(Exception):
# Salts pass
def generate_salt():
return os.urandom(16)
# Return string from bytes # Return string from bytes
@ -72,8 +70,7 @@ def derive_key(password, salt):
iterations=100000, iterations=100000,
backend=default_backend() backend=default_backend()
) )
r_key = base64.urlsafe_b64encode(kdf.derive(password)) return base64.urlsafe_b64encode(kdf.derive(password))
return r_key
# Encryption functions # Encryption functions
@ -99,94 +96,62 @@ def decrypt(token, key):
# encryption_function: encrypt(message, key) : decrypt(token, key): # encryption_function: encrypt(message, key) : decrypt(token, key):
# Returns settings_server_decrypted dictionary with Byte() values. Will need to use # Returns settings_server_decrypted dictionary with Byte() values. Will need to use
# ChangeEncodingDict to make them strings (recommended cfg file friendly) # ChangeEncodingDict to make them strings (recommended cfg file friendly)
def __settings_server(password, salt, settings_server, encryption_function): def encrypt_settings(settings_server, password, salt, encryption_function):
key = derive_key(password, salt) key = derive_key(password, salt)
settings_server_decrypted = OrderedDict() settings_server_decrypted = OrderedDict()
for setting in settings_server: for setting in settings_server:
settings_server_decrypted[setting] = encryption_function(settings_server[setting], key) settings_server_decrypted[setting] = encryption_function(settings_server[setting], key)
return settings_server_decrypted return settings_server_decrypted
def get_keyfile(keyfile=None):
# Returns (salt, settings_server) if keyfile is not None:
def _settings_server_encrypt(settings_server): with open(keyfile, "rb") as f:
salt = generate_salt() return f.read()
password = getpass.getpass("Enter password: ") return None
password2 = getpass.getpass("Retype Password: ")
if password != password2:
raise PasswordMismatch
settings_server_encrypted = __settings_server(password.encode(), salt, encode_dict(settings_server), encrypt)
return salt, settings_server_encrypted
# Returns (settings_server) def get_pass(q):
def _settings_server_decrypt(settings_server, settings_encrypt):
settings_server_encoded = encode_dict(settings_server)
if settings_encrypt["encrypt"]:
salt = salt_decode(settings_encrypt["salt"])
password = getpass.getpass("Enter password: ")
return __settings_server(password.encode(), salt, settings_server_encoded, decrypt)
else:
return settings_server_encoded
# Wrapper function that will catch exceptions and exit
def settings_server_new(function, **kwargs):
try: try:
return function(**kwargs) return getpass.getpass(q).encode()
# If the user cancels the login
except KeyboardInterrupt: except KeyboardInterrupt:
print("\nQuitting...") raise EncryptionFail("\nQuitting...")
# If the user passwords do not match (encrypt)
except PasswordMismatch:
print("Passwords do not match...")
# Incorrect password entered (decrypt) def settings_server_encrypt(settings_server, keyfile=None):
except InvalidToken: try:
print("Password or Token Incorrect...") settings_server = encode_dict(settings_server)
salt = os.urandom(16)
# Probably the salt value got modified password = get_keyfile(keyfile)
except base64.binascii.Error: if password is None:
print("Salt is invalid...") password = get_pass("Enter password: ")
password2 = get_pass("Enter password: ")
# Some other kind of fuck up if password != password2:
raise PasswordMismatch("Passwords do not match")
encrypted = encrypt_settings(settings_server, password, salt, encrypt)
return salt_encode(salt), decode_dict(encrypted)
except PasswordMismatch as e:
raise EncryptionFail(str(e))
except Exception as e: except Exception as e:
print("Unknown exception occurred...") err = str(e)
print(e) raise EncryptionFail("Encrypt Error: {}".format(err))
# Exit if an exception was thrown
sys.exit(1)
# Glue functions that package **kwargs automatically def settings_server_decrypt(settings_server, settings_encrypt, keyfile=None):
def settings_server_encrypt(settings_server): try:
kwargs = {"settings_server": settings_server} if not settings_encrypt["encrypt"]:
return settings_server_new(_settings_server_encrypt, **kwargs) return settings_server
settings_server = encode_dict(settings_server)
password = get_keyfile(keyfile or settings_encrypt["keyfile"]) or get_pass("Enter password: ")
def settings_server_decrypt(settings_server, settings_encrypt): salt = salt_decode(settings_encrypt["salt"])
kwargs = { decrypted = encrypt_settings(settings_server, password, salt, decrypt)
"settings_server": settings_server, return decode_dict(decrypted)
"settings_encrypt": settings_encrypt except base64.binascii.Error:
} raise EncryptionFail("Salt is invalid")
return settings_server_new(_settings_server_decrypt, **kwargs) except InvalidToken:
raise EncryptionFail("Password or token is incorrect")
except Exception as e:
# The _cfg functions should return a regular string err = str(e)
# These are the functions that should interface with the bot a return a plain string raise EncryptionFail("Decrypt Error: {}".format(err))
# settings_server ordered dictionary
def settings_server_encrypt_cfg(settings_server):
salt, settings_server = settings_server_encrypt(settings_server)
return salt_encode(salt), decode_dict(settings_server)
def settings_server_decrypt_cfg(settings_server, settings_encrypt):
settings_server = settings_server_decrypt(settings_server, settings_encrypt)
return decode_dict(settings_server)
def main(): def main():
@ -202,6 +167,7 @@ def main():
parser.add_argument("--encrypt", help="Generate encrypted authentication.", action="store_true") parser.add_argument("--encrypt", help="Generate encrypted authentication.", action="store_true")
parser.add_argument("--decrypt", help="Decrypt encrypted authentication", action="store_true") parser.add_argument("--decrypt", help="Decrypt encrypted authentication", action="store_true")
parser.add_argument("--recrypt", help="Recrypt encrypted authentication", action="store_true") parser.add_argument("--recrypt", help="Recrypt encrypted authentication", action="store_true")
parser.add_argument("-k", "--keyfile", help="Keyfile used for decryption", default=None)
parser.add_argument("-c", "--cfg", help="Specify config file.", default=default_cfg) parser.add_argument("-c", "--cfg", help="Specify config file.", default=default_cfg)
arguments = parser.parse_args() arguments = parser.parse_args()
@ -214,25 +180,28 @@ def main():
import importlib import importlib
cfg = importlib.import_module(arguments.cfg) cfg = importlib.import_module(arguments.cfg)
settings_server = cfg.settings_server settings_server = cfg.settings_server
settings_encrypt = None settings_encrypt = cfg.settings_encrypt
keyfile = arguments.keyfile or settings_encrypt["keyfile"]
if arguments.decrypt and arguments.encrypt: if arguments.decrypt and arguments.encrypt:
print("Re-encrypting") print("Re-encrypting")
if arguments.decrypt: # arguments.decrypt if arguments.decrypt: # arguments.decrypt
print("Decrypt...") print("Decrypt...")
settings_server = settings_server_decrypt_cfg(cfg.settings_server, cfg.settings_encrypt) settings_server = settings_server_decrypt(settings_server, settings_encrypt, arguments.keyfile)
settings_encrypt = OrderedDict([ settings_encrypt = OrderedDict([
("encrypt", False), ("encrypt", False),
("salt", cfg.settings_encrypt["encrypt"]) ("salt", settings_encrypt["salt"]),
("keyfile", arguments.keyfile)
]) ])
if arguments.encrypt: if arguments.encrypt:
print("Encrypt...") print("Encrypt...")
salt, settings_server = settings_server_encrypt_cfg(settings_server) salt, settings_server = settings_server_encrypt(settings_server, keyfile)
settings_encrypt = OrderedDict([ settings_encrypt = OrderedDict([
("encrypt", True), ("encrypt", True),
("salt", salt) ("salt", salt),
("keyfile", arguments.keyfile)
]) ])
print("settings_server = {}".format(pformat(settings_server))) print("settings_server = {}".format(pformat(settings_server)))
@ -241,4 +210,8 @@ def main():
if __name__ == "__main__": if __name__ == "__main__":
sys.exit(main()) try:
sys.exit(main())
except EncryptionFail as e:
print(e)
sys.exit(1)