import os.path
import sys
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePrivateKey

def print_privkey(private_key, public_key):
    if isinstance(private_key, Ed25519PrivateKey):
        raw_private = private_key.private_bytes(
            serialization.Encoding.Raw,
            serialization.PrivateFormat.Raw,
            serialization.NoEncryption(),
        )
        raw_public = public_key.public_bytes(
            serialization.Encoding.Raw,
            serialization.PublicFormat.Raw,
        )

        print( "    let keypair_bytes = hex!(")
        print(f"        \"{raw_private.hex()}\"")
        print(f"        \"{raw_public.hex()}\"")
        print( "    );")
        print( "    let signing = ed25519_dalek::SigningKey::from_keypair_bytes(&keypair_bytes).unwrap();")
        print( "    makiko::Privkey::Ed25519(signing.into())")
    elif isinstance(private_key, RSAPrivateKey):
        def print_num(name, x):
            x = to_be_bytes(x)
            if len(x) < 40:
                print(f"    let {name} = BigUint::from_bytes_be(&hex!(\"{x.hex()}\"));")
            else:
                print(f"    let {name} = BigUint::from_bytes_be(&hex!(")
                while x:
                    chunk, x = x[:32], x[32:]
                    print(f"        \"{chunk.hex()}\"")
                print("    ));")

        numbers = private_key.private_numbers()
        print_num("n", numbers.public_numbers.n)
        print_num("e", numbers.public_numbers.e)
        print_num("d", numbers.d)
        print_num("p", numbers.p)
        print_num("q", numbers.q)

        print("    let privkey = rsa::RsaPrivateKey::from_components(n, e, d, vec![p, q]).unwrap();")
        print("    makiko::Privkey::Rsa(privkey.into())")
    elif isinstance(private_key, EllipticCurvePrivateKey):
        if private_key.curve.name == "secp256r1":
            curve = "p256::NistP256"
            variant = "EcdsaP256"
        elif private_key.curve.name == "secp384r1":
            curve = "p384::NistP384"
            variant = "EcdsaP384"
        else:
            raise NotImplementedError()

        private_numbers = private_key.private_numbers()
        private_bytes = to_be_bytes(private_numbers.private_value)

        print(f"    let secret_key = elliptic_curve::SecretKey::<{curve}>::from_slice(&hex!(")
        print(f"        \"{private_bytes.hex()}\"")
        print( "    )).unwrap();")
        print(f"    let privkey = makiko::pubkey::EcdsaPrivkey::<{curve}>::from(secret_key);")
        print(f"    makiko::Privkey::{variant}(privkey)")
    else:
        raise NotImplementedError()

def to_be_bytes(x):
    return x.to_bytes((x.bit_length() + 7) // 8, "big")

base_dir = os.path.dirname(__file__)

print(f"// auto generated by {os.path.basename(__file__)}")
print("use num_bigint_dig::BigUint;")
print("use hex_literal::hex;")
print("use makiko::elliptic_curve;")
print()

def print_key(name, password=None, decode=True, format="ssh"):
    private_file = os.path.join(base_dir, name)
    public_file = os.path.join(base_dir, f"{name}.pub")
    private_bytes = open(private_file, "rb").read()
    public_bytes = open(public_file, "rb").read()

    if decode:
        if format == "ssh":
            private_key = serialization.load_ssh_private_key(private_bytes, password)
            public_key = serialization.load_ssh_public_key(public_bytes)
        elif format == "pem":
            private_key = serialization.load_pem_private_key(private_bytes, password)
            public_key = serialization.load_pem_public_key(public_bytes)

        print(f"pub fn {name}() -> makiko::Privkey {{")
        print_privkey(private_key, public_key)
        print(f"}}")

    print_key_file(f"{name.upper()}_PRIVKEY_FILE", private_bytes)
    print_key_file(f"{name.upper()}_PUBKEY_FILE", public_bytes)

    print()

def print_key_file(name, file_bytes):
    file_str = file_bytes.decode("utf-8")
    print(f"pub static {name}: &'static str = concat!(")
    for line in file_str.splitlines(keepends=True):
        while line:
            chunk, line = line[:80], line[80:]
            escaped_chunk = chunk.translate({
                ord("\n"): "\\n",
                ord("\\"): "\\\\",
                ord("\""): "\"",
            })
            print(f"    \"{escaped_chunk}\",")
    print(");")

for name in [
    "alice_ed25519",
    "edward_ed25519",
    "ruth_rsa_1024", "ruth_rsa_2048", "ruth_rsa_4096",
    "eda_ecdsa_p256", "eda_ecdsa_p384",
]:
     print_key(name)

for name in [
    "rsa_encrypted",
    "ed25519_encrypted",
    "ecdsa_p256_encrypted",
    "ecdsa_p384_encrypted",
]:
     print_key(name, b"password")

print_key("rsa_encrypted_aes128_gcm", b"password", decode=False)

print_key("pkcs1", None, format="pem")
for name in ["pkcs1_aes_256_ctr", "pkcs1_aes_128_cbc"]:
    print_key(name, b"password", format="pem")

for name in ["pkcs8_rsa", "pkcs8_ecdsa_p256", "pkcs8_ecdsa_p384", "pkcs8_ed25519"]:
    print_key(name, None, format="pem")
print_key("pkcs8v2_ed25519", None, format="pem", decode=False)

for name in ["pkcs8_rsa_encrypted"]:
    print_key(name, b"password", format="pem")