# SPDX-License-Identifier: MPL-2.0

"""
Test ISAKMP header parsing and composition.
"""

import collections
import unittest

from test.ike.classes import Ikev2PayloadBaseTest

from cryptodatahub.common.exception import InvalidValue
from cryptodatahub.ike.algorithm import (
    Ikev1ExchangeType,
    Ikev1PayloadType,
    Ikev2DiffieHellmanGroup,
    Ikev2ExchangeType,
    Ikev2PayloadType,
)
from cryptodatahub.ike.version import IkeVersion

from cryptoparser.common.exception import InvalidType, NotEnoughData
from cryptoparser.ike.common import IkePayloadTypeUnknown
from cryptoparser.ike.ikev1 import Ikev1PayloadKeyExchange, Ikev1PayloadNonce, Ikev1PayloadUnparsed
from cryptoparser.ike.ikev2 import Ikev2PayloadNonce, Ikev2PayloadKeyExchange, Ikev2PayloadUnparsed
from cryptoparser.ike.isakmp import IsakmpFlags, IsakmpMessage
from cryptoparser.ike.version import IsakmpProtocolVersion


class TestISAKMPHeader(unittest.TestCase):
    def setUp(self):
        self.header_dict = collections.OrderedDict([
            ('initiator_cookie', b'\x00' * 8),
            ('responder_cookie', b'\x00' * 8),
            ('next_payload', b'\x00'),  # NONE
            ('protocol_version', b'\x11'),  # ISAKMP v1.1
            ('exchange_type', b'\x01'),  # BASE
            ('flags', b'\x00'),
            ('message_id', b'\x00' * 4),
            ('length', b'\x00\x00\x00\x1c'),  # 28 bytes
        ])
        self.header_bytes = b''.join(self.header_dict.values())
        self.header = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V1, 1),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev1ExchangeType.BASE,
            flags=set(),
            message_id=0,
            payloads=[]
        )

    def test_error(self):
        with self.assertRaises(NotEnoughData) as context_manager:
            IsakmpMessage.parse_exact_size(b'')
        self.assertEqual(context_manager.exception.bytes_needed, IsakmpMessage.HEADER_SIZE)

        with self.assertRaises(InvalidType):
            IsakmpMessage.parse_exact_size(b'\x00' * (IsakmpMessage.HEADER_SIZE - 1) + b'\xff')

    def test_parse(self):
        header = IsakmpMessage.parse_exact_size(self.header_bytes)
        self.assertEqual(header.initiator_spi, 0)
        self.assertEqual(header.responder_spi, 0)
        self.assertEqual(header.version, IsakmpProtocolVersion(IkeVersion.V1, 1))
        self.assertEqual(header.exchange_type, Ikev1ExchangeType.BASE)
        self.assertEqual(header.flags, set())
        self.assertEqual(header.message_id, 0)

    def test_compose(self):
        self.assertEqual(self.header.compose(), self.header_bytes)

    def test_flags(self):
        header_with_flags = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V1, 1),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev1ExchangeType.BASE,
            flags={IsakmpFlags.ENCRYPTION, IsakmpFlags.COMMIT},
            message_id=0,
            payloads=[]
        )
        header_bytes_with_flags = bytearray(self.header_bytes)
        header_bytes_with_flags[19] = 0x03  # Set ENCRYPTION and COMMIT flags
        self.assertEqual(header_with_flags.compose(), bytes(header_bytes_with_flags))

    def test_ikev1_compose_payload_type(self):
        payload = Ikev1PayloadNonce(nonce_data=b'A' * 16)
        message = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V1, 0),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev1ExchangeType.INFORMATIONAL,
            flags=set(),
            message_id=0,
            payloads=[payload],
        )
        composed = message.compose()
        self.assertEqual(composed[16], Ikev1PayloadType.NONCE.value.code)

    def test_ikev2_parsing(self):
        header_dict_v2 = collections.OrderedDict([
            ('initiator_cookie', b'\x00' * 8),
            ('responder_cookie', b'\x00' * 8),
            ('next_payload', b'\x00'),  # NONE
            ('protocol_version', b'\x20'),  # ISAKMP v2.0
            ('exchange_type', b'\x22'),  # IKE_SA_INIT
            ('flags', b'\x00'),
            ('message_id', b'\x00' * 4),
            ('length', b'\x00\x00\x00\x1c'),  # 28 bytes
        ])
        header_bytes_v2 = b''.join(header_dict_v2.values())
        header = IsakmpMessage.parse_exact_size(header_bytes_v2)
        self.assertEqual(header.version, IsakmpProtocolVersion(IkeVersion.V2, 0))
        self.assertEqual(header.exchange_type, Ikev2ExchangeType.IKE_SA_INIT)

    def test_invalid_payload_type_ikev1(self):
        header_dict = collections.OrderedDict([
            ('initiator_cookie', b'\x00' * 8),
            ('responder_cookie', b'\x00' * 8),
            ('next_payload', b'\xff'),  # Invalid payload type
            ('protocol_version', b'\x11'),  # ISAKMP v1.1
            ('exchange_type', b'\x01'),  # BASE
            ('flags', b'\x00'),
            ('message_id', b'\x00' * 4),
            ('length', b'\x00\x00\x00\x1c'),  # 28 bytes
        ])
        header_bytes = b''.join(header_dict.values())
        with self.assertRaises(InvalidValue):
            IsakmpMessage.parse_exact_size(header_bytes)

    def test_ikev2_compose_with_empty_payloads(self):
        header_v2 = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V2, 0),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev2ExchangeType.IKE_SA_INIT,
            flags=set(),
            message_id=0,
            payloads=[]
        )

        # Empty payloads compose to a header-only 28-octet IKE message
        # (RFC 7296 §3.1) with ``next_payload = NONE``.
        composed = header_v2.compose()
        self.assertEqual(len(composed), IsakmpMessage.HEADER_SIZE)
        # Byte offset 16 (after 8-octet initiator + 8-octet responder SPI)
        # carries ``next_payload``; NONE is 0 for both IKEv1 / IKEv2.
        self.assertEqual(composed[16], 0)

    def test_payload_parsing_with_invalid_next_payload(self):
        header_dict = collections.OrderedDict([
            ('initiator_cookie', b'\x00' * 8),
            ('responder_cookie', b'\x00' * 8),
            ('next_payload', b'\x01'),  # SECURITY_ASSOCIATION
            ('protocol_version', b'\x11'),  # ISAKMP v1.1
            ('exchange_type', b'\x01'),  # BASE
            ('flags', b'\x00'),
            ('message_id', b'\x00' * 4),
            ('length', b'\x00\x00\x00\x20'),  # 32 bytes (header + 4 bytes payload)
        ])
        payload_data = b'\x00\x00\x00\x04'  # Invalid/minimal payload
        header_bytes = b''.join(header_dict.values()) + payload_data
        with self.assertRaises((NotEnoughData, TypeError, AttributeError)):
            IsakmpMessage.parse_exact_size(header_bytes)

    def test_payload_composition_loop(self):
        payload = Ikev2PayloadBaseTest(
            flags=set(),
            test_data=b'\x00\x01\x02\x03'
        )
        message = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V2, 0),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev2ExchangeType.IKE_SA_INIT,
            flags=set(),
            message_id=0,
            payloads=[payload]
        )
        message_bytes = message.compose()
        payload_bytes = payload.compose()
        self.assertEqual(message_bytes[-len(payload_bytes):], payload_bytes)

    def test_payload_parsing_loop(self):
        header_dict = collections.OrderedDict([
            ('initiator_cookie', b'\x00' * 8),
            ('responder_cookie', b'\x00' * 8),
            ('next_payload', b'\x28'),  # NONCE payload type (0x28 = 40 decimal)
            ('protocol_version', b'\x20'),  # ISAKMP v2.0
            ('exchange_type', b'\x22'),  # IKE_SA_INIT
            ('flags', b'\x00'),
            ('message_id', b'\x00' * 4),
            ('length', b'\x00\x00\x00\x30'),  # 48 bytes total (28 header + 20 payload)
        ])
        nonce_payload = (
            b'\x00' +        # Next payload = NONE (0x00)
            b'\x00' +        # Flags = 0 (no critical bit)
            b'\x00\x14' +    # Payload length = 20 bytes (header 4 + data 16)
            b'A' * 16       # 16 bytes of nonce data (minimum required for NONCE)
        )
        header_bytes = b''.join(header_dict.values()) + nonce_payload
        message = IsakmpMessage.parse_exact_size(header_bytes)
        self.assertEqual(len(message.payloads), 1)
        self.assertEqual(message.payloads[0].get_payload_type(), Ikev2PayloadType.NONCE)

    def test_invalid_payload_type_ikev2(self):
        header_dict = collections.OrderedDict([
            ('initiator_cookie', b'\x00' * 8),
            ('responder_cookie', b'\x00' * 8),
            ('next_payload', b'\xff'),  # Invalid payload type
            ('protocol_version', b'\x20'),  # ISAKMP v2.0
            ('exchange_type', b'\x22'),  # IKE_SA_INIT
            ('flags', b'\x00'),
            ('message_id', b'\x00' * 4),
            ('length', b'\x00\x00\x00\x1c'),  # 28 bytes
        ])
        header_bytes = b''.join(header_dict.values())
        with self.assertRaises(InvalidValue):
            IsakmpMessage.parse_exact_size(header_bytes)

    def test_unsupported_payload_type_ikev2(self):
        # ``next_payload`` is a valid enum (TSI) for which no parser
        # class is registered in :data:`IKEV2_PAYLOAD_CLASSES_BY_TYPE`.
        # Fall back to :class:`Ikev2PayloadUnparsed` and keep walking
        # the payload chain instead of raising ``InvalidType``.
        payload_body = b'\xab\xcd'
        payload_length = len(payload_body) + Ikev2PayloadUnparsed.HEADER_SIZE  # 6
        message_length = 28 + payload_length
        header_dict = collections.OrderedDict([
            ('initiator_cookie', b'\x00' * 8),
            ('responder_cookie', b'\x00' * 8),
            ('next_payload', bytes([Ikev2PayloadType.TSI.value.code])),
            ('protocol_version', b'\x20'),  # ISAKMP v2.0
            ('exchange_type', b'\x22'),  # IKE_SA_INIT
            ('flags', b'\x00'),
            ('message_id', b'\x00' * 4),
            ('length', message_length.to_bytes(4, 'big')),
        ])
        payload_bytes = b''.join([
            bytes([Ikev2PayloadType.NONE.value.code]),
            b'\x00',
            payload_length.to_bytes(2, 'big'),
            payload_body,
        ])
        message = IsakmpMessage.parse_exact_size(b''.join(header_dict.values()) + payload_bytes)
        self.assertEqual(len(message.payloads), 1)
        self.assertIsInstance(message.payloads[0], Ikev2PayloadUnparsed)
        self.assertEqual(bytes(message.payloads[0].payload_data), payload_body)

    def test_get_payload_by_type_ikev2(self):
        nonce_payload = Ikev2PayloadNonce(
            flags=set(),
            nonce_data=b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
        )
        ke_payload = Ikev2PayloadKeyExchange(
            flags=set(),
            dh_group=Ikev2DiffieHellmanGroup.MODP_GROUP_2048_BIT,
            key_exchange_data=b'\x00\x01\x02\x03\x04\x05\x06\x07'
        )
        message = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V2, 0),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev2ExchangeType.IKE_SA_INIT,
            flags=set(),
            message_id=0,
            payloads=[nonce_payload, ke_payload]
        )

        found_nonce = message.get_payload_by_type(Ikev2PayloadType.NONCE)
        self.assertEqual(found_nonce, nonce_payload)
        self.assertEqual(found_nonce.get_payload_type(), Ikev2PayloadType.NONCE)

        found_ke = message.get_payload_by_type(Ikev2PayloadType.KE)
        self.assertEqual(found_ke, ke_payload)
        self.assertEqual(found_ke.get_payload_type(), Ikev2PayloadType.KE)

    def test_get_payload_by_type_raises_index_error_on_multiple(self):
        first_nonce = Ikev2PayloadNonce(
            flags=set(),
            nonce_data=b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f',
        )
        second_nonce = Ikev2PayloadNonce(
            flags=set(),
            nonce_data=b'\xff\xfe\xfd\xfc\xfb\xfa\xf9\xf8\xf7\xf6\xf5\xf4\xf3\xf2\xf1\xf0',
        )
        message = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V2, 0),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev2ExchangeType.IKE_SA_INIT,
            flags=set(),
            message_id=0,
            payloads=[first_nonce, second_nonce],
        )

        with self.assertRaises(IndexError):
            message.get_payload_by_type(Ikev2PayloadType.NONCE)
        self.assertEqual(
            message.get_payloads_by_type(Ikev2PayloadType.NONCE),
            [first_nonce, second_nonce],
        )

    def test_get_payloads_by_type_returns_empty_list_when_missing(self):
        message = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V2, 0),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev2ExchangeType.IKE_SA_INIT,
            flags=set(),
            message_id=0,
            payloads=[],
        )

        self.assertEqual(
            message.get_payloads_by_type(Ikev2PayloadType.NONCE),
            [],
        )

    def test_get_payload_by_type_ikev1(self):
        ke_payload = Ikev1PayloadKeyExchange(key_exchange_data=b'\x00\x01\x02\x03')
        nonce_payload = Ikev1PayloadNonce(nonce_data=b'\x00\x01\x02\x03\x04\x05\x06\x07\x08')
        message = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V1, 1),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev1ExchangeType.BASE,
            flags=set(),
            message_id=0,
            payloads=[ke_payload, nonce_payload]
        )

        found_ke = message.get_payload_by_type(Ikev1PayloadType.KEY_EXCHANGE)
        self.assertEqual(found_ke, ke_payload)
        self.assertEqual(found_ke.get_payload_type(), Ikev1PayloadType.KEY_EXCHANGE)

        found_nonce = message.get_payload_by_type(Ikev1PayloadType.NONCE)
        self.assertEqual(found_nonce, nonce_payload)
        self.assertEqual(found_nonce.get_payload_type(), Ikev1PayloadType.NONCE)

    def test_get_payload_by_type_not_found(self):
        nonce_payload = Ikev2PayloadNonce(
            flags=set(),
            nonce_data=b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
        )
        message = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V2, 0),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev2ExchangeType.IKE_SA_INIT,
            flags=set(),
            message_id=0,
            payloads=[nonce_payload]
        )

        with self.assertRaises(KeyError) as context_manager:
            message.get_payload_by_type(Ikev2PayloadType.KE)
        self.assertEqual(context_manager.exception.args[0], Ikev2PayloadType.KE)

    def test_get_payload_by_type_empty_payloads(self):
        message = IsakmpMessage(
            version=IsakmpProtocolVersion(IkeVersion.V2, 0),
            initiator_spi=0,
            responder_spi=0,
            exchange_type=Ikev2ExchangeType.IKE_SA_INIT,
            flags=set(),
            message_id=0,
            payloads=[]
        )

        with self.assertRaises(KeyError) as context_manager:
            message.get_payload_by_type(Ikev2PayloadType.NONCE)
        self.assertEqual(context_manager.exception.args[0], Ikev2PayloadType.NONCE)


