Simplified encryption logic
This commit is contained in:
parent
9975eaab63
commit
08206db1d9
@ -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)
|
||||||
|
Reference in New Issue
Block a user