#!/usr/bin/env python3

# Libervia: an XMPP client
# Copyright (C) 2025-2026 Jérôme Poisosn (goffi@goffi.org)
# Copyright (C) 2022-2026 Tim Henkes (me@syndace.dev)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from typing import Any

import gpg

from .exceptions import (
    DecryptionFailed,
    GPGProviderError,
    InvalidPacket,
    UnknownKey,
    VerificationFailed,
)
from .models import GPGProviderBase, GPGPublicKeyBase, GPGSecretKeyBase


class GPGPublicKey(GPGPublicKeyBase):
    """
    GPG public key implementation based on GnuPG Made Easy (GPGME).
    """

    def __init__(self, key_obj: Any) -> None:
        """
        @param key_obj: The GPGME key object.
        """

        self.__key_obj = key_obj

    @property
    def fingerprint(self) -> str:
        return self.__key_obj.fpr

    @property
    def key_obj(self) -> Any:
        return self.__key_obj


class GPGSecretKey(GPGSecretKeyBase):
    """
    GPG secret key implementation based on GnuPG Made Easy (GPGME).
    """

    def __init__(self, public_key: GPGPublicKey) -> None:
        """
        @param public_key: The public key corresponding to this secret key.
        """

        self.__public_key = public_key

    @property
    def public_key(self) -> GPGPublicKey:
        return self.__public_key


