# -*- coding:utf-8 -*-
# !/usr/bin/env python 3.9.11
"""
@File    :  crypto.py
@Time    :  2020/9/25 19:55
@Author  :  YuYanQing
@Version :  1.0
@Contact :  mryu168@163.com
@License :  (C)Copyright 2022-2026
@Desc    :  None
"""

import base64
import copy
import logging
import os
import struct
from abc import abstractmethod

from Crypto import Random
from Crypto.Cipher import AES, PKCS1_OAEP, PKCS1_v1_5
from Crypto.PublicKey import RSA
from Crypto.Util import Counter
from six import text_type

from .streambody import StreamBody

logger = logging.getLogger(__name__)

_AES_CTR_COUNTER_BITS_LENGTH = 8 * 16
_AES_256_KEY_SIZE = 32
__all__ = ["BaseProvider", "AESProvider", "RSAProvider", "AESCTRCipher", "DataDecryptAdapter", "DataEncryptAdapter"]


def to_bytes(str_):
    """
    将字符串转为bytes
    Args:
        str_:

    Returns:

    """
    if isinstance(str_, text_type):
        try:
            return str_.encode("utf-8")
        except UnicodeEncodeError:
            detail = "your unicode strings can not encoded in utf8, utf8 support only!"
            raise CryptoException(detail)
    return str_


def random_key(key_len):
    """

    Args:
        key_len:

    Returns:

    """
    return Random.new().read(key_len)


def random_iv():
    """

    Returns:

    """
    iv = Random.new().read(16)
    return iv


def iv_to_big_int(iv):
    """

    Args:
        iv:

    Returns:

    """
    iv_pair = struct.unpack(">QQ", iv)
    iv_int = iv_pair[0] << 64 | iv_pair[1]
    return iv_int


class CryptoException(Exception):
    def __init__(self, message):
        self._message = message

    def __str__(self):
        return str(self._message)


