thesis-TA-resistant-anonymity-set/simulation/encrypt_message.py

# -*- coding: utf-8 -*-
import os

from cryptography.fernet import Fernet

SECRET_KEY = "secret.key"


def generate_key():
    key = Fernet.generate_key()
    with open(SECRET_KEY, "wb") as key_file:
        key_file.write(key)


def load_key():
    return open(SECRET_KEY, "rb").read()


def encrypt_message(message):
    """Encrypts a message string"""
    if not os.path.exists(SECRET_KEY):
        generate_key()
    key = load_key()
    encoded_message = message.encode("UTF-8")
    f = Fernet(key)
    encrypted_message = f.encrypt(encoded_message)
    return encrypted_message


def decrypt_message(encrypted_message):
    key = load_key()
    f = Fernet(key)
    decrypted_message = f.decrypt(encrypted_message)
    return decrypted_message.decode()