# SPDX-License-Identifier: MPL-2.0
# pylint: disable=too-many-lines

import collections
import ipaddress
import unittest

from cryptodatahub.common.algorithm import Signature
from cryptodatahub.common.exception import InvalidValue
from cryptodatahub.ike.algorithm import (
    Ikev2AuthenticationMethod,
    Ikev2CertificateType,
    Ikev2DiffieHellmanGroup,
    Ikev2IdType,
    Ikev2NotifyType,
    Ikev2ProtocolId,
    Ikev2TransformAttributeType,
)

from cryptoparser.common.exception import InvalidType, NotEnoughData, TooMuchData
from cryptoparser.ike.ikev2 import (
    Ikev2AuthDigitalSignatureEnvelope,
    Ikev2NotifyPayloadInvalidKe,
    Ikev2PayloadAuthentication,
    Ikev2PayloadCertificate,
    Ikev2PayloadCertificateRequest,
    Ikev2PayloadDelete,
    Ikev2PayloadEap,
    Ikev2PayloadEncryptedAndAuthenticated,
    Ikev2PayloadFlags,
    Ikev2PayloadIdentificationInitiator,
    Ikev2PayloadIdentificationInitiatorDerAsn1Gn,
    Ikev2PayloadIdentificationInitiatorFcName,
    Ikev2PayloadIdentificationInitiatorFqdn,
    Ikev2PayloadIdentificationInitiatorIpv6Addr,
    Ikev2PayloadIdentificationInitiatorKeyId,
    Ikev2PayloadIdentificationInitiatorNull,
    Ikev2PayloadIdentificationInitiatorRfc822Addr,
    Ikev2PayloadIdentificationResponder,
    Ikev2PayloadIdentificationResponderDerAsn1Dn,
    Ikev2PayloadIdentificationResponderFqdn,
    Ikev2PayloadIdentificationResponderIpv4Addr,
    Ikev2PayloadIdentificationResponderIpv6Addr,
    Ikev2PayloadKeyExchange,
    Ikev2PayloadNonce,
    Ikev2PayloadNotifyAuthenticationFailed,
    Ikev2PayloadNotifyUnparsed,
    Ikev2PayloadType,
    Ikev2PayloadVendorId,
    TransformAttributeKeyLength,

)

from .classes import Ikev2PayloadBaseTest, Ikev2PayloadNotifyNoDataTest


class TestIkev2PayloadBase(unittest.TestCase):
    def setUp(self):
        self.test_data = b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
        self.test_payload_minimal = Ikev2PayloadBaseTest(
            flags=set(),
            test_data=b'',
        )
        self.test_payload_minimal.next_payload = Ikev2PayloadType.NONE

        self.test_dict_minimal = collections.OrderedDict([
            ('next_payload', b'\x00'),
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x04'),
            ('test_data', b''),
        ])
        self.test_bytes_minimal = b''.join(self.test_dict_minimal.values())

        self.test_payload_with_data = Ikev2PayloadBaseTest(
            flags={Ikev2PayloadFlags.CRITICAL},
            test_data=self.test_data
        )
        self.test_payload_with_data.next_payload = Ikev2PayloadType.NONE

        self.test_dict_with_data = collections.OrderedDict([
            ('next_payload', b'\x00'),
            ('flags', b'\x80'),
            ('payload_length', b'\x00\x14'),
            ('test_data', self.test_data),
        ])
        self.test_bytes_with_data = b''.join(self.test_dict_with_data.values())

    def test_parse(self):
        parsed_payload = Ikev2PayloadBaseTest.parse_exact_size(self.test_bytes_minimal)
        self.assertEqual(parsed_payload.flags, self.test_payload_minimal.flags)
        self.assertEqual(parsed_payload.test_data, self.test_payload_minimal.test_data)
        self.assertEqual(parsed_payload.next_payload, self.test_payload_minimal.next_payload)

        parsed_payload = Ikev2PayloadBaseTest.parse_exact_size(self.test_bytes_with_data)
        self.assertEqual(parsed_payload.flags, self.test_payload_with_data.flags)
        self.assertEqual(parsed_payload.test_data, self.test_payload_with_data.test_data)
        self.assertEqual(parsed_payload.next_payload, self.test_payload_with_data.next_payload)

    def test_compose(self):
        self.assertEqual(self.test_payload_minimal.compose(), self.test_bytes_minimal)
        self.assertEqual(self.test_payload_with_data.compose(), self.test_bytes_with_data)

    def test_error_parse_not_enough_data(self):
        with self.assertRaises(NotEnoughData) as context_manager:
            Ikev2PayloadBaseTest.parse_exact_size(b'\x00')
        self.assertEqual(context_manager.exception.bytes_needed, 3)

    def test_error_payload_validation(self):
        with self.assertRaises(TypeError) as context_manager:
            Ikev2PayloadBaseTest(
                flags={'invalid_flag'},
                test_data=self.test_data
            )
        exception_str = str(context_manager.exception)
        self.assertIn("flags", exception_str)
        self.assertIn("must be", exception_str)
        self.assertIn("Ikev2PayloadFlags", exception_str)

    def test_next_payload(self):
        payload = Ikev2PayloadBaseTest(
            flags={Ikev2PayloadFlags.CRITICAL},
            test_data=self.test_data
        )
        payload.next_payload = Ikev2PayloadType.KE

        composed = payload.compose()
        self.assertEqual(composed[0], Ikev2PayloadType.KE.value.code)

        parsed, _ = Ikev2PayloadBaseTest._parse(composed)  # pylint: disable=protected-access
        self.assertEqual(parsed.next_payload, Ikev2PayloadType.KE)


class TestIkev2PayloadKeyExchange(unittest.TestCase):
    def setUp(self):
        self.dh_group = Ikev2DiffieHellmanGroup.MODP_GROUP_2048_BIT
        self.key_exchange_data = b'\x00\x01\x02\x03\x04\x05\x06\x07'
        self.key_exchange_payload = Ikev2PayloadKeyExchange(
            flags=set(),
            dh_group=self.dh_group,
            key_exchange_data=self.key_exchange_data
        )
        self.key_exchange_payload.next_payload = Ikev2PayloadType.NONE

        self.key_exchange_dict = collections.OrderedDict([
            ('next_payload', b'\x00'),
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x10'),
            ('dh_group', b'\x00\x0e'),
            ('reserved2', b'\x00\x00'),
            ('key_exchange_data', self.key_exchange_data),
        ])
        self.key_exchange_bytes = b''.join(self.key_exchange_dict.values())

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadKeyExchange.get_payload_type(), Ikev2PayloadType.KE)

    def test_parse(self):
        parsed_ke = Ikev2PayloadKeyExchange.parse_exact_size(self.key_exchange_bytes)
        self.assertEqual(parsed_ke.dh_group, self.dh_group)
        self.assertEqual(parsed_ke.key_exchange_data, self.key_exchange_data)

    def test_compose(self):
        composed_bytes = self.key_exchange_payload.compose()
        self.assertGreater(len(composed_bytes), Ikev2PayloadKeyExchange.HEADER_SIZE)

        parsed_ke = Ikev2PayloadKeyExchange.parse_exact_size(composed_bytes)
        self.assertEqual(parsed_ke.key_exchange_data, self.key_exchange_payload.key_exchange_data)

    def test_round_trip(self):
        composed_bytes = self.key_exchange_payload.compose()
        parsed_payload: Ikev2PayloadKeyExchange = Ikev2PayloadKeyExchange.parse_exact_size(composed_bytes)

        self.assertEqual(parsed_payload.dh_group, self.key_exchange_payload.dh_group)
        self.assertEqual(parsed_payload.key_exchange_data, self.key_exchange_payload.key_exchange_data)
        self.assertEqual(parsed_payload.flags, self.key_exchange_payload.flags)
        self.assertEqual(parsed_payload.next_payload, self.key_exchange_payload.next_payload)


