from Crypto import Random
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5
from Crypto.Hash import SHA
from Crypto.Signature import PKCS1_v1_5 as Signature_pkcs1_v1_5
from binascii import a2b_base64
from Crypto.Util.asn1 import DerSequence
import base64, os, json
from base64 import b64decode
from rsa import PublicKey, common, transform, core

##使用私鑰分段解密
def privKey_Decrypt(cipher_text, key_path="priv_path", default_length=128):
    with open(selection_key[key_path]) as f:
        key = f.read()
        rsa_key = RSA.importKey(key)
        cipher = Cipher_pkcs1_v1_5.new(rsa_key)
        data = base64.b64decode(cipher_text.encode())
        length = len(data)
        if length < default_length:
            decrypt_text = cipher.decrypt(data, 'ERROR')
            return decrypt_text
        res = []
        offset = 0
        while length - offset > 0:
            if length - offset > default_length:
                res.append(cipher.decrypt(data[offset:offset + default_length], "ERROR"))

        else:
            res.append(cipher.decrypt(data[offset:], "ERROR"))
        offset += default_length
    cipher_text = b""
    for r in res:
        cipher_text += r
    try:
        decrypt_text = json.loads(cipher_text.decode())
    except Exception as e:
        decrypt_text = cipher_text.decode()
    return str(decrypt_text) if type(decrypt_text) == int else decrypt_text

##使用公鑰分段加密
def pubKey_Encrypt(cipher_text, key_path="pub_path", default_length=100):
    """
    單次加密串的長度最大為 (key_size/8)-11
    1024bit的證書用100， 2048bit的證書用 200
    """
    cipher_text = str(cipher_text)
    with open(selection_key[key_path], "r") as f:
        key = f.read()
        rsa_key = RSA.importKey(key)
        cipher = Cipher_pkcs1_v1_5.new(rsa_key)
        data = cipher_text.encode()
        length = len(data)
        if length < default_length:
            cipher_text = base64.b64encode(cipher.encrypt(data)).decode()
            return cipher_text

    res = []
    for i in range(0, len(data), default_length):
        res.append(cipher.encrypt(data[i:i + default_length]))
    cipher_text = b''
    for res_c in res:
        cipher_text = cipher_text + res_c
    cipher_text = base64.b64encode(cipher_text).decode()
    return cipher_text

#使用私鑰進行簽名
def generate_signature(data, pri_key_path='priv_key'):
    with open(selection_key.get(pri_key_path)) as f:
        key = f.read()
        rsa_key = RSA.importKey(key)
        signer = Signature_pkcs1_v1_5.new(rsa_key)
        digest = SHA.new(data.encode())
        sign = signer.sign(digest)
        signature = base64.b64encode(sign).decode()
        return signature

#使用公鑰進行驗籤
def verify_sign(result, pub_key_path='pub_key'):
    sign = result.get("sign").encode()
    data = str(result.get("data")).encode()
    with open(selection_key.get(pub_key_path)) as f:
        key = f.read()
    rsa_key = RSA.importKey(key)
    signer = Signature_pkcs1_v1_5.new(rsa_key)
    digest = SHA.new(data)
    is_verify = signer.verify(digest, base64.b64decode(sign))
    return is_verify

#使用公鑰進行解密
def decrypt_with_public_key(cipher, pub_key_path='pub_key'):
    key = open(selection_key.get(pub_key_path), "rb").read()
    #key = b'-----BEGIN PUBLIC KEY-----...........-----END PUBLIC KEY-----\n'
    pk = PublicKey.load_pkcs1_openssl_pem(key)
    encrypted = transform.bytes2int(b64decode(cipher))
    decrypted = core.decrypt_int(encrypted, pk.e, pk.n)
    text = transform.int2bytes(decrypted)

    if len(text) > 0 and text[0] == 1:
        pos = text.find(b'\x00')
        if pos > 0:
            return json.loads(str(text[pos+1:], 'utf-8'))
        else:
            return None