使用Fernet进行加解密

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import base64
import hashlib
import os

from cryptography.fernet import Fernet


def get_secret_key():
    key = os.getenv("SECRET_KEY")
    if key is None:
        raise EnvironmentError("SECRET_KEY not set")

    hash_object = hashlib.sha256(key.encode())
    hash_digest = hash_object.digest()
    key_bytes = hash_digest[:32]
    base64_key = base64.urlsafe_b64encode(key_bytes)
    return base64_key.decode('utf-8')


def encrypt(data, key=None):
    key = key or get_secret_key()
    cipher_suite = Fernet(key)
    encrypted_data = cipher_suite.encrypt(data.encode())
    return encrypted_data


def decrypt(data, key=None):
    key = key or get_secret_key()
    cipher_suite = Fernet(key)
    decrypted_data = cipher_suite.decrypt(data)
    return decrypted_data.decode()


def simple_test():
    plain_data = 'Test_1234'
    cipher_data = encrypt(plain_data)
    print(cipher_data)
    assert decrypt(cipher_data) == plain_data


if __name__ == '__main__':
    simple_test()