class TestIkev2PayloadNonce(unittest.TestCase):  # pylint: disable=too-many-instance-attributes
    def setUp(self):
        self.nonce_data = b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
        self.nonce_payload = Ikev2PayloadNonce(
            flags=set(),
            nonce_data=self.nonce_data
        )
        self.nonce_payload.next_payload = Ikev2PayloadType.NONE

        self.nonce_dict = collections.OrderedDict([
            ('next_payload', b'\x00'),
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x14'),
            ('nonce_data', self.nonce_data),
        ])
        self.nonce_bytes = b''.join(self.nonce_dict.values())

        self.nonce_dict_too_small = collections.OrderedDict([
            ('next_payload', b'\x00'),
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x13'),
            ('nonce_data', b'\x00' * 15),
        ])
        self.nonce_bytes_too_small = b''.join(self.nonce_dict_too_small.values())

        self.nonce_dict_too_large = collections.OrderedDict([
            ('next_payload', b'\x00'),
            ('flags', b'\x00'),
            ('payload_length', b'\x01\x05'),
            ('nonce_data', b'\x00' * 257),
        ])
        self.nonce_bytes_too_large = b''.join(self.nonce_dict_too_large.values())

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadNonce.get_payload_type(), Ikev2PayloadType.NONCE)

    def test_parse(self):
        parsed_nonce = Ikev2PayloadNonce.parse_exact_size(self.nonce_bytes)
        self.assertEqual(parsed_nonce.nonce_data, self.nonce_data)
        self.assertEqual(parsed_nonce.next_payload, Ikev2PayloadType.NONE)
        self.assertEqual(parsed_nonce.flags, set())

    def test_compose(self):
        composed_bytes = self.nonce_payload.compose()
        self.assertIsInstance(composed_bytes, (bytes, bytearray))
        self.assertGreater(len(composed_bytes), Ikev2PayloadNonce.HEADER_SIZE)

        parsed_nonce = Ikev2PayloadNonce.parse_exact_size(composed_bytes)
        self.assertEqual(parsed_nonce.nonce_data, self.nonce_payload.nonce_data)
        self.assertEqual(parsed_nonce.next_payload, self.nonce_payload.next_payload)
        self.assertEqual(parsed_nonce.flags, self.nonce_payload.flags)

    def test_error_invalid_nonce_data(self):
        with self.assertRaises(Exception) as context_manager:
            Ikev2PayloadNonce(
                flags=set(),
                nonce_data="not_bytes"
            )
        self.assertTrue(len(str(context_manager.exception)) > 0)

    def test_minimal_nonce_data(self):
        nonce_empty = Ikev2PayloadNonce(
            flags=set(),
            nonce_data=b'\x00' * 16
        )
        nonce_empty.next_payload = Ikev2PayloadType.NONE
        composed_bytes = nonce_empty.compose()
        parsed_nonce = Ikev2PayloadNonce.parse_exact_size(composed_bytes)
        self.assertEqual(parsed_nonce.nonce_data, b'\x00' * 16)

    def test_too_small_nonce_data(self):
        with self.assertRaises(NotEnoughData) as context_manager:
            Ikev2PayloadNonce.parse_exact_size(self.nonce_bytes_too_small)
        self.assertEqual(context_manager.exception.bytes_needed, 16 - 15)

    def test_too_large_nonce_data(self):
        with self.assertRaises(TooMuchData) as context_manager:
            Ikev2PayloadNonce.parse_exact_size(self.nonce_bytes_too_large)
        self.assertEqual(context_manager.exception.bytes_needed, 257 - 256)

    def test_error_nonce_data_validation_min_length(self):
        with self.assertRaises(ValueError):
            Ikev2PayloadNonce(flags=set(), nonce_data=b'\x01\x02\x03\x04\x05\x06\x07\x08')

    def test_error_nonce_data_validation_max_length(self):
        with self.assertRaises(ValueError):
            Ikev2PayloadNonce(flags=set(), nonce_data=b'\x01' * 257)

    def test_round_trip(self):
        composed_bytes = self.nonce_payload.compose()
        parsed_payload: Ikev2PayloadNonce = Ikev2PayloadNonce.parse_exact_size(composed_bytes)

        self.assertEqual(parsed_payload.nonce_data, self.nonce_payload.nonce_data)
        self.assertEqual(parsed_payload.flags, self.nonce_payload.flags)
        self.assertEqual(parsed_payload.next_payload, self.nonce_payload.next_payload)