class TestIkePayloadTypeUnknown(unittest.TestCase):
    def test_get_byte_num(self):
        self.assertEqual(IkePayloadTypeUnknown.get_byte_num(), 1)

    def test_round_trip(self):
        wrapper = IkePayloadTypeUnknown(code=0xB8)
        composed = wrapper.compose()
        self.assertEqual(composed, b'\xb8')
        parsed = IkePayloadTypeUnknown.parse_exact_size(composed)
        self.assertEqual(parsed.code, 0xB8)


class TestIkev1PayloadUnparsed(unittest.TestCase):
    def test_round_trip(self):
        payload = Ikev1PayloadUnparsed(payload_data=b'\xaa\xbb\xcc\xdd')
        payload.payload_type = Ikev1PayloadType.HASH
        payload.next_payload = Ikev1PayloadType.NONE
        composed = payload.compose()
        parsed = Ikev1PayloadUnparsed.parse_exact_size(composed)
        self.assertEqual(parsed.payload_data, b'\xaa\xbb\xcc\xdd')
        self.assertEqual(parsed.next_payload, Ikev1PayloadType.NONE)

    def test_wire_payload_type_returned(self):
        payload = Ikev1PayloadUnparsed(payload_data=b'\x00')
        payload.payload_type = Ikev1PayloadType.HASH
        self.assertEqual(payload.get_payload_type(), Ikev1PayloadType.HASH)

    def test_unknown_next_payload_wrapped(self):
        # Header carries next_payload = 0xB8 (private-use in RFC 2408 §3.10).
        # Ikev1PayloadUnparsed._parse_header should route through
        # IkePayloadTypeUnknown instead of raising on the enum lookup.
        payload_body = b'\x11'
        header = b'\xb8\x00' + (4 + len(payload_body)).to_bytes(2, 'big') + payload_body
        parsed = Ikev1PayloadUnparsed.parse_exact_size(header)
        self.assertIsInstance(parsed.next_payload, IkePayloadTypeUnknown)
        self.assertEqual(parsed.next_payload.code, 0xB8)

    def test_compose_with_unknown_next_payload(self):
        payload = Ikev1PayloadUnparsed(payload_data=b'\x99')
        payload.payload_type = Ikev1PayloadType.HASH
        payload.next_payload = IkePayloadTypeUnknown(code=0xB8)
        composed = payload.compose()
        # First octet is the raw private-use code.
        self.assertEqual(composed[0], 0xB8)


