# -*- coding: utf-8 -*-
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA256
from Crypto.Cipher import PKCS1_OAEP
import base64

from ddans.data.color import DColor
from ddans.native.log import NLog

RSAPublicKey = (
    """MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEApiLWsPmoaiGdr02AmPYS
4/VgxyJNNJSK2VICirXl4P8qAP3b/6GIs7Hwnsmmuj7gledJ0XiZodFjDNgQEZEG
j4K9KX8NixcrJwnTQV5UebiERdj+64ihwqPVWo9dIOj43ZPJ8f4dritJv5tnQaBb
b9zTkU3Ofzk1g8RU00Ith6quOGZcU2IWC4OQ+5YiZvzrCBKEvjFwK2OvGq43tI51
GaZdI9H09jisypqefsMvv6esjEQQOtl+WGlI5fof7KZUt3dppvuDI4CYgeXIoKKL
9T9qDl8F9aNSs3XHx32YyA7L1FraFD6PjqRjueIFp2emsXQcIMeMqqgVyV+4qb+0
uwIDAQAB""")

RSAPrivateKey = (
    """MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCmItaw+ahqIZ2v
TYCY9hLj9WDHIk00lIrZUgKKteXg/yoA/dv/oYizsfCeyaa6PuCV50nReJmh0WMM
2BARkQaPgr0pfw2LFysnCdNBXlR5uIRF2P7riKHCo9Vaj10g6Pjdk8nx/h2uK0m/
m2dBoFtv3NORTc5/OTWDxFTTQi2Hqq44ZlxTYhYLg5D7liJm/OsIEoS+MXArY68a
rje0jnUZpl0j0fT2OKzKmp5+wy+/p6yMRBA62X5YaUjl+h/splS3d2mm+4MjgJiB
5cigoov1P2oOXwX1o1KzdcfHfZjIDsvUWtoUPo+OpGO54gWnZ6axdBwgx4yqqBXJ
X7ipv7S7AgMBAAECggEAE1kuWLkSsJ2OoDpt9iJ+a7cJmNd5V1zPWOTwr9fkWFJP
QYNPKlPVOFxKUivY84rcHAkuMyuQ9OQOXZISOQGDWpZW0mzeFmtR1r+Rr2S9gNmI
6huOAL9OW306HUxiA7GuPhv2omy+Zyjac2q1WVLEI2B78Um2+WEfxuNIH7xu0gRR
z/wluevMBQlv6hSeMESP6e9s2kB+FuM4bL7+1sexwS1TudKbHybujaiB8un9E2V3
C9BkE6yf00Rhvii6Q4E4w3PFQfcBEFnjOD0S2Watx8GIbYjlYvtquy8Yj7OtyZB6
w5sxyaYhKVVnkpGAbOgaCj94ABYF47VuwXr45I4PzQKBgQDp6svzD+QKlARjcWgO
rH7UEpFX6m49G6Ib4c63xEg9+GesTgfDxoYxtXDKsC7I8OwMrmffJGbyfjaAvzmg
icM3Y7kF6UPO9+q2mBRsyhoJSNONTeUVyoFSc+QDsVgndlTZIjqC5XLrONGaSPLJ
7jJzqKhp3NQcpO8+xY0d7Efm1QKBgQC10fCNuF1LUTYgsWvvGxTiCCAIBurZmHU+
GMb2Okfoco/P+QNl1jEba91AC7FudW3g/lRn0vTnDajlNoz2/CKunrTKN4/9lEBI
tiwEOvhc5FtapRifSfrrNwFqPM+LijvOkcteWVhAUWCEP1eLmpnXGqfmZxnTaaVe
5DyDgyAVTwKBgQDhEawfOQEANJ9sV2f6BnsZ7quuXcCoR3hZP5G5O8EO735BSesq
sk/YcG4QeCjr6vQU15LWH552YH5srpHyJ9uP5gHW8DTwZZaPDhFJA8S9fjcfVP3B
U9HO1fdknjTl83N8YPTPkZpbPZZfzyTcqrYclwFEqtfkn8YbLdW0Hu6n6QKBgDEv
jvS3uD1ExfDILlurptWSvx+Mx4F/1c+roj8/+u3t9+uaAIpPDE/PuTW3NBOens/r
3WeM0QWkU1lWjYHPKY61jlDaARDnRxCkf3nRfhibCf+Wg/erar6wzagwtTCZf55d
IHfwkYgsu4BrJkK02wMuydbxczIW9U6bTaeCFG7PAoGAN3oHYzTEBwZ7ZKv11QiQ
Ep/HeHBvD8fH/e1uVyBel4Co7NO5pV2s8qjS8+3+jxumKLNXCirjqwHo03YCdzxi
cy4GFDsukHmDQ7FF9VC4QV0fJ0F7jGYE9nNZ+K9qwS5KrihOSBPxj0HptrF6TSOw
BTSwr0byL3TC0m0StkSNKpo=""")