class TestIkev2PayloadDelete(unittest.TestCase):  # pylint: disable=too-many-instance-attributes
    def setUp(self):
        self.protocol_id = Ikev2ProtocolId.IKE
        self.spis = [0x1234567890abcdef, 0xfedcba0987654321]
        self.delete_payload = Ikev2PayloadDelete(
            flags=set(),
            protocol_id=self.protocol_id,
            spis=self.spis
        )
        self.delete_payload.next_payload = Ikev2PayloadType.NONE

        self.delete_dict = collections.OrderedDict([
            ('next_payload', b'\x00'),  # NONE
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x18'),  # 24 bytes (header + protocol_id + spi_size + num_spis + spis)
            ('protocol_id', b'\x01'),  # IKE protocol
            ('spi_size', b'\x10'),  # 16 bytes (2 SPIs * 8 bytes each)
            ('num_spis', b'\x00\x02'),  # 2 SPIs
            ('spis', b'\x12\x34\x56\x78\x90\xab\xcd\xef\xfe\xdc\xba\x09\x87\x65\x43\x21'),
        ])
        self.delete_bytes = b''.join(self.delete_dict.values())

        self.empty_delete_payload = Ikev2PayloadDelete(
            flags={Ikev2PayloadFlags.CRITICAL},
            protocol_id=self.protocol_id,
            spis=[]
        )
        self.empty_delete_payload.next_payload = Ikev2PayloadType.NONE

        self.empty_delete_dict = collections.OrderedDict([
            ('next_payload', b'\x00'),  # NONE
            ('flags', b'\x80'),  # CRITICAL (0x80)
            ('payload_length', b'\x00\x08'),  # 8 bytes (header + protocol_id + spi_size + num_spis)
            ('protocol_id', b'\x01'),  # IKE protocol
            ('spi_size', b'\x00'),
            ('num_spis', b'\x00\x00'),
        ])
        self.empty_delete_bytes = b''.join(self.empty_delete_dict.values())

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadDelete.get_payload_type(), Ikev2PayloadType.DELETE)

    def test_parse(self):
        parsed_delete = Ikev2PayloadDelete.parse_exact_size(self.delete_bytes)
        self.assertEqual(parsed_delete.protocol_id, self.protocol_id)
        self.assertEqual(parsed_delete.spis, self.spis)
        self.assertEqual(parsed_delete.next_payload, Ikev2PayloadType.NONE)
        self.assertEqual(parsed_delete.flags, set())

    def test_compose(self):
        composed_bytes = self.delete_payload.compose()
        self.assertIsInstance(composed_bytes, (bytes, bytearray))
        self.assertGreater(len(composed_bytes), Ikev2PayloadDelete.HEADER_SIZE)

        parsed_delete = Ikev2PayloadDelete.parse_exact_size(composed_bytes)
        self.assertEqual(parsed_delete.protocol_id, self.delete_payload.protocol_id)
        self.assertEqual(parsed_delete.spis, self.delete_payload.spis)
        self.assertEqual(parsed_delete.next_payload, self.delete_payload.next_payload)
        self.assertEqual(parsed_delete.flags, self.delete_payload.flags)

    def test_empty_spis(self):
        composed_bytes = self.empty_delete_payload.compose()
        parsed_delete = Ikev2PayloadDelete.parse_exact_size(composed_bytes)
        self.assertEqual(parsed_delete.spis, [])
        self.assertEqual(parsed_delete.protocol_id, self.protocol_id)
        self.assertEqual(parsed_delete.flags, {Ikev2PayloadFlags.CRITICAL})

    def test_error_parse_not_enough_data(self):
        with self.assertRaises(NotEnoughData) as context_manager:
            Ikev2PayloadDelete.parse_exact_size(b'\x00\x01\x02')
        self.assertEqual(context_manager.exception.bytes_needed, 1)

    def test_error_invalid_protocol_id(self):
        with self.assertRaises(TypeError):
            Ikev2PayloadDelete(
                flags=set(),
                protocol_id="invalid_protocol",
                spis=self.spis
            )

    def test_error_invalid_spis_type(self):
        with self.assertRaises(TypeError):
            Ikev2PayloadDelete(
                flags=set(),
                protocol_id=self.protocol_id,
                spis=["not_an_int"]
            )

    def test_round_trip(self):
        composed_bytes = self.delete_payload.compose()
        parsed_payload: Ikev2PayloadDelete = Ikev2PayloadDelete.parse_exact_size(composed_bytes)

        self.assertEqual(parsed_payload.protocol_id, self.delete_payload.protocol_id)
        self.assertEqual(parsed_payload.spis, self.delete_payload.spis)
        self.assertEqual(parsed_payload.flags, self.delete_payload.flags)
        self.assertEqual(parsed_payload.next_payload, self.delete_payload.next_payload)


class TestIkev2PayloadNotifyNoData(unittest.TestCase):  # pylint: disable=too-many-instance-attributes
    def setUp(self):
        self.protocol_id = Ikev2ProtocolId.IKE
        self.notify_type = Ikev2NotifyType.AUTHENTICATION_FAILED
        self.spi = b'\x00\x01\x02\x03\x04\x05\x06\x07'

        self.notify_payload_minimal = Ikev2PayloadNotifyNoDataTest(
            flags=set(),
            protocol_id=self.protocol_id,
            notify_type=self.notify_type,
            spi=b''
        )
        self.notify_payload_minimal.next_payload = Ikev2PayloadType.NONE

        self.notify_dict_minimal = collections.OrderedDict([
            ('next_payload', b'\x00'),  # NONE
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x08'),  # 8 bytes (header + protocol_id + spi_size + notify_type)
            ('protocol_id', b'\x01'),  # IKE protocol
            ('spi_size', b'\x00'),
            ('notify_type', b'\x00\x18'),  # AUTHENTICATION_FAILED (0x0018)
        ])
        self.notify_bytes_minimal = b''.join(self.notify_dict_minimal.values())

        self.notify_payload_with_spi = Ikev2PayloadNotifyNoDataTest(
            flags={Ikev2PayloadFlags.CRITICAL},
            protocol_id=self.protocol_id,
            notify_type=self.notify_type,
            spi=self.spi
        )
        self.notify_payload_with_spi.next_payload = Ikev2PayloadType.NONE

        self.notify_dict_with_spi = collections.OrderedDict([
            ('next_payload', b'\x00'),  # NONE
            ('flags', b'\x80'),  # CRITICAL (0x80)
            ('payload_length', b'\x00\x10'),  # 16 bytes (header + protocol_id + spi_size + notify_type + spi)
            ('protocol_id', b'\x01'),  # IKE protocol
            ('spi_size', b'\x08'),
            ('notify_type', b'\x00\x18'),  # AUTHENTICATION_FAILED (0x0018)
            ('spi', self.spi),
        ])
        self.notify_bytes_with_spi = b''.join(self.notify_dict_with_spi.values())

    def test_parse(self):
        parsed_notify = Ikev2PayloadNotifyNoDataTest.parse_exact_size(self.notify_bytes_minimal)
        self.assertEqual(parsed_notify.protocol_id, self.protocol_id)
        self.assertEqual(parsed_notify.type, self.notify_type)
        self.assertEqual(parsed_notify.spi, b'')
        self.assertEqual(parsed_notify.next_payload, Ikev2PayloadType.NONE)
        self.assertEqual(parsed_notify.flags, set())

        parsed_notify = Ikev2PayloadNotifyNoDataTest.parse_exact_size(self.notify_bytes_with_spi)
        self.assertEqual(parsed_notify.protocol_id, self.protocol_id)
        self.assertEqual(parsed_notify.type, self.notify_type)
        self.assertEqual(parsed_notify.spi, self.spi)
        self.assertEqual(parsed_notify.next_payload, Ikev2PayloadType.NONE)
        self.assertEqual(parsed_notify.flags, {Ikev2PayloadFlags.CRITICAL})

    def test_compose(self):
        self.assertEqual(self.notify_payload_minimal.compose(), self.notify_bytes_minimal)
        self.assertEqual(self.notify_payload_with_spi.compose(), self.notify_bytes_with_spi)

    def test_error_parse_not_enough_data(self):
        with self.assertRaises(NotEnoughData) as context_manager:
            Ikev2PayloadNotifyNoDataTest.parse_exact_size(b'\x00\x01\x02')
        self.assertEqual(context_manager.exception.bytes_needed, 1)

    def test_error_payload_validation(self):
        with self.assertRaises(TypeError) as context_manager:
            Ikev2PayloadNotifyNoDataTest(
                flags={'invalid_flag'},
                protocol_id=self.protocol_id,
                notify_type=self.notify_type,
                spi=self.spi
            )
        exception_str = str(context_manager.exception)
        self.assertIn("flags", exception_str)
        self.assertIn("must be", exception_str)
        self.assertIn("Ikev2PayloadFlags", exception_str)

    def test_error_invalid_protocol_id(self):
        with self.assertRaises(TypeError):
            Ikev2PayloadNotifyNoDataTest(
                flags=set(),
                protocol_id="invalid_protocol",
                notify_type=self.notify_type,
                spi=self.spi
            )

    def test_error_invalid_notify_type(self):
        with self.assertRaises(TypeError):
            Ikev2PayloadNotifyNoDataTest(
                flags=set(),
                protocol_id=self.protocol_id,
                notify_type="invalid_type",
                spi=self.spi
            )

    def test_error_invalid_spi(self):
        # The spi attribute has converter=bytes, so test with something that can't be converted
        with self.assertRaises((TypeError, ValueError)):
            Ikev2PayloadNotifyNoDataTest(
                flags=set(),
                protocol_id=self.protocol_id,
                notify_type=self.notify_type,
                spi=None
            )


