# -*- coding:utf-8 -*-
#############################
# pip install pycryptodome
#############################
import base64

from Crypto import Random
from Crypto.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5
from Crypto.Signature import PKCS1_v1_5 as Signature_pkcs1_v1_5
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA

# 1. 伪随机生成器
random_generator = Random.new().read
print(random_generator)
# 2. rsa实例
rsa = RSA.generate(1024, random_generator)
# 3. 生成私钥 公钥
private_key = rsa.export_key()
public_key = rsa.publickey().export_key()
print(private_key)
print(public_key)

# 4. 秘钥写入文件
with open('private_key.pem', 'wb') as f:
    f.write(private_key)

with open('public_key.pem', 'wb') as f:
    f.write(public_key)

# 5. 公钥加密
with open('public_key.pem', 'rb') as f:
    public_key = f.read()
    rsa_key = RSA.import_key(public_key)
    pkcs = Cipher_pkcs1_v1_5.new(rsa_key)
    text = base64.b64encode(pkcs.encrypt('1234'.encode(encoding='utf-8')))
    print(text)

# 6. 私钥解密
with open('private_key.pem', 'rb') as f:
    private_key = f.read()
    rsa_key = RSA.import_key(private_key)
    pkcs = Cipher_pkcs1_v1_5.new(rsa_key)
    text = pkcs.decrypt(base64.b64decode(text), random_generator).decode()
    print(text)

# 7. 签名和验签
# 7.1 私钥签名
print('私钥签名')
with open('private_key.pem') as f:
    key = f.read()
    rsa_key = RSA.importKey(key)
    signer = Signature_pkcs1_v1_5.new(rsa_key)
    digest = SHA.new()
    digest.update('签名和验签'.encode("utf8"))
    sign = signer.sign(digest)
    signature = base64.b64encode(sign)

print(signature)

# 7.2 公钥验签
print('公钥验签')
with open('public_key.pem') as f:
    key = f.read()
    rsa_key = RSA.importKey(key)
    verifier = Signature_pkcs1_v1_5.new(rsa_key)
    digest = SHA.new()
    digest.update('签名和验签'.encode("utf8"))
    is_verify = verifier.verify(digest, base64.b64decode(signature))

print(is_verify)