class AESCTRCipher(object):
    """数据加密类,用于加密用户数据"""

    def __init__(self):
        """初始化"""
        self.__block_size_len = AES.block_size
        self.__cipher = None
        self.__key_len = _AES_256_KEY_SIZE

    def new_cipher(self, key, start, offset=0):
        """
        初始化AES加解密对象
        Args:
            key: 对称密钥
            start: 对称加密初始随机值
            offset: 主要用于解密,想要解密文本的偏移,必须为16的整数倍

        Returns:
        """
        block_index_offset = self.__calc_offset(offset)
        new_start = start + block_index_offset
        my_counter = Counter.new(_AES_CTR_COUNTER_BITS_LENGTH, initial_value=new_start)
        self.__cipher = AES.new(key, AES.MODE_CTR, counter=my_counter)

    def encrypt(self, plaintext):
        """
        需要加密的数据
        Args:
            plaintext:

        Returns:

        """
        if self.__cipher is None:
            raise CryptoException("cipher is not initialized")
        return self.__cipher.encrypt(plaintext)

    def decrypt(self, plaintext):
        """
        需要解密的数据,数据的起始位置必须为16的整数倍
        Args:
            plaintext:

        Returns:

        """
        if self.__cipher is None:
            raise CryptoException("cipher is not initialized")
        return self.__cipher.decrypt(plaintext)

    # offset 必须为block_size的整数倍
    def __calc_offset(self, offset):
        """
        通过文本偏移计算counter的偏移
        Args:
            offset:

        Returns:

        """
        if not self.__is_block_aligned(offset):
            raise CryptoException("offset is not align to encrypt block")
        return offset // self.__block_size_len

    def __is_block_aligned(self, offset):
        """
        判断文本偏移是否是block_size对齐
        Args:
            offset:

        Returns:

        """
        if offset is None:
            offset = 0
        return offset % self.__block_size_len == 0

    def adjust_read_offset(self, offset):
        """
        用于调整读取的offset为block_size对齐
        Args:
            offset:

        Returns:

        """

        if offset:
            offset = (offset // self.__block_size_len) * self.__block_size_len
        return offset

    def get_key(self):
        """获取随机密钥"""
        return random_key(self.__key_len)

    def get_counter_iv(self):
        """获取对称加密初始随机值"""
        return random_iv()


class RSAKeyPair:
    """封装有公钥和私钥"""

    def __init__(self, public_key, private_key):
        self.public_key = public_key
        self.private_key = private_key


class RSAKeyPairPath:
    """封装有公钥和私钥的路径"""

    def __init__(self, public_key_path, private_key_path):
        self.public_key_path = public_key_path
        self.private_key_path = private_key_path


class BaseProvider(object):
    """客户端主密钥加密基类"""

    def __init__(self, cipher):
        """初始化
        :param cipher(an AES object): 数据加解密类
        """
        self.data_cipher = cipher

    def get_data_key(self):
        """随机获取数据加解密密钥"""
        return self.data_cipher.get_key()

    @abstractmethod
    def init_data_cipher(self):
        """初始化cipher"""

    @abstractmethod
    def init_data_cipher_by_user(self, encrypt_key, encrypt_iv, offset=0):
        """根据密钥初始化cipher"""

    def adjust_read_offset(self, start):
        """用于调整读取的offset为block_size对齐"""
        return self.data_cipher.adjust_read_offset(start)

    def make_data_encrypt_adapter(self, stream):
        """创建数据流加密适配器"""
        size = 0
        if hasattr(stream, "__len__"):
            size = len(stream)
        elif hasattr(stream, "tell") and hasattr(stream, "seek"):
            current = stream.tell()
            stream.seek(0, os.SEEK_END)
            size = stream.tell()
            stream.seek(current, os.SEEK_SET)
        else:
            return None
        return DataEncryptAdapter(stream, size, copy.copy(self.data_cipher))

    def make_data_decrypt_adapter(self, rt, offset):
        """创建数据流解密适配器"""
        return DataDecryptAdapter(rt, copy.copy(self.data_cipher), offset)


class RSAProvider(BaseProvider):
    """客户端非对称主密钥加密类"""

    def __init__(self, key_pair_info=None, cipher=AESCTRCipher(), passphrase=None):
        """初始化"""
        super(RSAProvider, self).__init__(cipher=cipher)

        default_rsa_dir = os.path.expanduser("~/.local_rsa")
        default_public_key_path = os.path.join(default_rsa_dir, ".public_key.pem")
        default_private_key_path = os.path.join(default_rsa_dir, ".private_key.pem")
        self.__encrypt_obj = None
        self.__decrypt_obj = None
        self.__data_key = None
        self.__data_iv = None

        if isinstance(key_pair_info, RSAKeyPair):
            self.__encrypt_obj = PKCS1_v1_5.new(RSA.importKey(key_pair_info.public_key, passphrase=passphrase))
            self.__decrypt_obj = PKCS1_v1_5.new(RSA.importKey(key_pair_info.private_key, passphrase=passphrase))
        elif isinstance(key_pair_info, RSAKeyPairPath):
            self.__encrypt_obj, self.__decrypt_obj = self.__get_key_by_path(
                key_pair_info.public_key_path,
                key_pair_info.private_key_path,
                passphrase,
            )
        else:
            logger.info("key_pair_info is None, try to get key from default path")
            if os.path.exists(default_private_key_path) and os.path.exists(default_public_key_path):
                self.__encrypt_obj, self.__decrypt_obj = self.__get_key_by_path(
                    default_public_key_path,
                    default_private_key_path,
                    passphrase,
                )

        # 为用户自动创建rsa
        if self.__encrypt_obj is None and self.__decrypt_obj is None:
            logger.warning("fail to get rsa key, will generate key")
            private_key = RSA.generate(2048)
            public_key = private_key.publickey()

            self.__encrypt_obj = PKCS1_OAEP.new(public_key)
            self.__decrypt_obj = PKCS1_OAEP.new(private_key)

            if not os.path.exists(default_rsa_dir):
                os.makedirs(default_rsa_dir)

            with open(default_private_key_path, "wb") as f:
                f.write(private_key.exportKey(passphrase=passphrase))

            with open(default_public_key_path, "wb") as f:
                f.write(public_key.exportKey(passphrase=passphrase))

    @staticmethod
    def get_rsa_key_pair(public_key, private_key):
        """开放给用户,用于生成RSAKeyPair"""
        if public_key is None or private_key is None:
            raise CryptoException("public_key or private_key is not allowed to be None !!!")
        return RSAKeyPair(public_key, private_key)

    @staticmethod
    def get_rsa_key_pair_path(public_key_path, private_key_path):
        """开放给用户,用于生成RSAKeyPairPath"""
        if public_key_path is None or private_key_path is None:
            raise CryptoException("public_key or private_key is not allowed to be None !!!")
        return RSAKeyPairPath(public_key_path, private_key_path)

    def __get_key_by_path(self, public_path=None, private_path=None, passphrase=None):
        """用于从提供的路径中获取公钥和私钥"""
        if public_path is None or private_path is None:
            return None, None

        encrypt_obj, decrypt_obj = None, None

        if os.path.exists(public_path) and os.path.exists(private_path):
            with open(public_path, "rb") as f:
                encrypt_obj = PKCS1_OAEP.new(RSA.importKey(f.read(), passphrase=passphrase))

            with open(private_path, "rb") as f:
                decrypt_obj = PKCS1_OAEP.new(RSA.importKey(f.read(), passphrase=passphrase))

        return encrypt_obj, decrypt_obj

    def init_data_cipher(self):
        """初始化cipher"""
        self.__data_key = self.get_data_key()
        self.__data_iv = self.data_cipher.get_counter_iv()
        start = iv_to_big_int(self.__data_iv)
        self.data_cipher.new_cipher(self.__data_key, start)
        encrypt_key = self.__encrypt_obj.encrypt(self.__data_key)
        encrypt_iv = self.__encrypt_obj.encrypt(self.__data_iv)
        return encrypt_key, encrypt_iv

    def init_data_cipher_by_user(self, encrypt_key, encrypt_iv, offset=0):
        """根据密钥初始化cipher"""
        self.__data_key = self.__decrypt_obj.decrypt(encrypt_key)
        self.__data_iv = self.__decrypt_obj.decrypt(encrypt_iv)
        start = iv_to_big_int(self.__data_iv)
        self.data_cipher.new_cipher(self.__data_key, start, offset)


class AESProvider(BaseProvider):
    """客户端对称主密钥加密类"""

    def __init__(self, aes_key=None, aes_key_path=None, cipher=AESCTRCipher()):
        """初始化"""
        super(AESProvider, self).__init__(cipher=cipher)
        self.__ed_obj = None
        self.__data_key = None
        self.__data_iv = None
        self.__my_counter = Counter.new(_AES_CTR_COUNTER_BITS_LENGTH, initial_value=0)
        self.__aes_key = aes_key
        self.__aes_key_path = aes_key_path

    def init_ed_obj(self):
        default_aes_dir = os.path.expanduser("~/.local_aes")
        default_key_path = os.path.join(default_aes_dir, ".aes_key.pem")
        if self.__aes_key:
            aes_key = to_bytes(base64.b64decode(to_bytes(self.__aes_key)))
            self.__ed_obj = AES.new(aes_key, AES.MODE_CTR, counter=self.__my_counter)
        elif self.__aes_key_path:
            if os.path.exists(self.__aes_key_path):
                with open(self.__aes_key_path, "rb") as f:
                    aes_key = f.read()
                    aes_key = to_bytes(base64.b64decode(to_bytes(aes_key)))
                    self.__ed_obj = AES.new(aes_key, AES.MODE_CTR, counter=self.__my_counter)
        else:
            logger.info("aes_key and aes_key_path is None, try to get key from default path")
            if os.path.exists(default_key_path):
                with open(default_key_path, "rb") as f:
                    aes_key = f.read()
                    aes_key = to_bytes(base64.b64decode(to_bytes(aes_key)))
                    self.__ed_obj = AES.new(aes_key, AES.MODE_CTR, counter=self.__my_counter)

        if self.__ed_obj is None:
            logger.warning("fail to get aes key, will generate key")
            aes_key = random_key(_AES_256_KEY_SIZE)
            self.__ed_obj = AES.new(aes_key, AES.MODE_CTR, counter=self.__my_counter)
            if not os.path.exists(default_aes_dir):
                os.makedirs(default_aes_dir)

            with open(default_key_path, "wb") as f:
                aes_key = to_bytes(base64.b64encode(to_bytes(aes_key)))
                f.write(aes_key)

    def init_data_cipher(self):
        """初始化cipher"""
        self.__data_key = self.get_data_key()
        self.__data_iv = self.data_cipher.get_counter_iv()
        start = iv_to_big_int(self.__data_iv)
        self.data_cipher.new_cipher(self.__data_key, start)

        self.init_ed_obj()
        encrypt_key = self.__ed_obj.encrypt(self.__data_key)
        encrypt_iv = self.__ed_obj.encrypt(self.__data_iv)
        return encrypt_key, encrypt_iv

    def init_data_cipher_by_user(self, encrypt_key, encrypt_iv, offset=0):
        """根据密钥初始化cipher"""
        self.init_ed_obj()
        self.__data_key = self.__ed_obj.decrypt(encrypt_key)
        self.__data_iv = self.__ed_obj.decrypt(encrypt_iv)
        start = iv_to_big_int(self.__data_iv)
        self.data_cipher.new_cipher(self.__data_key, start, offset)


class DataEncryptAdapter(object):
    """用于读取经过加密后的的数据"""

    def __init__(self, data, content_len, data_cipher):
        self._data = to_bytes(data)
        self._data_cipher = data_cipher
        self._content_len = content_len
        self._read_len = 0

    @property
    def len(self):
        return self._content_len

    def read(self, length):
        """读取加密后的数据"""
        if self._read_len >= self._content_len:
            return ""

        if length is None or length < 0:
            bytes_to_read = self._content_len - self._read_len
        else:
            bytes_to_read = min(length, self._content_len - self._read_len)

        if isinstance(self._data, bytes):
            content = self._data[self._read_len : self._read_len + bytes_to_read]
        else:
            content = self._data.read(bytes_to_read)

        self._read_len += bytes_to_read
        content = self._data_cipher.encrypt(content)
        return content


class DataDecryptAdapter(StreamBody):
    """用于读取经过解密后的数据"""

    def __init__(self, rt, data_cipher, offset=0):
        """初始化
        :param rt(request object): request请求返回的对象
        :param data_cipher(an AES object): 数据加解密类
        :param offset(int): 数据读取的偏移量
        """
        super(DataDecryptAdapter, self).__init__(rt)
        self._data_cipher = data_cipher
        self._offset = offset

    def read(self, length, auto_decompress=False):
        """读取解密后的数据"""
        if self._read_len >= self._content_len:
            return ""

        if self._use_encoding and not auto_decompress:
            content = self._rt.raw.read(length)
        else:
            try:
                content = next(self._rt.iter_content(length))
            except StopIteration:
                return ""

        content = self._data_cipher.decrypt(content)
        if self._read_len < self._offset < self._read_len + len(content):
            content = content[self._offset :]
            self._read_len = self._offset
        return content