class TestIkev2PayloadNotifyAuthenticationFailed(unittest.TestCase):
    def setUp(self):
        self.protocol_id = Ikev2ProtocolId.IKE

        self.notify_dict_authentication_failed = collections.OrderedDict([
            ('next_payload', b'\x00'),  # NONE
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x08'),  # 8 bytes (header + protocol_id + spi_size + notify_type)
            ('protocol_id', b'\x01'),  # IKE protocol
            ('spi_size', b'\x00'),
            ('notify_type', b'\x00\x18'),  # AUTHENTICATION_FAILED (0x0018)
        ])
        self.notify_bytes_authentication_failed = b''.join(self.notify_dict_authentication_failed.values())

    def test_get_message_type(self):
        self.assertEqual(Ikev2PayloadNotifyAuthenticationFailed._get_message_type(),  # pylint: disable=protected-access
                         Ikev2NotifyType.AUTHENTICATION_FAILED)

    def test_parse_authentication_failed(self):
        parsed_notify = Ikev2PayloadNotifyAuthenticationFailed.parse_exact_size(
            self.notify_bytes_authentication_failed)
        self.assertEqual(parsed_notify.type, Ikev2NotifyType.AUTHENTICATION_FAILED)
        self.assertEqual(parsed_notify.protocol_id, self.protocol_id)


class TestIkev2PayloadNotifyUnparsed(unittest.TestCase):
    def setUp(self):
        self.notify_type = Ikev2NotifyType.INVALID_SYNTAX
        self.notify_data = (
            b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
            b'\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f'
        )

        self.notify_dict_with_data = collections.OrderedDict([
            ('next_payload', b'\x00'),  # NONE
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x28'),  # 40 bytes (8 header + 32 data)
            ('protocol_id', b'\x01'),  # IKE protocol
            ('spi_size', b'\x00'),
            ('notify_type', b'\x00\x07'),  # INVALID_SYNTAX (0x0007)
        ])
        self.notify_bytes_with_data = b''.join(self.notify_dict_with_data.values()) + self.notify_data

    def test_parse(self):
        parsed_notify: Ikev2PayloadNotifyUnparsed = Ikev2PayloadNotifyUnparsed.parse_exact_size(
            self.notify_bytes_with_data)
        self.assertEqual(parsed_notify.data, self.notify_data)  # pylint: disable=no-member

    def test_compose(self):
        notify_payload = Ikev2PayloadNotifyUnparsed(
            flags=set(),
            protocol_id=Ikev2ProtocolId.IKE,
            type=self.notify_type,
            spi=b'',
            data=self.notify_data
        )
        notify_payload.next_payload = Ikev2PayloadType.NONE

        composed_bytes = notify_payload.compose()
        self.assertEqual(composed_bytes, self.notify_bytes_with_data)

    def test_round_trip(self):
        original_payload = Ikev2PayloadNotifyUnparsed(
            flags=set(),
            protocol_id=Ikev2ProtocolId.IKE,
            type=self.notify_type,
            spi=b'',
            data=self.notify_data
        )
        original_payload.next_payload = Ikev2PayloadType.NONE

        composed_bytes = original_payload.compose()
        parsed_payload: Ikev2PayloadNotifyUnparsed = Ikev2PayloadNotifyUnparsed.parse_exact_size(composed_bytes)

        self.assertEqual(parsed_payload.data, original_payload.data)  # pylint: disable=no-member


class TestIkev2NotifyPayloadInvalidKe(unittest.TestCase):
    def setUp(self):
        self.dh_group = Ikev2DiffieHellmanGroup.MODP_GROUP_2048_BIT

        self.invalid_ke_dict = collections.OrderedDict([
            ('next_payload', b'\x00'),  # NONE
            ('flags', b'\x00'),
            ('payload_length', b'\x00\x0a'),  # 10 bytes (8 header + 2 dh_group)
            ('protocol_id', b'\x01'),  # IKE protocol
            ('spi_size', b'\x00'),
            ('notify_type', b'\x00\x11'),  # INVALID_KE_PAYLOAD (0x0011)
            ('dh_group', b'\x00\x0e'),  # MODP_GROUP_2048_BIT (0x000e)
        ])
        self.invalid_ke_bytes = b''.join(self.invalid_ke_dict.values())
        self.invalid_ke_payload = Ikev2NotifyPayloadInvalidKe(
            flags=set(),
            protocol_id=Ikev2ProtocolId.IKE,
            type=Ikev2NotifyType.INVALID_KE_PAYLOAD,
            spi=b'',
            dh_group=self.dh_group
        )
        self.invalid_ke_payload.next_payload = Ikev2PayloadType.NONE

    def test_get_message_type(self):
        self.assertEqual(Ikev2NotifyPayloadInvalidKe._get_message_type(),  # pylint: disable=protected-access
                         Ikev2NotifyType.INVALID_KE_PAYLOAD)

    def test_parse(self):
        parsed_notify: Ikev2NotifyPayloadInvalidKe = Ikev2NotifyPayloadInvalidKe.parse_exact_size(self.invalid_ke_bytes)
        self.assertEqual(parsed_notify.type, Ikev2NotifyType.INVALID_KE_PAYLOAD)
        self.assertEqual(parsed_notify.protocol_id, Ikev2ProtocolId.IKE)
        self.assertEqual(parsed_notify.dh_group, self.dh_group)  # pylint: disable=no-member

    def test_compose(self):
        composed_bytes = self.invalid_ke_payload.compose()
        self.assertEqual(composed_bytes, self.invalid_ke_bytes)

    def test_round_trip(self):
        composed_bytes = self.invalid_ke_payload.compose()
        parsed_payload: Ikev2NotifyPayloadInvalidKe = Ikev2NotifyPayloadInvalidKe.parse_exact_size(composed_bytes)

        # Verify all attributes are preserved
        self.assertEqual(parsed_payload.dh_group, self.invalid_ke_payload.dh_group)  # pylint: disable=no-member
        self.assertEqual(parsed_payload.flags, self.invalid_ke_payload.flags)
        self.assertEqual(parsed_payload.next_payload, self.invalid_ke_payload.next_payload)
        self.assertEqual(parsed_payload.spi, self.invalid_ke_payload.spi)


