# -*- coding: utf-8 -*-
# Author: areful

import base64

import rsa.common
from Crypto.PublicKey import RSA


def rsa_encrypt_bytes(pub_key, bytes_str):
    if not isinstance(bytes_str, bytes):
        return None

    pubkey = rsa.PublicKey(pub_key.n, pub_key.e)
    key_length = rsa.common.byte_size(pub_key.n)
    max_msg_length = key_length - 11
    count = len(bytes_str) // max_msg_length
    if len(bytes_str) % max_msg_length > 0:
        count = count + 1
    cry_bytes = b''

    # rsa加密要以max_msg_length分组, 每组分别加密(加密的数据长度为key_length, 解密时每key_length长度解密然后相加)
    for i in range(count):
        start = max_msg_length * i
        size = max_msg_length
        content = bytes_str[start: start + size]
        # rsa 分组 加密
        crypto = rsa.encrypt(content, pubkey)
        cry_bytes = cry_bytes + crypto
    return cry_bytes


# rsa 解密, bytes_string是rsa_encrypt_hex, rsa_encrypt_bytes的返回值
def rsa_decrypt(pri_key, bytes_string):
    # 导入rsa库
    import rsa.common
    pri_key = rsa.PrivateKey(pri_key.n, pri_key.e, pri_key.d, pri_key.p, pri_key.q)
    key_length = rsa.common.byte_size(pri_key.n)
    if len(bytes_string) % key_length != 0:
        # 如果数据长度不是key_length的整数倍, 则数据是无效的
        return None

    count = len(bytes_string) // key_length
    d_cty_bytes = b''

    # 分组解密
    for i in range(count):
        start = key_length * i
        size = key_length
        content = bytes_string[start: start + size]
        # rsa 分组 解密
        d_crypto = rsa.decrypt(content, pri_key)
        d_cty_bytes = d_cty_bytes + d_crypto
    return d_cty_bytes


# rsa 加密, 注意: 这里是传递的是16进制字符串
def rsa_encrypt_hex(pub_key, hex_string):
    return rsa_encrypt_bytes(pub_key, bytes.fromhex(hex_string))


# rsa 库的测试
def test_encrypt_decrypt():
    # 产生公钥私钥
    (pub, pri) = rsa.newkeys(256)
    # 构建新的公钥私钥
    pub_key = rsa.PublicKey(pri.n, pri.e)
    pri_key = rsa.PrivateKey(pri.n, pri.e, pri.d, pri.p, pri.q)

    message = b'\x00\x00\x00\x00\x01'
    # 加密 message
    crypto = rsa.encrypt(message, pub_key)
    # 解密
    d_crypto = rsa.decrypt(crypto, pri_key)
    print(d_crypto)


def read_keys_from_pem_file(pub_key_file, pri_key_file):
    with open(pub_key_file, 'r') as f:
        _key = f.read()
        _pub_key = RSA.importKey(_key)  # 导入读取到的公钥
    with open(pri_key_file, 'r') as f:
        _key = f.read()
        _pri_key = RSA.importKey(_key)  # 导入读取到的私钥
        return _pub_key, _pri_key


if __name__ == '__main__':
    # pubKey, priKey = rsa.newkeys(256)
    pubKey, priKey = read_keys_from_pem_file('rsa_pubkey.pem', 'rsa_private_key.pem')

    bts_str = ('hello world 中文 hello world 中文\n' * 30).encode()
    crypto_bytes = rsa_encrypt_bytes(pubKey, bts_str)
    cipher_text = base64.b64encode(crypto_bytes)

    encrypted_bytes = base64.b64decode(cipher_text)
    d_crypto_bytes = rsa_decrypt(priKey, encrypted_bytes)
    print(d_crypto_bytes.decode())

    hex_str = '001122334455AAff' * 40
    crypto_bytes = rsa_encrypt_hex(pubKey, hex_str)
    d_crypto_bytes = rsa_decrypt(priKey, crypto_bytes)
    print(d_crypto_bytes.hex())