SamAuthenticator/SamAuthenticator/Authenticator.py

import onetimepass as otp
import json
import cryptography.fernet
import argon2
import base64
import os
import copy

_salt = "V5RlhpuwACffXuUNLex7Al9ulPy4SRHbyaAxWigjX9Z01OVaCO"


def get_default_salt():
    return copy.deepcopy(_salt)


def encrypt_data(data_bytes, password, salt):
    password_hash = argon2.argon2_hash(password=password, salt=salt)
    encoded_hash = base64.urlsafe_b64encode(password_hash[:32])
    encryptor = cryptography.fernet.Fernet(encoded_hash)
    return encryptor.encrypt(data_bytes)


def decrypt_data(cipher_bytes, password, salt):
    password_hash = argon2.argon2_hash(password=password, salt=salt)
    encoded_hash = base64.urlsafe_b64encode(password_hash[:32])
    decryptor = cryptography.fernet.Fernet(encoded_hash)
    return decryptor.decrypt(cipher_bytes)


def write_keys_to_file(auth_keys, password, file_name="data.dat"):
    backup_file = file_name + ".bak"
    if os.path.exists(file_name):
        os.rename(file_name, backup_file)
    with open(file_name, 'wb') as f:
        f.write(encrypt_data(auth_keys.dump_data().encode('utf-8'), password, _salt))
        if os.path.exists(backup_file):
            os.remove(backup_file)


def read_keys_from_file(password, file_name="data.dat"):
    with open(file_name, 'rb') as f:
        ciphered_data = f.read()
        readable_data = decrypt_data(ciphered_data, password, _salt)
        keys_object = AuthenticatorKeys()
        keys_object.read_dump(readable_data.decode('utf-8'))
        return keys_object


class AuthenticatorKeys:
    def __init__(self):
        self.data = {'secrets': {},
                     'version': '1.0'}

    def set_secret(self, name, secret):
        self.data['secrets'][name] = {'secret': secret}

    def get_secret(self, name):
        return self.data['secrets'][name]['secret']

    def get_token(self, name):
        return otp.get_totp(self.get_secret(name))

    def remove_secret(self, name):
        del self.data['secrets'][name]

    def get_names(self):
        return self.data['secrets'].keys()

    def get_size(self):
        return len(self.data['secrets'])

    def dump_data(self):
        return json.dumps(self.data)

    def read_dump(self, dump_data):
        self.data = json.loads(dump_data)

    @staticmethod
    def test_secret_validity(secret):
        otp.get_totp(secret)