class TestIkev2PayloadCertificateRequest(unittest.TestCase):
    _CERT_ENCODING = Ikev2CertificateType.X509_CERTIFICATE_SIGNATURE
    _CERTIFICATION_AUTHORITY = (
        b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
        b'\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f'
    )
    _CERTREQ_BYTES = bytes.fromhex(
        '00'      # next_payload = NONE
        '00'      # flags
        '0025'    # payload_length = 37 (4 header + 1 cert_encoding + 32 certification authority)
        '04'      # cert_encoding = X.509 certificate signature
    ) + _CERTIFICATION_AUTHORITY

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadCertificateRequest.get_payload_type(), Ikev2PayloadType.CERTREQ)

    def test_parse(self):
        parsed_certreq: Ikev2PayloadCertificateRequest = Ikev2PayloadCertificateRequest.parse_exact_size(
            self._CERTREQ_BYTES)
        self.assertEqual(parsed_certreq.next_payload, Ikev2PayloadType.NONE)
        self.assertEqual(parsed_certreq.flags, set())
        self.assertEqual(parsed_certreq.cert_encoding, self._CERT_ENCODING)
        self.assertEqual(parsed_certreq.certification_authority, self._CERTIFICATION_AUTHORITY)

    def test_compose(self):
        certreq_payload = Ikev2PayloadCertificateRequest(
            flags=set(),
            cert_encoding=self._CERT_ENCODING,
            certification_authority=self._CERTIFICATION_AUTHORITY,
        )
        certreq_payload.next_payload = Ikev2PayloadType.NONE
        self.assertEqual(certreq_payload.compose(), self._CERTREQ_BYTES)

    def test_round_trip(self):
        certreq_payload = Ikev2PayloadCertificateRequest(
            flags=set(),
            cert_encoding=self._CERT_ENCODING,
            certification_authority=self._CERTIFICATION_AUTHORITY,
        )
        certreq_payload.next_payload = Ikev2PayloadType.NONE
        composed_bytes = certreq_payload.compose()
        parsed_payload: Ikev2PayloadCertificateRequest = Ikev2PayloadCertificateRequest.parse_exact_size(
            composed_bytes)

        self.assertEqual(parsed_payload.cert_encoding, certreq_payload.cert_encoding)
        self.assertEqual(parsed_payload.certification_authority, certreq_payload.certification_authority)
        self.assertEqual(parsed_payload.flags, certreq_payload.flags)
        self.assertEqual(parsed_payload.next_payload, certreq_payload.next_payload)

    def test_error_parse_not_enough_data(self):
        incomplete_data = self._CERTREQ_BYTES[:-5]

        with self.assertRaises(NotEnoughData) as context_manager:
            Ikev2PayloadCertificateRequest.parse_exact_size(incomplete_data)
        self.assertEqual(context_manager.exception.bytes_needed, 5)

    def test_get_authority_hashes_returns_empty_list_for_empty(self):
        payload = Ikev2PayloadCertificateRequest(
            flags=set(), cert_encoding=self._CERT_ENCODING, certification_authority=b'',
        )
        self.assertEqual(payload.get_authority_hashes(), [])

    def test_get_authority_hashes_splits_single_20_byte_hash(self):
        hash_a = b'\x00' * 20
        payload = Ikev2PayloadCertificateRequest(
            flags=set(), cert_encoding=self._CERT_ENCODING, certification_authority=hash_a,
        )
        self.assertEqual(payload.get_authority_hashes(), [hash_a])

    def test_get_authority_hashes_splits_multiple_concatenated_hashes(self):
        hash_a = b'\x00' * 20
        hash_b = b'\xff' * 20
        hash_c = b'\x42' * 20
        payload = Ikev2PayloadCertificateRequest(
            flags=set(),
            cert_encoding=self._CERT_ENCODING,
            certification_authority=hash_a + hash_b + hash_c,
        )
        self.assertEqual(payload.get_authority_hashes(), [hash_a, hash_b, hash_c])

    def test_get_authority_hashes_drops_trailing_partial_chunk(self):
        hash_a = b'\x00' * 20
        payload = Ikev2PayloadCertificateRequest(
            flags=set(),
            cert_encoding=self._CERT_ENCODING,
            certification_authority=hash_a + b'\xaa\xbb\xcc',  # 3 stray octets
        )
        self.assertEqual(payload.get_authority_hashes(), [hash_a])

    def test_get_authority_hashes_accepts_bytearray_certification_authority(self):
        hash_a = b'\x00' * 20
        payload = Ikev2PayloadCertificateRequest(
            flags=set(),
            cert_encoding=self._CERT_ENCODING,
            certification_authority=bytearray(hash_a),
        )
        self.assertEqual(payload.get_authority_hashes(), [hash_a])


class TestIkev2PayloadCertificate(unittest.TestCase):
    _CERT_ENCODING = Ikev2CertificateType.X509_CERTIFICATE_SIGNATURE
    _CERTIFICATE_DATA = (
        b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
        b'\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f'
    )
    _CERT_BYTES = bytes.fromhex(
        '00'      # next_payload = NONE
        '00'      # flags
        '0025'    # payload_length = 37 (4 header + 1 cert_encoding + 32 data)
        '04'      # cert_encoding = X.509 certificate signature
    ) + _CERTIFICATE_DATA

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadCertificate.get_payload_type(), Ikev2PayloadType.CERT)

    def test_parse(self):
        parsed_cert: Ikev2PayloadCertificate = Ikev2PayloadCertificate.parse_exact_size(self._CERT_BYTES)
        self.assertEqual(parsed_cert.next_payload, Ikev2PayloadType.NONE)
        self.assertEqual(parsed_cert.flags, set())
        self.assertEqual(parsed_cert.cert_encoding, self._CERT_ENCODING)
        self.assertEqual(parsed_cert.certificate_data, self._CERTIFICATE_DATA)

    def test_compose(self):
        cert_payload = Ikev2PayloadCertificate(
            flags=set(),
            cert_encoding=self._CERT_ENCODING,
            certificate_data=self._CERTIFICATE_DATA,
        )
        cert_payload.next_payload = Ikev2PayloadType.NONE
        self.assertEqual(cert_payload.compose(), self._CERT_BYTES)

    def test_round_trip(self):
        cert_payload = Ikev2PayloadCertificate(
            flags=set(),
            cert_encoding=self._CERT_ENCODING,
            certificate_data=self._CERTIFICATE_DATA,
        )
        cert_payload.next_payload = Ikev2PayloadType.NONE
        composed_bytes = cert_payload.compose()
        parsed_payload: Ikev2PayloadCertificate = Ikev2PayloadCertificate.parse_exact_size(composed_bytes)

        self.assertEqual(parsed_payload.cert_encoding, cert_payload.cert_encoding)
        self.assertEqual(parsed_payload.certificate_data, cert_payload.certificate_data)
        self.assertEqual(parsed_payload.flags, cert_payload.flags)
        self.assertEqual(parsed_payload.next_payload, cert_payload.next_payload)

    def test_error_parse_not_enough_data(self):
        incomplete_data = self._CERT_BYTES[:-5]

        with self.assertRaises(NotEnoughData) as context_manager:
            Ikev2PayloadCertificate.parse_exact_size(incomplete_data)
        self.assertEqual(context_manager.exception.bytes_needed, 5)


