import sys
import subprocess
import os
import argparse
import hashlib
from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from datetime import datetime, timedelta, timezone
from cryptography.x509.oid import ObjectIdentifier
from pyasn1.type import univ, namedtype
from pyasn1.codec.der import encoder as der_encoder
from cryptography.x509 import load_pem_x509_certificate
import binascii
import base64
import requests


class MyFileManager:
    @staticmethod
    def read_file(file_path, description):
        try:
            with open(file_path, 'rb') as f:
                return f.read()
        except FileNotFoundError:
            print(f"Error: {description} file not found - {file_path}")
        except Exception as e:
            print(f"Error: Failed to read {description} file - {e}")
        MySystemUtils.custom_exit(1)

    @staticmethod
    def write_file(file_path, data, description):
        try:
            if os.path.isdir(file_path):
                print(f"Error: {file_path} is a directory, expected a file path.")
                MySystemUtils.custom_exit(1)

            if os.path.dirname(file_path):
                os.makedirs(os.path.dirname(file_path), exist_ok=True)

            with open(file_path, "wb") as f:
                f.write(data)
            print(f"{description} generated successfully: {file_path}")
        except Exception as e:
            print(f"Error: Failed to write {description} - {e}")
            MySystemUtils.custom_exit(1)


class MySystemUtils:
    cert_file_path = "x509_cert.der"

    @staticmethod
    def welcome_print():
        print("=======================================================================")
        print("                       Starting sign_tool.py                           ")
        print("=======================================================================")

    @staticmethod
    def custom_exit(code=0):
        # 删除证书
        os.remove(MySystemUtils.cert_file_path) if os.path.exists(MySystemUtils.cert_file_path) else None
        print("=======================================================================")
        print("                         SUCCESS" if code == 0 else "                         FAILED")
        print("=======================================================================")
        sys.exit(code)

    @staticmethod
    def calc_sha256(file_path):
        print(f"Calculating SHA256 for file: {file_path}")
        try:
            with open(file_path, 'rb') as f:
                return hashlib.sha256(f.read()).hexdigest()
        except FileNotFoundError:
            print(f"Error: File not found - {file_path}")
        except Exception as e:
            print(f"Error: Failed to calculate SHA256 - {e}")
        MySystemUtils.custom_exit(1)
    @staticmethod
    def load_private_key(priv_key_path):
        print("Loading private key...")
        try:
            with open(priv_key_path, "rb") as key_file:
                return serialization.load_pem_private_key(key_file.read(), password=None)
        except Exception as e:
            print(f"Error: Failed to load private key - {e}")
            MySystemUtils.custom_exit(1)
    @staticmethod
    def install_missing_packages():
        """Check and install missing packages."""
        required_packages = [
            "cryptography",
            "pyasn1",
            "requests"
        ]
        for package in required_packages:
            try:
                __import__(package)
            except ImportError:
                print(f"Package '{package}' is not installed. Installing...")
                subprocess.check_call([sys.executable, "-m", "pip", "install", package])

class MyCertificateBuilder:
    def __init__(self, private_key, sha256):
        print("Initializing MyCertificateBuilder...")
        self.private_key = private_key
        self.sha256 = sha256
        self.__subject__ = self.__issuer__ = x509.Name([
            x509.NameAttribute(NameOID.COMMON_NAME, "QUECTEL_ASR_SECFOTA"),
        ])
        print("MyCertificateBuilder initialized successfully.")

    def build(self):
        print("Building certificate...")
        try:
            builder = x509.CertificateBuilder().subject_name(
                self.__subject__
            ).issuer_name(
                self.__issuer__
            ).public_key(
                self.private_key.public_key()
            ).serial_number(
                x509.random_serial_number()
            ).not_valid_before(
                datetime.now(timezone.utc)
            ).not_valid_after(
                datetime.now(timezone.utc) + timedelta(days=3650)
            )
            # .add_extension(
            #     x509.BasicConstraints(ca=False, path_length=None), critical=True
            # ).add_extension(
            #     x509.SubjectKeyIdentifier.from_public_key(self.private_key.public_key()), critical=False
            # ).add_extension(
            #     x509.AuthorityKeyIdentifier.from_issuer_public_key(self.private_key.public_key()), critical=False
            # )

            custom_der = self._build_custom_extension()
            builder = builder.add_extension(
                x509.UnrecognizedExtension(ObjectIdentifier("1.3.6.1.4.1.4128.2100.8005"), custom_der),
                critical=False
            )

            print("Certificate built successfully.")
            return builder.sign(self.private_key, hashes.SHA256())
        except Exception as e:
            print(f"Error: Failed to build certificate - {e}")
            MySystemUtils.custom_exit(1)

    def _build_custom_extension(self):
        print("Preparing custom extension value...")
        sha256_oid = "2.16.840.1.101.3.4.2.1"

        class AlgorithmIdentifier(univ.Sequence):
            componentType = namedtype.NamedTypes(
                namedtype.NamedType('algorithm', univ.ObjectIdentifier()),
                namedtype.NamedType('parameters', univ.Null())
            )

        class CustomExtensionValue(univ.Sequence):
            componentType = namedtype.NamedTypes(
                namedtype.NamedType('algorithm', AlgorithmIdentifier()),
                namedtype.NamedType('hashValue', univ.OctetString())
            )

        sha256_bin = binascii.unhexlify(self.sha256) if isinstance(self.sha256, str) else self.sha256
        custom_value = CustomExtensionValue()
        custom_value['algorithm']['algorithm'] = sha256_oid
        custom_value['algorithm']['parameters'] = univ.Null()
        custom_value['hashValue'] = sha256_bin

        print("Custom extension value prepared successfully.")
        return der_encoder.encode(custom_value)


