# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.
from twisted.trial import unittest

from twisted.python import components
from twisted.pair import ip, raw
from zope import interface


@interface.implementer(raw.IRawDatagramProtocol)
class MyProtocol:
    def __init__(self, expecting):
        self.expecting = list(expecting)

    def datagramReceived(
        self,
        data,
        partial,
        source,
        dest,
        protocol,
        version,
        ihl,
        tos,
        tot_len,
        fragment_id,
        fragment_offset,
        dont_fragment,
        more_fragments,
        ttl,
    ):
        assert self.expecting, "Got a packet when not expecting anymore."
        expectData, expectKw = self.expecting.pop(0)

        expectKwKeys = expectKw.keys()
        expectKwKeys = list(sorted(expectKwKeys))
        localVariables = locals()

        for k in expectKwKeys:
            assert (
                expectKw[k] == localVariables[k]
            ), "Expected {}={!r}, got {!r}".format(k, expectKw[k], localVariables[k])
        assert expectData == data, f"Expected {expectData!r}, got {data!r}"

    def addProto(self, num, proto):
        # IRawDatagramProtocol.addProto
        pass


class IPTests(unittest.TestCase):
    def testPacketParsing(self):
        proto = ip.IPProtocol()
        p1 = MyProtocol(
            [
                (
                    b"foobar",
                    {
                        "partial": 0,
                        "dest": "1.2.3.4",
                        "source": "5.6.7.8",
                        "protocol": 0x0F,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
            ]
        )
        proto.addProto(0x0F, p1)

        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0F"  # protocol
            + b"FE"  # checksum
            + b"\x05\x06\x07\x08"
            + b"\x01\x02\x03\x04"
            + b"foobar",
            partial=0,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )

        assert not p1.expecting, (
            "Should not expect any more packets, but still want %r" % p1.expecting
        )

    def testMultiplePackets(self):
        proto = ip.IPProtocol()
        p1 = MyProtocol(
            [
                (
                    b"foobar",
                    {
                        "partial": 0,
                        "dest": "1.2.3.4",
                        "source": "5.6.7.8",
                        "protocol": 0x0F,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
                (
                    b"quux",
                    {
                        "partial": 1,
                        "dest": "5.4.3.2",
                        "source": "6.7.8.9",
                        "protocol": 0x0F,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
            ]
        )
        proto.addProto(0x0F, p1)
        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0F"  # protocol
            + b"FE"  # checksum
            + b"\x05\x06\x07\x08"
            + b"\x01\x02\x03\x04"
            + b"foobar",
            partial=0,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )
        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0F"  # protocol
            + b"FE"  # checksum
            + b"\x06\x07\x08\x09"
            + b"\x05\x04\x03\x02"
            + b"quux",
            partial=1,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )

        assert not p1.expecting, (
            "Should not expect any more packets, but still want %r" % p1.expecting
        )

    def testMultipleSameProtos(self):
        proto = ip.IPProtocol()
        p1 = MyProtocol(
            [
                (
                    b"foobar",
                    {
                        "partial": 0,
                        "dest": "1.2.3.4",
                        "source": "5.6.7.8",
                        "protocol": 0x0F,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
            ]
        )

        p2 = MyProtocol(
            [
                (
                    b"foobar",
                    {
                        "partial": 0,
                        "dest": "1.2.3.4",
                        "source": "5.6.7.8",
                        "protocol": 0x0F,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
            ]
        )

        proto.addProto(0x0F, p1)
        proto.addProto(0x0F, p2)

        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0F"  # protocol
            + b"FE"  # checksum
            + b"\x05\x06\x07\x08"
            + b"\x01\x02\x03\x04"
            + b"foobar",
            partial=0,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )

        assert not p1.expecting, (
            "Should not expect any more packets, but still want %r" % p1.expecting
        )
        assert not p2.expecting, (
            "Should not expect any more packets, but still want %r" % p2.expecting
        )

    def testWrongProtoNotSeen(self):
        proto = ip.IPProtocol()
        p1 = MyProtocol([])
        proto.addProto(1, p1)

        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0F"  # protocol
            + b"FE"  # checksum
            + b"\x05\x06\x07\x08"
            + b"\x01\x02\x03\x04"
            + b"foobar",
            partial=0,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )

    def testDemuxing(self):
        proto = ip.IPProtocol()
        p1 = MyProtocol(
            [
                (
                    b"foobar",
                    {
                        "partial": 0,
                        "dest": "1.2.3.4",
                        "source": "5.6.7.8",
                        "protocol": 0x0F,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
                (
                    b"quux",
                    {
                        "partial": 1,
                        "dest": "5.4.3.2",
                        "source": "6.7.8.9",
                        "protocol": 0x0F,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
            ]
        )
        proto.addProto(0x0F, p1)

        p2 = MyProtocol(
            [
                (
                    b"quux",
                    {
                        "partial": 1,
                        "dest": "5.4.3.2",
                        "source": "6.7.8.9",
                        "protocol": 0x0A,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
                (
                    b"foobar",
                    {
                        "partial": 0,
                        "dest": "1.2.3.4",
                        "source": "5.6.7.8",
                        "protocol": 0x0A,
                        "version": 4,
                        "ihl": 20,
                        "tos": 7,
                        "tot_len": 20 + 6,
                        "fragment_id": 0xDEAD,
                        "fragment_offset": 0x1EEF,
                        "dont_fragment": 0,
                        "more_fragments": 1,
                        "ttl": 0xC0,
                    },
                ),
            ]
        )
        proto.addProto(0x0A, p2)

        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0A"  # protocol
            + b"FE"  # checksum
            + b"\x06\x07\x08\x09"
            + b"\x05\x04\x03\x02"
            + b"quux",
            partial=1,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )
        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0F"  # protocol
            + b"FE"  # checksum
            + b"\x05\x06\x07\x08"
            + b"\x01\x02\x03\x04"
            + b"foobar",
            partial=0,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )
        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0F"  # protocol
            + b"FE"  # checksum
            + b"\x06\x07\x08\x09"
            + b"\x05\x04\x03\x02"
            + b"quux",
            partial=1,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )
        proto.datagramReceived(
            b"\x54"  # ihl version
            + b"\x07"  # tos
            + b"\x00\x1a"  # tot_len
            + b"\xDE\xAD"  # id
            + b"\xBE\xEF"  # frag_off
            + b"\xC0"  # ttl
            + b"\x0A"  # protocol
            + b"FE"  # checksum
            + b"\x05\x06\x07\x08"
            + b"\x01\x02\x03\x04"
            + b"foobar",
            partial=0,
            dest="dummy",
            source="dummy",
            protocol="dummy",
        )

        assert not p1.expecting, (
            "Should not expect any more packets, but still want %r" % p1.expecting
        )
        assert not p2.expecting, (
            "Should not expect any more packets, but still want %r" % p2.expecting
        )

    def testAddingBadProtos_WrongLevel(self):
        """Adding a wrong level protocol raises an exception."""
        e = ip.IPProtocol()
        try:
            e.addProto(42, "silliness")
        except components.CannotAdapt:
            pass
        else:
            raise AssertionError("addProto must raise an exception for bad protocols")

    def testAddingBadProtos_TooSmall(self):
        """Adding a protocol with a negative number raises an exception."""
        e = ip.IPProtocol()
        try:
            e.addProto(-1, MyProtocol([]))
        except TypeError as e:
            if e.args == ("Added protocol must be positive or zero",):
                pass
            else:
                raise
        else:
            raise AssertionError("addProto must raise an exception for bad protocols")

    def testAddingBadProtos_TooBig(self):
        """Adding a protocol with a number >=2**32 raises an exception."""
        e = ip.IPProtocol()
        try:
            e.addProto(2 ** 32, MyProtocol([]))
        except TypeError as e:
            if e.args == ("Added protocol must fit in 32 bits",):
                pass
            else:
                raise
        else:
            raise AssertionError("addProto must raise an exception for bad protocols")

    def testAddingBadProtos_TooBig2(self):
        """Adding a protocol with a number >=2**32 raises an exception."""
        e = ip.IPProtocol()
        try:
            e.addProto(2 ** 32 + 1, MyProtocol([]))
        except TypeError as e:
            if e.args == ("Added protocol must fit in 32 bits",):
                pass
            else:
                raise
        else:
            raise AssertionError("addProto must raise an exception for bad protocols")