class TestIkev2PayloadIdentificationInitiator(unittest.TestCase):
    _IDENTIFIER = 'scanner@cryptolyzer'
    _ID_BYTES = bytes.fromhex(
        '00'        # next_payload = NONE
        '00'        # flags
        '001b'      # payload_length = 27 (4 header + 1 id_type + 3 reserved + 19 id_data)
        '03'        # id_type = RFC822_ADDR
        '000000'    # reserved
    ) + _IDENTIFIER.encode('ascii')

    def test_get_payload_type(self):
        self.assertEqual(
            Ikev2PayloadIdentificationInitiatorRfc822Addr.get_payload_type(),
            Ikev2PayloadType.IDI,
        )

    def test_parse(self):
        parsed = Ikev2PayloadIdentificationInitiatorRfc822Addr.parse_exact_size(self._ID_BYTES)
        self.assertEqual(parsed.next_payload, Ikev2PayloadType.NONE)
        self.assertEqual(parsed.flags, set())
        self.assertEqual(parsed.id_type, Ikev2IdType.RFC822_ADDR)
        self.assertEqual(parsed.identifier, self._IDENTIFIER)  # pylint: disable=no-member

    def test_compose(self):
        id_payload = Ikev2PayloadIdentificationInitiatorRfc822Addr(
            flags=set(),
            identifier=self._IDENTIFIER,
        )
        id_payload.next_payload = Ikev2PayloadType.NONE
        self.assertEqual(id_payload.compose(), self._ID_BYTES)

    def test_round_trip(self):
        id_payload = Ikev2PayloadIdentificationInitiatorRfc822Addr(
            flags=set(),
            identifier=self._IDENTIFIER,
        )
        id_payload.next_payload = Ikev2PayloadType.NONE
        composed = id_payload.compose()
        parsed = Ikev2PayloadIdentificationInitiatorRfc822Addr.parse_exact_size(composed)

        self.assertEqual(parsed.id_type, id_payload.id_type)
        identifier_parsed = parsed.identifier  # pylint: disable=no-member
        identifier_expected = id_payload.identifier  # pylint: disable=no-member
        self.assertEqual(identifier_parsed, identifier_expected)
        self.assertEqual(parsed.flags, id_payload.flags)
        self.assertEqual(parsed.next_payload, id_payload.next_payload)


class TestIkev2PayloadIdentificationResponder(unittest.TestCase):
    _IDENTIFIER = ipaddress.IPv4Address('192.0.2.1')

    def test_get_payload_type(self):
        self.assertEqual(
            Ikev2PayloadIdentificationResponderIpv4Addr.get_payload_type(),
            Ikev2PayloadType.IDR,
        )

    def test_round_trip(self):
        id_payload = Ikev2PayloadIdentificationResponderIpv4Addr(
            flags=set(),
            identifier=self._IDENTIFIER,
        )
        id_payload.next_payload = Ikev2PayloadType.NONE
        composed = id_payload.compose()
        parsed = Ikev2PayloadIdentificationResponderIpv4Addr.parse_exact_size(composed)
        self.assertEqual(parsed.id_type, Ikev2IdType.IPV4_ADDR)
        self.assertEqual(parsed.identifier, self._IDENTIFIER)  # pylint: disable=no-member

    def test_round_trip_der_asn1_dn(self):
        dn_bytes = (
            b'\x30\x1a\x31\x18\x30\x16\x06\x03\x55\x04\x03\x0c\x0f'
            b'gateway.test.org'
        )
        payload = Ikev2PayloadIdentificationResponderDerAsn1Dn(
            flags=set(),
            identifier=dn_bytes,
        )
        payload.next_payload = Ikev2PayloadType.NONE

        composed = payload.compose()
        parsed = Ikev2PayloadIdentificationResponderDerAsn1Dn.parse_exact_size(composed)
        self.assertEqual(parsed.id_type, Ikev2IdType.DER_ASN1_DN)
        self.assertEqual(parsed.identifier, dn_bytes)  # pylint: disable=no-member

    def test_ipv4_decode_rejects_wrong_length(self):
        with self.assertRaises(InvalidType):
            Ikev2PayloadIdentificationResponderIpv4Addr._decode_identifier(  # pylint: disable=protected-access
                b'\xc0\x00\x02',
            )

    def test_ipv6_decode_rejects_wrong_length(self):
        with self.assertRaises(InvalidType):
            Ikev2PayloadIdentificationResponderIpv6Addr._decode_identifier(  # pylint: disable=protected-access
                b'\x00' * 15,
            )


class TestIkev2PayloadIdentificationVariants(unittest.TestCase):
    def test_initiator_fqdn_round_trip(self):
        payload = Ikev2PayloadIdentificationInitiatorFqdn(
            flags=set(), identifier='gateway.example.com',
        )
        payload.next_payload = Ikev2PayloadType.NONE
        parsed = Ikev2PayloadIdentificationInitiatorFqdn.parse_exact_size(payload.compose())
        self.assertEqual(parsed.id_type, Ikev2IdType.FQDN)
        self.assertEqual(parsed.identifier, 'gateway.example.com')  # pylint: disable=no-member

    def test_initiator_ipv6_round_trip(self):
        addr = ipaddress.IPv6Address('2001:db8::1')
        payload = Ikev2PayloadIdentificationInitiatorIpv6Addr(
            flags=set(), identifier=addr,
        )
        payload.next_payload = Ikev2PayloadType.NONE
        parsed = Ikev2PayloadIdentificationInitiatorIpv6Addr.parse_exact_size(payload.compose())
        self.assertEqual(parsed.id_type, Ikev2IdType.IPV6_ADDR)
        self.assertEqual(parsed.identifier, addr)  # pylint: disable=no-member

    def test_initiator_key_id_round_trip(self):
        payload = Ikev2PayloadIdentificationInitiatorKeyId(
            flags=set(), identifier=b'\x01\x02\x03',
        )
        payload.next_payload = Ikev2PayloadType.NONE
        parsed = Ikev2PayloadIdentificationInitiatorKeyId.parse_exact_size(payload.compose())
        self.assertEqual(parsed.id_type, Ikev2IdType.KEY_ID)
        self.assertEqual(parsed.identifier, b'\x01\x02\x03')  # pylint: disable=no-member

    def test_initiator_der_asn1_gn_round_trip(self):
        payload = Ikev2PayloadIdentificationInitiatorDerAsn1Gn(
            flags=set(), identifier=b'\xab\xcd\xef',
        )
        payload.next_payload = Ikev2PayloadType.NONE
        parsed = Ikev2PayloadIdentificationInitiatorDerAsn1Gn.parse_exact_size(payload.compose())
        self.assertEqual(parsed.id_type, Ikev2IdType.DER_ASN1_GN)
        self.assertEqual(parsed.identifier, b'\xab\xcd\xef')  # pylint: disable=no-member

    def test_initiator_fc_name_round_trip(self):
        payload = Ikev2PayloadIdentificationInitiatorFcName(
            flags=set(), identifier=b'\x10\x20\x30\x40\x50\x60\x70\x80',
        )
        payload.next_payload = Ikev2PayloadType.NONE
        parsed = Ikev2PayloadIdentificationInitiatorFcName.parse_exact_size(payload.compose())
        self.assertEqual(parsed.id_type, Ikev2IdType.FC_NAME)
        self.assertEqual(parsed.identifier, b'\x10\x20\x30\x40\x50\x60\x70\x80')  # pylint: disable=no-member

    def test_initiator_null_round_trip(self):
        payload = Ikev2PayloadIdentificationInitiatorNull(
            flags=set(), identifier=b'',
        )
        payload.next_payload = Ikev2PayloadType.NONE
        parsed = Ikev2PayloadIdentificationInitiatorNull.parse_exact_size(payload.compose())
        self.assertEqual(parsed.id_type, Ikev2IdType.NULL)
        self.assertEqual(parsed.identifier, b'')  # pylint: disable=no-member

    def test_id_payload_rejects_negative_id_data_length(self):
        # payload_length=6 → id_data_length = 6 - 4 - 4 = -2 → NotEnoughData
        bogus = b'\x00\x00\x00\x06\x02\x00\x00\x00\x00\x00'
        with self.assertRaises(NotEnoughData):
            Ikev2PayloadIdentificationInitiatorFqdn.parse_exact_size(bogus)

    def test_id_payload_rejects_id_type_mismatch(self):
        # Wire id_type = 0x01 (IPV4_ADDR); caller uses FQDN subclass.
        wire = b'\x00\x00\x00\x0c\x01\x00\x00\x00\xc0\x00\x02\x01'
        with self.assertRaises(InvalidType):
            Ikev2PayloadIdentificationInitiatorFqdn.parse_exact_size(wire)

    def test_initiator_variant_dispatch_picks_subclass(self):
        source = Ikev2PayloadIdentificationInitiatorFqdn(
            flags=set(), identifier='initiator.example',
        )
        source.next_payload = Ikev2PayloadType.NONE
        parsed = Ikev2PayloadIdentificationInitiator.parse_exact_size(source.compose())
        self.assertIsInstance(parsed, Ikev2PayloadIdentificationInitiatorFqdn)

    def test_responder_variant_dispatch_picks_subclass(self):
        source = Ikev2PayloadIdentificationResponderFqdn(
            flags=set(), identifier='responder.example',
        )
        source.next_payload = Ikev2PayloadType.NONE
        parsed = Ikev2PayloadIdentificationResponder.parse_exact_size(source.compose())
        self.assertIsInstance(parsed, Ikev2PayloadIdentificationResponderFqdn)

    def test_rfc822_mixin_returns_none_for_ikev1(self):
        self.assertIsNone(Ikev2PayloadIdentificationInitiatorRfc822Addr.get_id_type_ikev1())

    def test_fc_name_mixin_returns_none_for_ikev1(self):
        self.assertIsNone(Ikev2PayloadIdentificationInitiatorFcName.get_id_type_ikev1())

    def test_null_mixin_returns_none_for_ikev1(self):
        self.assertIsNone(Ikev2PayloadIdentificationInitiatorNull.get_id_type_ikev1())


