# SPDX-License-Identifier: MPL-2.0
"""ISAKMP header parser."""

import enum
import typing

import attr

from cryptodatahub.ike.algorithm import Ikev1PayloadType, Ikev2PayloadType, Ikev1ExchangeType, Ikev2ExchangeType
from cryptodatahub.ike.version import IkeVersion

from cryptoparser.common.parse import ParsableBase, ParserBinary, ComposerBinary
from cryptoparser.common.exception import NotEnoughData

from cryptoparser.ike.version import IsakmpProtocolVersion
from cryptoparser.ike.ikev1 import (
    Ikev1PayloadBase,
    Ikev1PayloadUnparsed,
    IKEV1_PAYLOAD_CLASSES_BY_TYPE,
)
from cryptoparser.ike.ikev2 import (
    Ikev2PayloadBase,
    Ikev2PayloadUnparsed,
    IKEV2_PAYLOAD_CLASSES_BY_TYPE,
)


class IsakmpFlags(enum.IntFlag):
    """ISAKMP flags."""
    ENCRYPTION = 1 << 0
    COMMIT = 1 << 1
    AUTHENTICATION_ONLY = 1 << 2
    INITIATOR = 1 << 3
    VERSION = 1 << 4
    RESPONSE = 1 << 5


@attr.s
class IsakmpMessage(ParsableBase):
    """ISAKMP message parser.

    .. code-block:: text

                        1                   2                   3
         0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                       IKE SA Initiator's SPI                  |
        |                                                               |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                       IKE SA Responder's SPI                  |
        |                                                               |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |  Next Payload | MjVer | MnVer | Exchange Type |     Flags     |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                          Message ID                           |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                            Length                             |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                       Payload Data (variable)                 |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

    :ivar initiator_spi: Initiator SPI
    :ivar responder_spi: Responder SPI
    :ivar exchange_type: Exchange Type
    :ivar flags: Flags
    :ivar message_id: Message ID
    :ivar payloads: Payloads
    :ivar version: Version
    """
    HEADER_SIZE = 28

    version: IsakmpProtocolVersion = attr.ib(validator=attr.validators.instance_of(IsakmpProtocolVersion))
    initiator_spi: int = attr.ib(validator=attr.validators.and_(
        attr.validators.instance_of(int),
        attr.validators.ge(0),
        attr.validators.lt(2**64)
    ))
    responder_spi: int = attr.ib(validator=attr.validators.and_(
        attr.validators.instance_of(int),
        attr.validators.ge(0),
        attr.validators.lt(2**64)))
    exchange_type: typing.Union[Ikev1ExchangeType, Ikev2ExchangeType] = attr.ib(
        validator=attr.validators.instance_of((Ikev1ExchangeType, Ikev2ExchangeType))
    )
    flags: list[IsakmpFlags] = attr.ib(validator=attr.validators.deep_iterable(
        member_validator=attr.validators.instance_of(IsakmpFlags)
    ))
    message_id: int = attr.ib(validator=attr.validators.and_(
        attr.validators.instance_of(int),
        attr.validators.ge(0),
        attr.validators.lt(2**32)
    ))
    payloads: list[typing.Union[Ikev1PayloadBase, Ikev2PayloadBase]] = attr.ib(
        validator=attr.validators.deep_iterable(
            member_validator=attr.validators.instance_of((Ikev1PayloadBase, Ikev2PayloadBase))
        )
    )

    def get_payloads_by_type(
        self, payload_type: typing.Union[Ikev1PayloadType, Ikev2PayloadType]
    ) -> list[typing.Union[Ikev1PayloadBase, Ikev2PayloadBase]]:
        return [
            payload for payload in self.payloads
            if payload.get_payload_type() == payload_type
        ]

    def get_payload_by_type(
        self, payload_type: typing.Union[Ikev1PayloadType, Ikev2PayloadType]
    ) -> typing.Union[Ikev1PayloadBase, Ikev2PayloadBase]:
        payloads = self.get_payloads_by_type(payload_type)
        if not payloads:
            raise KeyError(payload_type)
        if len(payloads) > 1:
            raise IndexError(payload_type)
        return payloads[0]

    @classmethod
    def parse_header_and_body(cls, parsable):
        """Parse the 28-byte ISAKMP header and slice out the payload body.

        Returns ``(header, body_bytes, next_payload)``. ``header`` is a dict
        of the concrete header fields; ``body_bytes`` is
        ``parsable[HEADER_SIZE:header.length]``; ``next_payload`` is the
        first payload-type code (enum member or
        :class:`IkePayloadTypeUnknown`).

        Callers walk the plaintext payload chain with
        :meth:`parse_payload_chain`; encryption-aware callers (IKEv1
        ``ENCRYPTION`` flag per RFC 2408 §3.1, wire-level state that
        cryptoparser does not track) wrap ``body_bytes`` opaquely instead.
        """
        if len(parsable) < cls.HEADER_SIZE:
            raise NotEnoughData(cls.HEADER_SIZE - len(parsable))

        parser = ParserBinary(parsable)
        parser.parse_numeric('initiator_spi', 8)
        parser.parse_numeric('responder_spi', 8)
        parser.parse_raw('next_payload', 1)
        parser.parse_parsable('version', IsakmpProtocolVersion)
        if parser['version'].major == IkeVersion.V1:
            next_payload_type = Ikev1PayloadType
            parser.parse_numeric_enum_coded('exchange_type', Ikev1ExchangeType)
        elif parser['version'].major == IkeVersion.V2:
            next_payload_type = Ikev2PayloadType
            parser.parse_numeric_enum_coded('exchange_type', Ikev2ExchangeType)
        else:
            raise NotImplementedError(parser['version'])

        next_payload_parser = ParserBinary(parser['next_payload'])
        next_payload_parser.parse_numeric_enum_coded('value', next_payload_type)
        next_payload = next_payload_parser['value']

        parser.parse_numeric_flags('flags', 1, IsakmpFlags)
        parser.parse_numeric('message_id', 4)
        parser.parse_numeric('length', 4)

        body_bytes = bytes(parsable[parser.parsed_length:parser['length']])
        header = {
            'initiator_spi': parser['initiator_spi'],
            'responder_spi': parser['responder_spi'],
            'version': parser['version'],
            'exchange_type': parser['exchange_type'],
            'flags': parser['flags'],
            'message_id': parser['message_id'],
        }
        return header, body_bytes, next_payload

    @classmethod
    def parse_payload_chain(cls, version, body_bytes, next_payload):
        """Walk the payload chain from ``body_bytes`` starting at
        ``next_payload``. Returns ``(payloads, consumed_length)``.

        Do NOT call this for IKEv1 messages with the ``ENCRYPTION`` flag
        set: ``body_bytes`` is ciphertext under SKEYID_e (RFC 2408 §3.1)
        and walking it as plaintext misreads the first block's
        payload-length field. That case is a caller-side concern.
        """
        if version.major == IkeVersion.V1:
            payload_classes = IKEV1_PAYLOAD_CLASSES_BY_TYPE
            payload_none = Ikev1PayloadType.NONE
            unparsed_class = Ikev1PayloadUnparsed
        elif version.major == IkeVersion.V2:
            payload_classes = IKEV2_PAYLOAD_CLASSES_BY_TYPE
            payload_none = Ikev2PayloadType.NONE
            unparsed_class = Ikev2PayloadUnparsed
        else:
            raise NotImplementedError(version)

        payloads = []
        parser_payload = ParserBinary(body_bytes)
        while parser_payload.unparsed_length > 0 and next_payload != payload_none:
            # ``next_payload`` is either a known enum member or — for
            # private-use / unimplemented codes (IANA "IKE Payload
            # Types" 128-255 per RFC 2408 §3.10 / RFC 7296 §3.2) — an
            # :class:`IkePayloadTypeUnknown` wrapper produced by
            # ``Ikev*PayloadTypeFactory``. Both cases fall through the
            # payload-class map's ``.get`` fallback to
            # ``Ikev*PayloadUnparsed`` so the payload walk continues.
            payload_class = payload_classes.get(next_payload, unparsed_class)
            parser_payload.parse_parsable('payload', payload_class)

            payload = parser_payload['payload']
            if isinstance(payload, (Ikev1PayloadUnparsed, Ikev2PayloadUnparsed)):
                # Preserve the on-wire payload-type code so
                # ``compose()`` can reproduce it in the previous
                # payload's ``next_payload`` field.
                payload.payload_type = next_payload
            payloads.append(payload)

            next_payload = payload.next_payload

        return payloads, parser_payload.parsed_length

    @classmethod
    def _parse(cls, parsable):
        header, body_bytes, next_payload = cls.parse_header_and_body(parsable)
        payloads, consumed_length = cls.parse_payload_chain(header['version'], body_bytes, next_payload)
        return cls(payloads=payloads, **header), cls.HEADER_SIZE + consumed_length

    @classmethod
    def compose_payload_chain(cls, version, payloads):
        """Compose the payload chain: link each payload's ``next_payload``
        to its successor's type code and concatenate the encoded bytes.
        """
        if version.major == IkeVersion.V1:
            payload_none = Ikev1PayloadType.NONE
        elif version.major == IkeVersion.V2:
            payload_none = Ikev2PayloadType.NONE
        else:
            raise NotImplementedError(version)

        payload_composer = ComposerBinary()
        for i, payload in enumerate(payloads):
            payload.next_payload = (
                payload_none
                if i == len(payloads) - 1
                else payloads[i + 1].get_payload_type()
            )
            payload_composer.compose_parsable(payload)
        return payload_composer.composed_bytes

    @classmethod
    def compose_header(  # pylint: disable=too-many-arguments,too-many-positional-arguments
            cls,
            initiator_spi,
            responder_spi,
            version,
            first_payload_type,
            exchange_type,
            flags,
            message_id,
            ike_total_length,
    ):
        """Compose the 28-octet IKE message header (RFC 7296 §3.1).

        Standalone so callers can build an outer header without touching
        the payload chain — needed by IKEv2 SK-envelope compose paths that
        supply the encrypted body directly and must not have
        ``payload.next_payload`` overwritten (see
        :meth:`compose_payload_chain`).
        """
        header_composer = ComposerBinary()
        header_composer.compose_numeric(initiator_spi, 8)
        header_composer.compose_numeric(responder_spi, 8)
        header_composer.compose_numeric_enum_coded(first_payload_type)
        header_composer.compose_parsable(version)
        header_composer.compose_numeric_enum_coded(exchange_type)
        header_composer.compose_numeric_flags(flags, 1)
        header_composer.compose_numeric(message_id, 4)
        header_composer.compose_numeric(ike_total_length, 4)
        return header_composer.composed_bytes

    def compose(self):
        if self.version.major == IkeVersion.V1:
            payload_none = Ikev1PayloadType.NONE
        elif self.version.major == IkeVersion.V2:
            payload_none = Ikev2PayloadType.NONE
        else:
            raise NotImplementedError(self.version)

        payload_bytes = self.compose_payload_chain(self.version, self.payloads)
        first_payload_type = (
            self.payloads[0].get_payload_type() if self.payloads else payload_none
        )
        header_bytes = self.compose_header(
            self.initiator_spi, self.responder_spi, self.version, first_payload_type,
            self.exchange_type, self.flags, self.message_id,
            self.HEADER_SIZE + len(payload_bytes),
        )
        return header_bytes + payload_bytes
