#
#   Copyright © 2022 Evilham <contact@evilham.com>
#
#   This file is part of DD
#
#   DD 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.
#
#   DD 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 DD. If not, see <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: AGPL-3.0-or-later


"""
This file only depends on:
    - attrs
    - cryptography
    - requests

Integrations should be able to copy this file and use it verbatim on their
codebase without using anything else from DD.

To check for changes or to update this file, head to:
    https://gitlab.com/DD-workspace/DD/-/blob/main/dd-sso/admin/src/admin/lib/keys.py
"""

import json
import stat
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

import requests
from attr import field, frozen
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.serialization.base import (
    PRIVATE_KEY_TYPES,
    PUBLIC_KEY_TYPES,
)
from jose import jwe, jws, jwt
from jose.backends.rsa_backend import RSAKey
from jose.constants import ALGORITHMS

try:
    from functools import cache  # type: ignore  # Python 3.9+
except ImportError:
    try:
        from functools import cached_property as cache  # type: ignore  # Python 3.8
    except ImportError:
        from functools import lru_cache

        def cache(call: Callable) -> property:  # type: ignore  # Python 3.7
            return property(lru_cache()(call))


Data = Union[str, bytes]
Json = Union[Dict, str, List[Dict]]


@frozen(slots=False)
class ThirdPartyIntegrationKeys(object):
    """
    This represents the keys for a a third-party integration that interacts
    with us.


    These services must publish their public key on an agreed endpoint.
    See: {their_remote_pubkey_url}

    Their key is cached in the file system: '{key_store}/their.pubkey'.
    Key rotation can happen if requested by the third-party service, by
    manually removing the public key file from the data store.


    We also generate a private key for each third-party service.
    This key is read from or generated to if needed in the file system:
    '{key_store}/our.privkey'

    Rotating our private key can only happen by removing the private key file
    and restarting the service.

    In order to use the private key, you get full access to the cryptography
    primitive with {our_privkey}.

    In order to publish your public key, you get it serialised with
    {our_pubkey_pem}.
    """

    #
    # Standard attributes
    #
    our_name: str
    """
    The identifier of our service.
    """

    key_store: Path = field(converter=Path)
    """
    Local path to a directory containing service keys.
    """

    #
    # Attributes related to the third party
    #
    their_name: str
    """
    The identifier of the third-party.
    """

    their_service_domain: str
    """
    The domain on which the third party integration is running.
    """

    their_pubkey_endpoint: str = "/pubkey/"
    """
    The absolute path (starting with /) on {their_service_domain} that will
    serve the pubkey.
    """
    #
    # Cryptography attributes
    #

    key_size_bits: int = 4096
    """
    When generating private keys, how many bits will be used.
    """

    key_chmod: int = 0o600
    """
    Permissions to apply to private keys.
    """

    encrypt_algorithm: str = ALGORITHMS.RSA_OAEP_256
    """
    The encryption algorithm.
    """

    sign_algorithm: str = ALGORITHMS.RS512
    """
    The signature algorithm.
    """

    validity: timedelta = timedelta(minutes=1)
    """
    The validity of a signed token.
    """

    @property
    def validity_half(self) -> timedelta:
        return self.validity / 2

    #
    # Properties related to the third party keys
    #
    @property
    def their_pubkey_path(self) -> Path:
        """
        Helper returning the full path to the cached pubkey.
        """
        return self.key_store.joinpath("their.pubkey")

    @property
    def their_remote_pubkey_url(self) -> str:
        """
        Helper returning the full remote URL to the service's pubkey endpoint.
        """
        return f"https://{self.their_service_domain}{self.their_pubkey_endpoint}"

    @property
    def their_pubkey_bytes(self) -> bytes:
        """
        Return the service's pubkey by fetching it only if necessary.
        That means the key is fetched exactly once and saved to
        their_pubkey_path.
        In order to rotate keys, their_pubkey_path must be manually deleted.

        If the key has never been fetched and downloading it fails, an empty
        bytestring will be returned.

        Note we do not cache this property since there might be temporary
        failures.
        """
        if not self.their_pubkey_path.exists():
            # Download if not available
            try:
                res = requests.get(self.their_remote_pubkey_url)
            except:
                # The service is currently unavailable
                return b""
            self.their_pubkey_path.write_bytes(res.content)
        return self.their_pubkey_path.read_bytes()

    @property
    def their_pubkey_jwk(self) -> Dict[str, Any]:
        """
        Return the service's PublicKey in JWK form fetching if necessary or
        empty dict if it was not in the store and we were unable to fetch it.
        """
        pk_b = self.their_pubkey_bytes
        if not pk_b:
            return dict()
        rsak = RSAKey(json.loads(pk_b), self.sign_algorithm)
        return rsak.to_dict()

    #
    # Properties related to our own keys
    #
    @property
    def our_privkey_path(self) -> Path:
        """
        Helper returning the full path to our private key.
        """
        return self.key_store.joinpath("our.privkey")

    @cache
    def our_privkey_bytes(self) -> bytes:
        """
        Helper that returns our private key, generating it if necessary.

        This property is cached to avoid expensive operations.
        That does mean that on key rotation the service must be restarted.
        """
        return self._load_privkey(self.our_privkey_path)

    @cache
    def our_privkey(self) -> RSAKey:
        """
        Helper that returns our private key in cryptography form, generating
        it if necessary.
        """
        pk_b = self.our_privkey_bytes
        return RSAKey(json.loads(pk_b), self.sign_algorithm)

    @cache
    def our_privkey_jwk(self) -> Dict[str, Any]:
        """
        Return our PrivateKey in JWK for this third-party service, generating
        it if necessary.
        """
        return self.our_privkey.to_dict()

    @cache
    def our_pubkey(self) -> RSAKey:
        """
        Helper that returns our public key, generating the privkey if necessary.
        """
        return self.our_privkey.public_key()

    @cache
    def our_pubkey_jwk(self) -> Dict[str, Any]:
        """
        Helper that returns our public key in JWK form.
        """
        return self.our_pubkey.to_dict()

    #
    # Message-passing methods
    #
    def encrypt_outgoing_data(self, data: bytes) -> str:
        return jwe.encrypt(
            data,
            self.their_pubkey_jwk,
            algorithm=self.encrypt_algorithm,
            kid=self.our_name,
        ).decode("utf-8")

    def encrypt_outgoing_json(self, data: Json) -> str:
        return self.encrypt_outgoing_data(json.dumps(data).encode("utf-8"))

    def decrypt_incoming_data(self, enc_data: Data) -> bytes:
        return jwe.decrypt(enc_data, self.our_privkey_jwk)  # type: ignore  # Wrong hint

    def decrypt_incoming_json(self, enc_data: Data) -> Json:
        d: Json = json.loads(self.decrypt_incoming_data(enc_data))
        return d

    def sign_outgoing_json(self, data: Json) -> str:
        now = datetime.utcnow()
        claims = {
            "data": data,
            "aud": self.their_name,
            "iss": self.our_name,
            "iat": now,
            "nbf": now - self.validity_half,
            "exp": now + self.validity_half,
        }
        return jwt.encode(
            claims,
            self.our_privkey_jwk,
            algorithm=self.sign_algorithm,
            headers={
                "kid": self.our_name,
            },
        )

    def sign_outgoing_data(self, data: Json) -> str:
        return self.sign_outgoing_json(data)

    def verify_incoming_data(self, data: Data) -> Any:
        signed_data: Dict = jwt.decode(
            data,
            self.their_pubkey_jwk,
            algorithms=[self.sign_algorithm],
            audience=self.our_name,
            issuer=self.their_name,
            options={
                "require_aud": True,
                "require_iat": True,
                "require_iss": True,
                "require_nbf": True,
                "require_exp": True,
            },
        )
        return signed_data["data"]

    def verify_incoming_json(self, data: bytes) -> Json:
        d: Json = json.loads(self.verify_incoming_data(data))
        return d

    def sign_and_encrypt_outgoing_data(self, data: bytes) -> str:
        return self.sign_outgoing_data(self.encrypt_outgoing_data(data))

    def sign_and_encrypt_outgoing_json(self, data: Json) -> str:
        return self.sign_outgoing_data(self.encrypt_outgoing_json(data))

    def verify_and_decrypt_incoming_data(self, data: Data) -> bytes:
        enc_data: str = self.verify_incoming_data(data)
        return self.decrypt_incoming_data(enc_data.encode("utf-8"))

    def verify_and_decrypt_incoming_json(self, data: Data) -> Json:
        enc_data: str = self.verify_incoming_data(data)
        return self.decrypt_incoming_json(enc_data.encode("utf-8"))

    def get_outgoing_request_headers(self) -> Dict[str, str]:
        # Use current time as ever-changing payload
        now = datetime.utcnow().isoformat()
        return {"Authorization": self.sign_outgoing_data(now)}

    #
    # Helper methods
    #
    def _load_privkey(self, path: Path, force_generation: bool = False) -> bytes:
        # Check recommendations here
        # https://cryptography.io/en/latest/hazmat/primitives/asymmetric/rsa/#cryptography.hazmat.primitives.asymmetric.rsa.generate_private_key
        #
        needs_generation = force_generation or not path.exists()

        # Perform further sanity checks
        # by re-using needs_generation we save checking the file system
        if not needs_generation:
            path_st = path.stat()
            # Check and fix permissions if necessary
            if stat.S_IMODE(path_st.st_mode) != self.key_chmod:
                path.touch(mode=self.key_chmod, exist_ok=True)
            # Force generation if file is empty
            needs_generation = path_st == 0

        if needs_generation:
            # Generate the key as needed
            gpk = rsa.generate_private_key(
                public_exponent=65537, key_size=self.key_size_bits
            )
            enc_gpk = gpk.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption(),
            )
            enc_jwk = json.dumps(RSAKey(enc_gpk, self.sign_algorithm).to_dict())
            # Ensure permissions
            path.touch(mode=self.key_chmod, exist_ok=True)
            # Write private key
            path.write_text(enc_jwk)

        # Return key
        return path.read_bytes()