class MyPackBuilder:
    def __init__(self, pack_path, der_cert_path):
        print(f"Initializing MyPackBuilder with pack_path: {pack_path}, der_cert_path: {der_cert_path}")
        self.cert = MyFileManager.read_file(der_cert_path, "Certificate")
        self.pack = MyFileManager.read_file(pack_path, "Original pack")
        print("MyPackBuilder initialized successfully.")

    def build(self):
        print("Building signed pack...")
        try:
            self.pack += self.cert
            cert_len = len(self.cert)
            pack_len = len(self.pack)
            print(f"Certificate length: {cert_len}, Certificate length % 4096 = {cert_len % 4096}")
            if cert_len % 4096 != 0:
                padding16_size = (16 - (pack_len % 16))%16
                print(f"Padding16 size: {padding16_size}")
                padding16_tag = b'\0' * padding16_size
                padding_size = (4096 - padding16_size) - (cert_len  % (4096 - padding16_size))
                padding_tag = b"QUEC_SECFOTA\0\0\0\0"
                padding_data = padding16_tag + padding_tag * (padding_size // len(padding_tag)) + padding_tag[:padding_size % len(padding_tag)]
                self.pack += padding_data
                print(f"Padding added: {padding_size} bytes")

            print("Signed pack built successfully.")
            return self.pack
        except Exception as e:
            print(f"Error: Failed to build pack - {e}")
            MySystemUtils.custom_exit(1)


class MyRemoteSigner:
    def __init__(self, hsm_url, hsm_req_token, sha256):
        self.hsm_url = hsm_url
        self.hsm_req_token = hsm_req_token
        self.sha256 = sha256

    def handle_remote_signing(self, orig_pack_path, sign_pack_path, cert_file_path):
        print("Building request...")
        custom_extension = MyCertificateBuilder(None, self.sha256)._build_custom_extension().hex()
        print(f"custom_extension: {custom_extension}")
        headers = {
            "Content-Type": "application/json",
            "token": self.hsm_req_token
        }
        data = {
            "oidKey": "1.3.6.1.4.1.4128.2100.8005",
            "oidValue": custom_extension
        }
        print(f"request data: {data}")
        try:
            response = requests.post(self.hsm_url, json=data, headers=headers)
            print(f"request response, data = {response.json()}")
            if response.status_code == 200:
                response_json = response.json()
                data_base64 = response_json.get("data")
                if data_base64:
                    cert_pem_data = base64.b64decode(data_base64)
                    cert_der_data = load_pem_x509_certificate(cert_pem_data).public_bytes(serialization.Encoding.DER)
                    # print(f"request cert_pem_data: {cert_pem_data}")
                    # print(f"convert cert_der_data: {cert_der_data}")
                    MyFileManager.write_file(cert_file_path, cert_der_data, "[Remote certificate]")
                    sign_pack = MyPackBuilder(orig_pack_path, cert_file_path).build()
                    MyFileManager.write_file(sign_pack_path, sign_pack, "[Signed pack]")
                else:
                    print("Error: No certificate data found in response")
                    MySystemUtils.custom_exit(1)
            else:
                print(f"Request error, status code: {response.status_code}")
                MySystemUtils.custom_exit(1)
        except requests.RequestException as e:
            print(f"Request API error: {e}")
            MySystemUtils.custom_exit(1)


class PubKey2CConverter:
    @staticmethod
    def detect_key_type(key_data):
        """检测密钥类型"""
        try:
            # Attempt to load as a public key
            serialization.load_pem_public_key(key_data)
            return "public"
        except ValueError:
            try:
                # Attempt to load as a private key
                serialization.load_pem_private_key(key_data, password=None)
                return "private"
            except ValueError:
                return "unknown"

    @staticmethod
    def extract_public_key_from_private(key_data):
        """从私钥中提取公钥"""
        try:
            priv_key = serialization.load_pem_private_key(key_data, password=None)
            return priv_key.public_key()
        except Exception as e:
            print(f"Error: Failed to extract public key from private key - {e}")
            MySystemUtils.custom_exit(1)

    @staticmethod
    def convert_pubkey_to_c(input_path, output_path, array_name):
        try:
            # Read PEM public key file
            with open(input_path, "rb") as pem_file:
                key_data = pem_file.read()

            # Detect key type
            key_type = PubKey2CConverter.detect_key_type(key_data)
            print(f"Detected key type: {key_type}")

            public_key = None
            
            if key_type == "public":
                # Directly load the public key
                public_key = serialization.load_pem_public_key(key_data)
                print("Processing public key directly")
                
            elif key_type == "private":
                # Extract the public key from the private key
                public_key = PubKey2CConverter.extract_public_key_from_private(key_data)
                print("Extracted public key from private key")
                
            elif key_type == "unknown":
                print(f"Error: The input file is neither a valid public key nor private key")
                MySystemUtils.custom_exit(1)
            else:
                print(f"Error: Unknown key type detected")
                MySystemUtils.custom_exit(1)

            # Convert the public key to DER format
            der_data = public_key.public_bytes(
                serialization.Encoding.DER,
                serialization.PublicFormat.SubjectPublicKeyInfo
            )

            # Create output header file
            with open(output_path, "w") as out_file:
                out_file.write("// Auto-generated by pubkey2c\n")
                out_file.write(f"// Key type: {key_type}\n")
                out_file.write(f"const unsigned char {array_name}[] = {{\n")

                # Write the data in hexadecimal format, with 16 bytes per line.
                for i, byte in enumerate(der_data):
                    if i % 16 == 0:
                        out_file.write("    ")
                    out_file.write(f"0x{byte:02x}")
                    if i != len(der_data) - 1:
                        out_file.write(",")
                    if i % 16 == 15 or i == len(der_data) - 1:
                        out_file.write("\n")
                    else:
                        out_file.write(" ")

                # Define the size of the array to be written into
                out_file.write(f"}};\n\n")
                out_file.write(f"const unsigned int {array_name}_len = sizeof({array_name});\n")

            print(f"Conversion successful!")
            print(f"Input:  {input_path} ({key_type} key)")
            print(f"Output: {output_path} (array '{array_name}')")
            print(f"DER data length: {len(der_data)} bytes")
            
        except FileNotFoundError:
            print(f"Error: File not found - {input_path}")
            MySystemUtils.custom_exit(1)
        except Exception as e:
            print(f"Error: Failed to convert public key to C header - {e}")
            MySystemUtils.custom_exit(1)  


# Ensure all required packages are installed before running the script
MySystemUtils.install_missing_packages()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Sign a package locally or remotely.")
    subparsers = parser.add_subparsers(dest="mode", required=True, help="Mode of operation: 'local', 'remote', or 'pubkey2c'")

    # Local signing arguments
    local_parser = subparsers.add_parser("local", help="Perform local signing")
    local_parser.add_argument("-orig", required=True, help="Path to the original package")
    local_parser.add_argument("-out", required=True, help="Path to save the output package")
    local_parser.add_argument("-priv", required=True, help="Path to the private key")

    # Remote signing arguments
    remote_parser = subparsers.add_parser("remote", help="Perform remote signing")
    remote_parser.add_argument("-orig", required=True, help="Path to the original package")
    remote_parser.add_argument("-out", required=True, help="Path to save the output package")
    remote_parser.add_argument("-url", required=True, help="HSM URL for remote signing")
    remote_parser.add_argument("-token", required=True, help="HSM request token")

    # Bin2C arguments
    pubkey2c_parser = subparsers.add_parser("pubkey2c", help="Convert PEM public key file to C header file")
    pubkey2c_parser.add_argument("-orig", required=True, help="Path to the input PEM  public key or private key file")
    pubkey2c_parser.add_argument("-out", required=True, help="Path to the output C header file")
    pubkey2c_parser.add_argument("-array", required=True, help="Name of the array in the C header file")

    args = parser.parse_args()
    MySystemUtils.welcome_print()  

    if args.mode == "local":
        orig_pack_path = args.orig
        out_pack_path = args.out
        priv_key_path = args.priv

        sha256 = MySystemUtils.calc_sha256(orig_pack_path)
        print(f"SHA256 value: {sha256}")
        priv_key = MySystemUtils.load_private_key(priv_key_path)
        x509_cert = MyCertificateBuilder(priv_key, sha256).build()
        x509_cert_der = x509_cert.public_bytes(serialization.Encoding.DER)
        MyFileManager.write_file(MySystemUtils.cert_file_path, x509_cert_der, "Certificate")
        out_pack = MyPackBuilder(orig_pack_path, MySystemUtils.cert_file_path).build()  # Changed from sign_pack to out_pack
        MyFileManager.write_file(out_pack_path, out_pack, "Output pack")  # Changed description
    elif args.mode == "remote":
        orig_pack_path = args.orig
        out_pack_path = args.out
        hsm_url = args.url
        hsm_req_token = args.token

        sha256 = MySystemUtils.calc_sha256(orig_pack_path)
        print(f"SHA256 value: {sha256}")
        MyRemoteSigner(hsm_url, hsm_req_token, sha256).handle_remote_signing(orig_pack_path, out_pack_path, MySystemUtils.cert_file_path)  # Changed sign_pack_path to out_pack_path
    elif args.mode == "pubkey2c":
        orig_pubkey_path = args.orig
        out_c_path = args.out
        array_name = args.array
        PubKey2CConverter.convert_pubkey_to_c(orig_pubkey_path, out_c_path, array_name)

    MySystemUtils.custom_exit(0)