class GPGProvider(GPGProviderBase):
    """
    GPG provider implementation based on GnuPG Made Easy (GPGME).
    """

    ALGORITHM = "ed25519/cert,sign+cv25519/encr"

    def __init__(self, home_dir: str | None = None) -> None:
        """
        @param home_dir: Optional GPG home directory path to use for all operations.
        """

        self.__home_dir = home_dir

    def export_public_key(self, public_key: GPGPublicKeyBase) -> bytes:
        assert isinstance(public_key, GPGPublicKey)

        pattern = public_key.fingerprint

        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                result = c.key_export_minimal(pattern)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e

            if result is None:
                raise UnknownKey(f"Public key {pattern} not found.")

            return result

    def import_public_key(self, packet: bytes) -> GPGPublicKeyBase:
        # TODO
        # - Reject packets older than version 4
        # - Check whether it's actually a public key (through packet inspection?)

        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                result = c.key_import(packet)
            except gpg.errors.GPGMEError as e:
                # From looking at the code, `key_import` never raises. The documentation
                # says it does though, so this is included for future-proofness.
                raise GPGProviderError("Internal GPGME error") from e

            if not hasattr(result, "considered"):
                raise InvalidPacket(
                    f"Data not considered for public key import: {result}"
                )

            if len(result.imports) != 1:
                raise InvalidPacket(
                    "Public key packet does not contain exactly one public key (not"
                    " counting subkeys)."
                )

            try:
                key_obj = c.get_key(result.imports[0].fpr, secret=False)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.errors.KeyError as e:
                raise GPGProviderError("Newly imported public key not found") from e

            return GPGPublicKey(key_obj)

    def backup_secret_key(self, secret_key: GPGSecretKeyBase) -> bytes:
        assert isinstance(secret_key, GPGSecretKey)
        # TODO
        # - Handle password protection/pinentry
        # - Make sure the key is exported unencrypted

        pattern = secret_key.public_key.fingerprint

        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                result = c.key_export_secret(pattern)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e

            if result is None:
                raise UnknownKey(f"Secret key {pattern} not found.")

            return result

    def restore_secret_keys(self, data: bytes) -> set[GPGSecretKeyBase]:
        # TODO
        # - Reject packets older than version 4
        # - Check whether it's actually secret keys (through packet inspection?)

        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                result = c.key_import(data)
            except gpg.errors.GPGMEError as e:
                # From looking at the code, `key_import` never raises. The documentation
                # says it does though, so this is included for future-proofness.
                raise GPGProviderError("Internal GPGME error") from e

            if not hasattr(result, "considered"):
                raise InvalidPacket(
                    f"Data not considered for secret key import: {result}"
                )

            if len(result.imports) == 0:
                raise InvalidPacket("Secret key packet does not contain a secret key.")

            secret_keys = set()
            for import_status in result.imports:
                try:
                    key_obj = c.get_key(import_status.fpr, secret=True)
                except gpg.errors.GPGMEError as e:
                    raise GPGProviderError("Internal GPGME error") from e
                except gpg.errors.KeyError as e:
                    raise GPGProviderError("Newly imported secret key not found") from e

                secret_keys.add(GPGSecretKey(GPGPublicKey(key_obj)))

            return secret_keys

    def encrypt_symmetrically(self, plaintext: bytes, password: str) -> bytes:
        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                ciphertext, __, __ = c.encrypt(plaintext, passphrase=password, sign=False)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e

            return ciphertext

    def decrypt_symmetrically(self, ciphertext: bytes, password: str) -> bytes:
        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                plaintext, __, __ = c.decrypt(
                    ciphertext, passphrase=password, verify=False
                )
            except gpg.errors.GPGMEError as e:
                # TODO: Find out what kind of error is raised if the password is wrong and
                # re-raise it as DecryptionFailed instead.
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.UnsupportedAlgorithm as e:
                raise DecryptionFailed("Unsupported algorithm") from e

            return plaintext

    def sign(self, data: bytes, secret_keys: set[GPGSecretKeyBase]) -> bytes:
        signers = []
        for secret_key in secret_keys:
            assert isinstance(secret_key, GPGSecretKey)

            signers.append(secret_key.public_key.key_obj)

        with gpg.Context(home_dir=self.__home_dir, signers=signers) as c:
            try:
                signed_data, __ = c.sign(data)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.errors.InvalidSigners as e:
                raise GPGProviderError(
                    "At least one of the secret keys is invalid for signing"
                ) from e

            return signed_data

    def sign_detached(self, data: bytes, secret_keys: set[GPGSecretKeyBase]) -> bytes:
        signers = []
        for secret_key in secret_keys:
            assert isinstance(secret_key, GPGSecretKey)

            signers.append(secret_key.public_key.key_obj)

        with gpg.Context(home_dir=self.__home_dir, signers=signers) as c:
            try:
                signature, __ = c.sign(data, mode=gpg.constants.sig.mode.DETACH)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.errors.InvalidSigners as e:
                raise GPGProviderError(
                    "At least one of the secret keys is invalid for signing"
                ) from e

            return signature

    def verify(self, signed_data: bytes, public_keys: set[GPGPublicKeyBase]) -> bytes:
        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                data, result = c.verify(signed_data)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.errors.BadSignatures as e:
                raise VerificationFailed("Bad signatures on signed data") from e

            valid_signature_found = False
            for public_key in public_keys:
                assert isinstance(public_key, GPGPublicKey)

                for subkey in public_key.key_obj.subkeys:
                    for sig in result.signatures:
                        if subkey.can_sign and subkey.fpr == sig.fpr:
                            valid_signature_found = True

            if not valid_signature_found:
                raise VerificationFailed(
                    "Data not signed by one of the expected public keys"
                )

            return data

    def verify_detached(
        self, data: bytes, signature: bytes, public_keys: set[GPGPublicKeyBase]
    ) -> None:
        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                __, result = c.verify(data, signature=signature)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.errors.BadSignatures as e:
                raise VerificationFailed("Bad signatures on signed data") from e

            valid_signature_found = False
            for public_key in public_keys:
                assert isinstance(public_key, GPGPublicKey)

                for subkey in public_key.key_obj.subkeys:
                    for sig in result.signatures:
                        if subkey.can_sign and subkey.fpr == sig.fpr:
                            valid_signature_found = True

            if not valid_signature_found:
                raise VerificationFailed(
                    "Data not signed by one of the expected public keys"
                )

    def encrypt(
        self,
        plaintext: bytes,
        public_keys: set[GPGPublicKeyBase],
        signing_keys: set[GPGSecretKeyBase] | None = None,
    ) -> bytes:
        recipients = []
        for public_key in public_keys:
            assert isinstance(public_key, GPGPublicKey)

            recipients.append(public_key.key_obj)

        signers = []
        if signing_keys is not None:
            for secret_key in signing_keys:
                assert isinstance(secret_key, GPGSecretKey)

                signers.append(secret_key.public_key.key_obj)

        sign = signing_keys is not None

        kwargs = {"home_dir": self.__home_dir, "signers": signers}

        with gpg.Context(**kwargs) as c:
            try:
                ciphertext, __, __ = c.encrypt(
                    plaintext,
                    recipients=recipients,
                    sign=sign,
                    always_trust=True,
                    add_encrypt_to=True,
                )
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.errors.InvalidRecipients as e:
                raise GPGProviderError(
                    "At least one of the public keys is invalid for encryption"
                ) from e
            except gpg.errors.InvalidSigners as e:
                raise GPGProviderError(
                    "At least one of the signing keys is invalid for signing"
                ) from e

            return ciphertext

    def decrypt(
        self,
        ciphertext: bytes,
        secret_keys: set[GPGSecretKeyBase],
        public_keys: set[GPGPublicKeyBase] | None = None,
    ) -> bytes:
        verify = public_keys is not None

        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                plaintext, result, verify_result = c.decrypt(ciphertext, verify=verify)
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.UnsupportedAlgorithm as e:
                raise DecryptionFailed("Unsupported algorithm") from e

            # TODO: Check whether the data was decrypted using one of the expected secret
            # keys

            if public_keys is not None:
                valid_signature_found = False
                for public_key in public_keys:
                    assert isinstance(public_key, GPGPublicKey)

                    for subkey in public_key.key_obj.subkeys:
                        for sig in verify_result.signatures:
                            if subkey.can_sign and subkey.fpr == sig.fpr:
                                valid_signature_found = True

                if not valid_signature_found:
                    raise VerificationFailed(
                        "Data not signed by one of the expected public keys"
                    )

            return plaintext

    def list_public_keys(self, user_id: str) -> set[GPGPublicKeyBase]:
        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                return {
                    GPGPublicKey(key) for key in c.keylist(pattern=user_id, secret=False)
                }
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e

    def list_secret_keys(self, user_id: str) -> set[GPGSecretKeyBase]:
        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                return {
                    GPGSecretKey(GPGPublicKey(key))
                    for key in c.keylist(pattern=user_id, secret=True)
                }
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e

    def can_sign(self, public_key: GPGPublicKeyBase) -> bool:
        assert isinstance(public_key, GPGPublicKey)

        return any(subkey.can_sign for subkey in public_key.key_obj.subkeys)

    def can_encrypt(self, public_key: GPGPublicKeyBase) -> bool:
        assert isinstance(public_key, GPGPublicKey)

        return any(subkey.can_encrypt for subkey in public_key.key_obj.subkeys)

    def create_key(self, user_id: str) -> GPGSecretKeyBase:
        with gpg.Context(home_dir=self.__home_dir) as c:
            try:
                # A single ECC (primary) key with signing and encryption capabilities is
                # not supported, a subkey is required for encryption. Create a primary key
                # only capable of certifying subkeys
                primary_create_result = c.create_key(
                    user_id,
                    algorithm=self.ALGORITHM,
                    expires=False,
                    certify=True,
                    force=True,
                )

                primary_key_obj = c.get_key(primary_create_result.fpr, secret=True)

                c.create_subkey(
                    primary_key_obj, algorithm=self.ALGORITHM, expires=False, sign=True
                )

                c.create_subkey(
                    primary_key_obj,
                    algorithm=self.ALGORITHM,
                    expires=False,
                    encrypt=True,
                )
            except gpg.errors.GPGMEError as e:
                raise GPGProviderError("Internal GPGME error") from e
            except gpg.errors.KeyNotFound as e:
                raise GPGProviderError("Newly created key not found") from e

            return GPGSecretKey(GPGPublicKey(primary_key_obj))