if __name__ == "__main__":
    # TODO: convert into real tests
    a = ThirdPartyIntegrationKeys(
        our_name="a", their_name="b", their_service_domain="b.com", key_store="tmpa"
    )
    b = ThirdPartyIntegrationKeys(
        our_name="b", their_name="a", their_service_domain="a.com", key_store="tmpb"
    )
    Path("tmpa/their.pubkey").write_text(json.dumps(b.our_pubkey_jwk))
    Path("tmpb/their.pubkey").write_text(json.dumps(a.our_pubkey_jwk))

    m1 = b"test message"
    print("out", m1)
    o1 = a.sign_and_encrypt_outgoing_data(m1)
    print("enc", o1)
    r1 = b.verify_and_decrypt_incoming_data(o1)
    print("got", r1)
    assert m1 == r1, f"Wrong result. Got: {r1!r}. Expected: {m1!r}"

    m2 = {"test": 1, "hello": "world"}
    m2 = {}
    print("out", m2)
    o2 = a.sign_and_encrypt_outgoing_json(m2)
    print(jwt.get_unverified_headers(o2))
    print("enc", o2)
    r2 = b.verify_and_decrypt_incoming_json(o2)
    print("got", r2)
    assert m2 == r2, f"Wrong result. Got: {r2!r}. Expected: {m2!r}"