class NRsa:

    @staticmethod
    def rsa_encrypt(source: str, public_pem: str):
        cipher = PKCS1_OAEP.new(key=RSA.import_key(public_pem),
                                hashAlgo=SHA256)
        cipherText = cipher.encrypt(source.encode('utf-8'))
        return base64.b64encode(cipherText).decode('utf-8')

    @staticmethod
    def rsa_decrypt(encrpyted: str, private_pem: str):
        rsaDecrypt = PKCS1_OAEP.new(key=RSA.import_key(private_pem),
                                    hashAlgo=SHA256)
        return rsaDecrypt.decrypt(base64.b64decode(
            encrpyted.encode('utf-8'))).decode('utf-8')

    @staticmethod
    def encode(source: str, publicKey: str = RSAPublicKey):
        _key = NRsa.format_pem_key(publicKey)
        cipher = PKCS1_OAEP.new(key=RSA.import_key(_key), hashAlgo=SHA256)

        dataBytes = source.encode('utf-8')
        encrypted_chunks = []
        chunk_size = 190
        # 按 max_chunk_size 对消息进行分块
        for i in range(0, len(dataBytes), chunk_size):
            chunk = dataBytes[i:i + chunk_size]

            # 加密每一个块
            encrypted_chunk = cipher.encrypt(chunk)

            # 添加加密后的块到列表中
            encrypted_chunks.append(encrypted_chunk)

        # 将所有加密块拼接在一起
        encrypted = b''.join(encrypted_chunks)
        return base64.b64encode(encrypted).decode('utf-8')

    @staticmethod
    def decode(encrpyted: str, privateKey: str = RSAPrivateKey):
        _key = NRsa.format_pem_key(privateKey, "RSA PRIVATE KEY")
        rsaDecrypt = PKCS1_OAEP.new(key=RSA.import_key(_key), hashAlgo=SHA256)

        dataBytes = base64.b64decode(encrpyted.encode('utf-8'))

        chunk_size = 256  # RSA 2048位的块大小
        decrypted_chunks = []

        # 分块处理
        for i in range(0, len(dataBytes), chunk_size):
            chunk = dataBytes[i:i + chunk_size]

            # 解密每一个块
            decrypted_chunk = rsaDecrypt.decrypt(chunk)

            # 将解密后的数据存储
            decrypted_chunks.append(decrypted_chunk)

        # 将解密后的块拼接成完整的字节数组
        decrypted = b''.join(decrypted_chunks)
        return decrypted.decode('utf-8')

    @staticmethod
    def format_pem_key(key_str: str, label: str = 'PUBLIC KEY'):
        # 确保每行64字符
        formatted_key = '\n'.join(
            [key_str[i:i + 64] for i in range(0, len(key_str), 64)])
        # 添加头尾
        pem_key = (f"-----BEGIN {label}-----\n"
                   f"{formatted_key}"
                   f"\n-----END {label}-----")
        return pem_key

    @staticmethod
    def remove_pem(pem: str, label: str = 'PUBLIC KEY'):
        return pem.replace(f"-----BEGIN {label}-----\n",
                           "").replace(f"\n-----END {label}-----",
                                       "").replace("\n", "")

    @staticmethod
    def generate_rsa_key():
        try:
            # Generating private key (RsaKey object) of key length of 1024 bits
            private_key = RSA.generate(2048)
            # Generating the public key (RsaKey object) from the private key
            public_key = private_key.publickey()
            # Converting the RsaKey objects to string
            private_pem = private_key.export_key().decode()
            # -----BEGIN RSA PRIVATE KEY-----
            # xxx
            # -----END RSA PRIVATE KEY-----
            public_pem = public_key.export_key().decode()
            # -----BEGIN PUBLIC KEY-----
            # xxx
            # -----END PUBLIC KEY-----
            private_pem_stripped = NRsa.remove_pem(private_pem,
                                                   "RSA PRIVATE KEY")

            # 去掉公钥的 "BEGIN" 和 "END" 部分
            public_pem_stripped = NRsa.remove_pem(private_pem)
            # 返回包含公钥和私钥的字典对象
            return {
                "private_key": private_pem_stripped,
                "public_key": public_pem_stripped,
                "public_pem": public_pem,
                "private_pem": private_pem,
            }
        except Exception as e:
            NLog.error(f'generate_rsa_key failed! {e}')

    @staticmethod
    def parse_rsa_public_key(private_pem: str):
        try:
            rsa_key = RSA.import_key(private_pem)
            public_key = rsa_key.publickey()
            public_pem = public_key.export_key().decode()
            NLog.printf(f'{public_pem}', fc=DColor.green)
        except Exception as e:
            NLog.error(f'parse_rsa_public_key failed! {e}')
