import os
import base64

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher


# 获取密钥
def get_key(**kwargs):
    # 密钥文件
    key_file = kwargs['key_file']
    # 密钥文件绝对路径
    rsa_key_path = os.path.dirname(__file__) + '/static/rsa/' + key_file
    with open(rsa_key_path) as f:
        data = f.read()
        key = RSA.importKey(data)
    return key


# 用公钥加密
def encrypt_data(**kwargs):
    # 待加密信息
    msg = kwargs['msg']
    # 获取密钥
    public_key = get_key(key_file='rsa_public_key.pem')
    cipher = PKCS1_cipher.new(public_key)
    encrypt_text = base64.b64encode(cipher.encrypt(bytes(msg.encode("utf8"))))
    return encrypt_text.decode('utf-8')


# 用私钥解密
def decrypt_data(**kwargs):
    # 待加密信息
    encrypt_msg = kwargs['encrypt_msg']
    # 获取密钥
    private_key = get_key(key_file='rsa_private_key.pem')
    cipher = PKCS1_cipher.new(private_key)
    back_text = cipher.decrypt(base64.b64decode(encrypt_msg), 0)
    return back_text.decode('utf-8')


# 检验
def check_encrypt_decrypt(**kwargs):
    msg = kwargs['msg']
    decrypt_msg = decrypt_data(encrypt_msg=kwargs['encrypt_msg'])
    if msg == decrypt_msg:
        return True
    else:
        return False