class TestIkev2PayloadAuthentication(unittest.TestCase):
    _AUTH_METHOD = Ikev2AuthenticationMethod.RSA_DIGITAL_SIGNATURE
    _AUTH_DATA = b'\xab' * 256  # 2048-bit RSA signature
    _AUTH_BYTES = bytes.fromhex(
        '00'        # next_payload = NONE
        '00'        # flags
        '0108'      # payload_length = 264 (4 header + 1 auth_method + 3 reserved + 256 auth_data)
        '01'        # auth_method = RSA_DIGITAL_SIGNATURE
        '000000'    # reserved
    ) + _AUTH_DATA

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadAuthentication.get_payload_type(), Ikev2PayloadType.AUTH)

    def test_parse(self):
        parsed: Ikev2PayloadAuthentication = Ikev2PayloadAuthentication.parse_exact_size(self._AUTH_BYTES)
        self.assertEqual(parsed.auth_method, self._AUTH_METHOD)
        self.assertEqual(parsed.auth_data, self._AUTH_DATA)

    def test_round_trip_rsa(self):
        auth_payload = Ikev2PayloadAuthentication(
            flags=set(),
            auth_method=self._AUTH_METHOD,
            auth_data=self._AUTH_DATA,
        )
        auth_payload.next_payload = Ikev2PayloadType.NONE
        composed = auth_payload.compose()
        parsed: Ikev2PayloadAuthentication = Ikev2PayloadAuthentication.parse_exact_size(composed)
        self.assertEqual(parsed.auth_method, self._AUTH_METHOD)
        self.assertEqual(parsed.auth_data, self._AUTH_DATA)

    def test_round_trip_ecdsa_256(self):
        payload = Ikev2PayloadAuthentication(
            flags=set(),
            auth_method=Ikev2AuthenticationMethod.ECDSA_SHA_256_P_256,
            auth_data=b'\xcd' * 64,  # r||s for P-256
        )
        payload.next_payload = Ikev2PayloadType.NONE
        composed = payload.compose()
        parsed: Ikev2PayloadAuthentication = Ikev2PayloadAuthentication.parse_exact_size(composed)
        self.assertEqual(parsed.auth_method, Ikev2AuthenticationMethod.ECDSA_SHA_256_P_256)
        self.assertEqual(len(parsed.auth_data), 64)

    def test_error_parse_negative_auth_data_length(self):
        # payload_length=6 → auth_data_length = 6 - 4 - 4 = -2 → NotEnoughData
        bogus = b'\x00\x00\x00\x06\x01\x00\x00\x00\x00\x00'
        with self.assertRaises(NotEnoughData):
            Ikev2PayloadAuthentication.parse_exact_size(bogus)

    def test_round_trip_null(self):
        payload = Ikev2PayloadAuthentication(
            flags=set(),
            auth_method=Ikev2AuthenticationMethod.NULL_AUTHENTICATION,
            auth_data=b'\xef' * 32,  # HMAC(SK_pi, ...) output sized to the prf
        )
        payload.next_payload = Ikev2PayloadType.NONE
        composed = payload.compose()
        parsed: Ikev2PayloadAuthentication = Ikev2PayloadAuthentication.parse_exact_size(composed)
        self.assertEqual(parsed.auth_method, Ikev2AuthenticationMethod.NULL_AUTHENTICATION)


class TestIkev2AuthDigitalSignatureEnvelope(unittest.TestCase):
    # AlgorithmIdentifier for ecdsa-with-SHA256 (OID 1.2.840.10045.4.3.2),
    # encoded as a DER SEQUENCE so asn1crypto can decode the OID.
    _SIGNATURE_ALGORITHM = Signature.ECDSA_WITH_SHA2_256
    _ALGORITHM_IDENTIFIER = bytes.fromhex('300a06082a8648ce3d040302')
    _SIGNATURE = b'\xab' * 64  # P-256 r||s
    _ENVELOPE_BYTES = bytes([len(_ALGORITHM_IDENTIFIER)]) + _ALGORITHM_IDENTIFIER + _SIGNATURE

    def test_round_trip(self):
        envelope = Ikev2AuthDigitalSignatureEnvelope(
            signature_algorithm=self._SIGNATURE_ALGORITHM,
            signature=self._SIGNATURE,
        )
        composed = envelope.compose()
        self.assertEqual(composed, self._ENVELOPE_BYTES)

        parsed = Ikev2AuthDigitalSignatureEnvelope.from_bytes(composed)
        self.assertEqual(parsed.signature_algorithm, self._SIGNATURE_ALGORITHM)
        self.assertEqual(parsed.signature, self._SIGNATURE)

    def test_from_auth_payload(self):
        payload = Ikev2PayloadAuthentication(
            flags=set(),
            auth_method=Ikev2AuthenticationMethod.DIGITAL_SIGNATURE,
            auth_data=self._ENVELOPE_BYTES,
        )
        payload.next_payload = Ikev2PayloadType.NONE

        envelope: Ikev2AuthDigitalSignatureEnvelope = payload.parse_digital_signature_envelope()
        self.assertEqual(envelope.signature_algorithm, self._SIGNATURE_ALGORITHM)
        self.assertEqual(envelope.signature, self._SIGNATURE)

    def test_wrong_auth_method_rejected(self):
        payload = Ikev2PayloadAuthentication(
            flags=set(),
            auth_method=Ikev2AuthenticationMethod.RSA_DIGITAL_SIGNATURE,
            auth_data=self._ENVELOPE_BYTES,
        )
        payload.next_payload = Ikev2PayloadType.NONE

        with self.assertRaises(InvalidType):
            payload.parse_digital_signature_envelope()

    def test_parse_unknown_oid_raises_invalid_value(self):
        # AlgorithmIdentifier for the made-up OID 1.2.3.4.5 (not in
        # cryptodatahub's Signature registry).
        # DER: SEQUENCE(0x30) len=6 OID(0x06) len=4 first-two(0x2A) 03 04 05
        algorithm_identifier = bytes.fromhex('300606042a030405')
        envelope_bytes = (
            bytes([len(algorithm_identifier)]) + algorithm_identifier + b'\xcc' * 32
        )
        with self.assertRaises(InvalidValue):
            Ikev2AuthDigitalSignatureEnvelope.from_bytes(envelope_bytes)


