from cryptography.exceptions import UnsupportedAlgorithm
from cryptography.hazmat.primitives.asymmetric import ed25519, rsa
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.serialization import (
    SSHCertificateBuilder,
    SSHCertificateType,
    SSHCertificate,
    Encoding,
    PrivateFormat,
    PublicFormat,
    NoEncryption,
    load_ssh_private_key,
    load_ssh_public_key,
)

import logging
import os
from datetime import datetime, timedelta, timezone
from collections import namedtuple
from dataclasses import dataclass

logger = logging.getLogger("crypto")

PubKey = Ed25519PublicKey
PrivKey = Ed25519PrivateKey


@dataclass
class KeyPair:
    private: PrivKey
    public: PubKey


def generate_key() -> KeyPair:
    private_key = Ed25519PrivateKey.generate()
    public_key = private_key.public_key()

    return KeyPair(private_key, public_key)


def load_priv_key(private_key_path: str) -> PrivKey:
    try:
        with open(private_key_path, "rb") as f:
            private_key = load_ssh_private_key(f.read(), None)

        assert isinstance(private_key, Ed25519PrivateKey)
        return private_key
    except ValueError as e:
        logger.error("Unable to load key: %s", e)
        raise
    except (UnsupportedAlgorithm, AssertionError) as e:
        logger.error("Key type unsupported: %s", e)
        raise


def load_pub_key(public_key_path: str) -> PubKey:
    try:
        with open(public_key_path, "rb") as f:
            public_key = load_ssh_public_key(f.read())

        assert isinstance(public_key, Ed25519PublicKey)
        return public_key
    except ValueError as e:
        logger.error("Unable to load key: %s", e)
        raise
    except (UnsupportedAlgorithm, AssertionError) as e:
        logger.error("Key type unsupported: %s", e)
        raise


def load_key(private_key_path: str, public_key_path: str) -> KeyPair:
    return KeyPair(load_priv_key(private_key_path), load_pub_key(public_key_path))


def loads_pub_key(public_key_bytes: bytes):
    """
    Try to load a public key from raw bytes.
    Raises ValueError if an invalid key is passed.
    """
    try:
        public_key = load_ssh_public_key(public_key_bytes)
        assert isinstance(public_key, Ed25519PublicKey)
        return public_key
    except ValueError as e:
        logger.error("Unable to load key: %s", e)
        raise
    except (UnsupportedAlgorithm, AssertionError) as e:
        logger.error("Key type unsupported: %s", e)
        raise ValueError(f"Unable to load key: {e}") from e


def secure_open(path, flags):
    """Open a private key file securely, ensuring the file mode is 0600."""
    return os.open(path, flags, 0o600)


def save_key(key: KeyPair, private_key_path: str, public_key_path: str) -> None:
    """Save a key pair in the appropriate formats for OpenSSH."""
    with open(private_key_path, "wb", opener=secure_open) as f:
        f.write(
            key.private.private_bytes(
                Encoding.PEM, PrivateFormat.OpenSSH, NoEncryption()
            )
        )

    with open(public_key_path, "wb") as f:
        f.write(key.public.public_bytes(Encoding.OpenSSH, PublicFormat.OpenSSH))


def sign_key(
    ca_key: KeyPair,
    public_key: PubKey,
    certificate_type: SSHCertificateType,
    *,
    key_id: bytes,
    principals: list[bytes],
    valid_time: timedelta,
) -> SSHCertificate:
    now = datetime.now(tz=timezone.utc)

    builder = (
        SSHCertificateBuilder()
        .public_key(public_key)
        .type(certificate_type)
        .valid_after(now.timestamp())
        .valid_before((now + valid_time).timestamp())
        .key_id(key_id)
        .valid_principals(principals)
        .add_extension(b"permit-agent-forwarding", b"")
        .add_extension(b"permit-pty", b"")
        .add_extension(b"permit-X11-forwarding", b"")
        .add_extension(b"permit-port-forwarding", b"")
        .add_extension(b"permit-user-rc", b"")
    )

    return builder.sign(ca_key.private)