class TestIkev2PayloadUnparsedRoundTrip(unittest.TestCase):
    def test_round_trip(self):
        payload = Ikev2PayloadUnparsed(flags=set(), payload_data=b'\xaa\xbb\xcc\xdd')
        payload.payload_type = Ikev2PayloadType.SK
        payload.next_payload = Ikev2PayloadType.NONE
        composed = payload.compose()
        parsed = Ikev2PayloadUnparsed.parse_exact_size(composed)
        self.assertEqual(parsed.payload_data, b'\xaa\xbb\xcc\xdd')
        self.assertEqual(parsed.next_payload, Ikev2PayloadType.NONE)

    def test_wire_payload_type_returned(self):
        payload = Ikev2PayloadUnparsed(flags=set(), payload_data=b'\x00')
        payload.payload_type = Ikev2PayloadType.SK
        self.assertEqual(payload.get_payload_type(), Ikev2PayloadType.SK)

    def test_unknown_next_payload_wrapped(self):
        payload_body = b'\x11'
        header = b'\xb8\x00' + (4 + len(payload_body)).to_bytes(2, 'big') + payload_body
        parsed = Ikev2PayloadUnparsed.parse_exact_size(header)
        self.assertIsInstance(parsed.next_payload, IkePayloadTypeUnknown)
        self.assertEqual(parsed.next_payload.code, 0xB8)

    def test_compose_with_unknown_next_payload(self):
        payload = Ikev2PayloadUnparsed(flags=set(), payload_data=b'\x99')
        payload.payload_type = Ikev2PayloadType.SK
        payload.next_payload = IkePayloadTypeUnknown(code=0xB8)
        self.assertEqual(payload.compose()[0], 0xB8)