class TestIkev2PayloadEncryptedAndAuthenticated(unittest.TestCase):
    # 16-octet IV + 16 octets of "encrypted" body + 1 octet pad length + 16 octet ICV
    _ENCRYPTED_DATA = b'\x11' * 16 + b'\x22' * 16 + b'\x00' + b'\x33' * 16
    _SK_BYTES = bytes.fromhex(
        '23'      # next_payload = IDI (35 = inner first payload after decrypt)
        '00'      # flags
        '0035'    # payload_length = 53 (4 header + 49 encrypted_data)
    ) + _ENCRYPTED_DATA

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadEncryptedAndAuthenticated.get_payload_type(), Ikev2PayloadType.SK)

    def test_parse(self):
        parsed: Ikev2PayloadEncryptedAndAuthenticated
        parsed = Ikev2PayloadEncryptedAndAuthenticated.parse_exact_size(self._SK_BYTES)
        self.assertEqual(parsed.next_payload, Ikev2PayloadType.IDI)
        self.assertEqual(parsed.flags, set())
        self.assertEqual(parsed.encrypted_data, self._ENCRYPTED_DATA)

    def test_round_trip(self):
        sk_payload = Ikev2PayloadEncryptedAndAuthenticated(
            flags=set(),
            encrypted_data=self._ENCRYPTED_DATA,
        )
        sk_payload.next_payload = Ikev2PayloadType.IDI
        composed = sk_payload.compose()
        self.assertEqual(composed, self._SK_BYTES)

        parsed: Ikev2PayloadEncryptedAndAuthenticated
        parsed = Ikev2PayloadEncryptedAndAuthenticated.parse_exact_size(composed)
        self.assertEqual(parsed.encrypted_data, self._ENCRYPTED_DATA)
        self.assertEqual(parsed.next_payload, Ikev2PayloadType.IDI)


class TestIkev2PayloadEap(unittest.TestCase):
    _EAP_DATA = b'\x01\x01\x00\x05\x01'  # EAP-Request/Identity-ish
    _EAP_BYTES = bytes.fromhex(
        '00'      # next_payload = NONE
        '00'      # flags
        '0009'    # payload_length = 9 (4 header + 5 eap_data)
    ) + _EAP_DATA

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadEap.get_payload_type(), Ikev2PayloadType.EAP)

    def test_round_trip(self):
        eap_payload = Ikev2PayloadEap(
            flags=set(),
            eap_data=self._EAP_DATA,
        )
        eap_payload.next_payload = Ikev2PayloadType.NONE
        composed = eap_payload.compose()
        self.assertEqual(composed, self._EAP_BYTES)
        parsed: Ikev2PayloadEap = Ikev2PayloadEap.parse_exact_size(composed)
        self.assertEqual(parsed.eap_data, self._EAP_DATA)


class TestIkev2PayloadVendorId(unittest.TestCase):
    def setUp(self):
        self.vendor_id = (
            b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
            b'\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f'
        )

        self.vendor_dict = collections.OrderedDict([
            ('next_payload', b'\x00'),  # NONE
            ('flags', b'\x80'),  # CRITICAL (0x80)
            ('payload_length', b'\x00\x24'),  # 36 bytes (4 header + 32 vendor_id)
            ('vendor_id', self.vendor_id),
        ])
        self.vendor_bytes = b''.join(self.vendor_dict.values())

        self.vendor_payload = Ikev2PayloadVendorId(
            flags={Ikev2PayloadFlags.CRITICAL},
            vendor_id=self.vendor_id
        )
        self.vendor_payload.next_payload = Ikev2PayloadType.NONE

    def test_get_payload_type(self):
        self.assertEqual(Ikev2PayloadVendorId.get_payload_type(), Ikev2PayloadType.VENDOR_ID)

    def test_parse(self):
        parsed_vendor: Ikev2PayloadVendorId = Ikev2PayloadVendorId.parse_exact_size(self.vendor_bytes)
        self.assertEqual(parsed_vendor.next_payload, Ikev2PayloadType.NONE)
        self.assertEqual(parsed_vendor.flags, {Ikev2PayloadFlags.CRITICAL})
        self.assertEqual(parsed_vendor.vendor_id, self.vendor_id)

    def test_compose(self):
        composed_bytes = self.vendor_payload.compose()
        self.assertEqual(composed_bytes, self.vendor_bytes)

    def test_round_trip(self):
        composed_bytes = self.vendor_payload.compose()
        parsed_payload: Ikev2PayloadVendorId = Ikev2PayloadVendorId.parse_exact_size(composed_bytes)

        self.assertEqual(parsed_payload.vendor_id, self.vendor_payload.vendor_id)
        self.assertEqual(parsed_payload.flags, self.vendor_payload.flags)
        self.assertEqual(parsed_payload.next_payload, self.vendor_payload.next_payload)

    def test_error_parse_not_enough_data(self):
        incomplete_data = self.vendor_bytes[:-10]

        with self.assertRaises(NotEnoughData) as context_manager:
            Ikev2PayloadVendorId.parse_exact_size(incomplete_data)
        self.assertEqual(context_manager.exception.bytes_needed, 10)


class TestTransformAttributeKeyLength(unittest.TestCase):
    def setUp(self):
        self.key_length_value = 128
        self.key_length_attribute = TransformAttributeKeyLength(value=self.key_length_value)

    def test_parse(self):
        composed_bytes = self.key_length_attribute.compose()
        parsed_attribute = TransformAttributeKeyLength.parse_exact_size(composed_bytes)

        self.assertEqual(parsed_attribute.value, self.key_length_attribute.value)

    def test_compose(self):
        composed_bytes = self.key_length_attribute.compose()
        self.assertEqual(len(composed_bytes), 4)
        self.assertEqual(composed_bytes[0], 0x80)
        self.assertEqual(composed_bytes[1], Ikev2TransformAttributeType.KEY_LENGTH.value.code)
        self.assertEqual(composed_bytes[2:4], b'\x00\x80')

    def test_round_trip(self):
        composed_bytes = self.key_length_attribute.compose()
        parsed_attribute = TransformAttributeKeyLength.parse_exact_size(composed_bytes)

        self.assertEqual(parsed_attribute.value, self.key_length_attribute.value)

    def test_error_parse_not_enough_data(self):
        incomplete_data = b'\x80\x0e\x00'

        with self.assertRaises(NotEnoughData) as context_manager:
            TransformAttributeKeyLength.parse_exact_size(incomplete_data)
        self.assertEqual(context_manager.exception.bytes_needed, 1)


if __name__ == '__main__':
    unittest.main()
