import os
import base64

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


def get_key(**kwargs):
    """
    获取密钥
    Parameters:
        key_file str 密钥文件名称
    Returns:
        key 密钥
    """
    key_file = kwargs['key_file']  # 密钥文件
    key_path = os.path.abspath(os.path.abspath(os.path.dirname(__file__)) + '/keys/' + key_file)  # 密钥文件绝对路径
    with open(key_path) as f:
        data = f.read()
        key = RSA.importKey(data)
    return key


def encrypt_data(**kwargs):
    """
    用公钥加密
    Parameters:
        msg str 待加密信息
    Returns:
        加密后结果
    """
    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):
    """
    用私钥解密
    Parameters:
        encrypt_msg str 加密信息
    Returns:
        执行正确 解密后结果
        执行错误 False
    """
    try:
        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')
    except Exception:
        return False


def check_data(**kwargs):
    """
    对比数据和加密数据
    Parameters:
        msg str 数据
        encrypt_msg str 加密数据
    Returns:
        True 对比正确
        False 对比错误
    """
    msg = kwargs['msg']
    decrypt_msg = decrypt_data(encrypt_msg=kwargs['encrypt_msg'])
    if decrypt_msg is False:
        return False
    elif msg == decrypt_msg:
        return True
    else:
        return False