import os

from Crypto.Cipher import DES
import rsa
import hashlib


# 生成公私钥对
from rsa import common, transform, core


def generate_keypair(user_name):
    (pubkey, prikey) = rsa.newkeys(1024)
    # 存储公钥
    with open(user_name + "_pub.pem", "wb") as fp:
        fp.write(pubkey.save_pkcs1())
    # 存储私钥
    with open(user_name + "_pri.pem", "wb") as fp:
        fp.write(prikey.save_pkcs1())


# 使用rsa加密
def rsa_encrypt(msg, user_name):
    # 读取公钥
    with open(user_name + "_pub.pem", "rb") as fp:
        pubkey = rsa.PublicKey.load_pkcs1(fp.read())
    # 使用该用户公钥加密
    return rsa.encrypt(msg, pubkey)


# 使用rsa解密
def rsa_decrypt(crypto, user_name):
    # 读取私钥
    with open(user_name + "_pri.pem", "rb") as fp:
        prikey = rsa.PrivateKey.load_pkcs1(fp.read())
    # 解密
    return rsa.decrypt(crypto, prikey)


# 使用rsa签名
def rsa_sign(msg: bytes, user_name: str):
    # 读取私钥
    with open(user_name + "_pri.pem", "rb") as fp:
        prikey = rsa.PrivateKey.load_pkcs1(fp.read())
    # 签名
    return rsa.sign_hash(msg, prikey, "MD5")


# 通过签名计算哈希
def rsa_verify(sign: bytes, user_name: str):
    # 读取公钥
    with open(user_name + "_pub.pem", "rb") as fp:
        pubkey = rsa.PublicKey.load_pkcs1(fp.read())
    key_length = common.byte_size(pubkey.n)
    encrypted = transform.bytes2int(sign)
    decrypted = core.decrypt_int(encrypted, pubkey.e, pubkey.n)
    return transform.int2bytes(decrypted, key_length)[-32:]


# 计算散列值
def cal_md5(data: bytes) -> str:
    h = hashlib.md5()
    h.update(data)
    return h.hexdigest()


# 生成DES密钥
def generate_des_key() -> bytes:
    return os.urandom(8)


# 使用DES加密
def des_encrypt(des_key: bytes, data: bytes):
    # 生成DES对象，使用ECB加密模式
    des = DES.new(des_key, DES.MODE_ECB)
    padded_text = pad(data)
    return des.encrypt(padded_text)


# 使用DES解密
def des_decrypt(des_key: bytes, data: bytes):
    # 生成DES对象，使用ECB加密模式
    des = DES.new(des_key, DES.MODE_ECB)
    return des.decrypt(data)


# 如果需要加密的内容不是16的倍数，则需要补足
def pad(text: bytes):
    while len(text) % 8 != 0:
        text += b" "
    return text