#! /usr/bin/env python
# -*- coding: utf-8 -*-
# __author__: vv
# Date: 2019/11/7

import os
import base64
import logging
from django.conf import settings

# 摘要算法，默认为sha256(参照配置文件)
from Crypto import Random
from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA
from Crypto.Signature import PKCS1_v1_5
from Crypto.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# AES根据16位对齐
BS = 16

# 伪随机数生成器
random_generator = Random.new().read


def sign_rsa(content, private_key=None):
    """
    非对称签名
    :param content:
    :param private_key:
    :return:
    """
    if private_key:
        private_key = RSA.importKey(private_key)
    else:
        private_key = RSA.importKey(settings.YIBAO_PRIVATE_KEY)

    h = SHA256.new(content)
    signer = PKCS1_v1_5.new(private_key)
    signature = signer.sign(h)
    return encode_base64(signature)


def verify_rsa(content, signature, public_key=None, digest_alg=SHA256):
    """
    非对称验签
    :param content:
    :param signature:
    :param public_key:
    :return:
    """
    if not public_key:
        public_key = RSA.importKey(settings.YIBAO_PUBLIC_KEY)
    h = SHA256.new(content)
    verifier = PKCS1_v1_5.new(public_key)
    signature = signature.rstrip('')

    if verifier.verify(h, decode_base64(signature)):
        return True
    return False


# 封装数字信封
def encrypt(content, private_key=None, public_key=None):
    if not public_key:
        public_key = RSA.importKey(settings.YIBAO_PUBLIC_KEY)

    if not private_key:
        private_key = RSA.importKey(settings.YIBAO_PRIVATE_KEY)
    # 生成随机密钥
    random_key = get_random_key_readable(16)
    # 对数据进行签名
    sign_to_base64 = sign_rsa(content, private_key)

    # 用随机密钥对数据和签名进行加密
    cipher = AES.new(random_key, AES.MODE_ECB)
    encrypted_data_to_base64 = cipher.encrypt(base64.urlsafe_b64encode(content + '$' + decode_base64(sign_to_base64)))
    encrypted_data_to_base64 = base64.urlsafe_b64encode(encrypted_data_to_base64)
    # 对密钥加密
    cipher = Cipher_pkcs1_v1_5.new(public_key)
    encrypted_random_key_to_base64 = base64.urlsafe_b64encode(cipher.encrypt(random_key))
    cigher_text = [encrypted_random_key_to_base64]
    cigher_text.append(encrypted_data_to_base64)
    cigher_text.append('AES')
    cigher_text.append('SHA256')
    return '$'.join(cigher_text)


# 拆开数字信封
def decrypt(content, private_key=None):

    if not private_key:
        private_key = RSA.importKey(settings.YIBAO_PRIVATE_KEY)

    args = content.split('$')
    if len(args) != 4:
        raise Exception("source invalid", args)
    # 分解参数
    encrypted_random_key_to_base64 = str(args[0])
    encrypted_date_to_base64 = str(args[1])
    symmetric_encrypt_alg = args[2]
    digest_alg = args[3]

    # 用私钥对随机密钥进行解密
    cipher = Cipher_pkcs1_v1_5.new(private_key)
    random_key = cipher.decrypt(decode_base64(encrypted_random_key_to_base64), random_generator)

    cipher = AES.new(random_key, AES.MODE_ECB)
    encryped_data = cipher.decrypt(decode_base64(encrypted_date_to_base64))
    # 分解参数
    data = encryped_data.split('$')
    source_data = data[0]
    sign_to_base64 = data[1]

    return source_data, sign_to_base64


# base64解码
def decode_base64(data):
    missing_padding = 4 - len(data) % 4
    if missing_padding:
        data += b'=' * 3
    return base64.urlsafe_b64decode(data)


def encode_base64(data):
    data = base64.urlsafe_b64encode(data)
    for i in range(3):
        if data.endswith('='):
            data = data[:-1]
    return data


# 生成随机密钥
def get_random_key_readable(key_size=16):
    ulen = int(key_size / 4 * 3)
    key = base64.b64encode(os.urandom(ulen))
    return key