import base64

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher
from Crypto.Cipher import AES
from functools import wraps
from utils.http import make_json_response
import json

IV = '16-Bytes--String'

with open('public.rsa') as f:
    key = f.read()
    public_key = RSA.import_key(key)
    public_cipher = PKCS1_cipher.new(public_key)
    print(public_key.exportKey().decode(encoding='utf-8'))

with open('private.rsa') as f:
    key = f.read()
    private_key = RSA.import_key(key)
    private_cipher = PKCS1_cipher.new(private_key)
    # print(private_key.exportKey().decode(encoding='utf-8'))


# 安全传输decorator
def secure_transport(view_func):
    @wraps(view_func)
    def _wrapped_view(request, *args, **kwargs):
        data = request.POST
        enc_key = data.get('enc_key')
        cipher_text = data.get('cipher_text')
        if not enc_key or not cipher_text:
            print('无加密')
            return view_func(request, *args, **kwargs)
        aes_key = private_cipher.decrypt(base64.b64decode(enc_key.encode('utf-8')), b'error').decode('utf-8')
        print(f'key={aes_key}')

        aes_cipher = AES.new(aes_key.encode('utf-8'), AES.MODE_CBC, IV.encode('utf-8'))
        decrypted = aes_cipher.decrypt(base64.b64decode(cipher_text.encode('utf-8')))
        # print(decrypted)
        decrypted = decrypted[:-decrypted[-1]]
        # print(decrypted)
        plain_text = decrypted.decode('utf-8')
        print(plain_text)

        try:
            loaded = json.loads(plain_text)
        except:
            print('不是json')
            loaded = {}
            for p in map(lambda s: s.split('='), plain_text.split('&')):
                loaded[p[0]] = p[1]
        dec_request = request
        dec_request.POST = {**request.POST, **loaded}

        raw_response = view_func(dec_request, *args, **kwargs)

        content = json.dumps({'data': raw_response.content.decode('utf-8')})
        padding = 16 - len(content) % 16
        content += bytes([padding] * padding)
        print(content)

        aes_cipher = AES.new(aes_key.encode('utf-8'), AES.MODE_CBC, IV.encode('utf-8'))
        enc_content = base64.b64encode(aes_cipher.encrypt(content)).decode('utf-8')
        print(enc_content)
        return make_json_response(enc_content=enc_content)
    return _wrapped_view


def test():
    plain_text = '{"username": "user1"}'
    encrypted = public_cipher.encrypt(bytes(plain_text.encode('utf8')))
    cipher_text = base64.b64encode(encrypted)
    print(cipher_text.decode('utf8'))

    decrypted = private_cipher.decrypt(base64.b64decode(cipher_text), b'error')
    print(decrypted.decode('utf8'))

# print(public_key.exportKey().decode('utf-8'))

# random_generator = Random.new().read
# rsa = RSA.generate(2048, random_generator)

# private_key = rsa.exportKey()
# print(private_key.decode('utf-8'))