# Copyright 2014-2021 The aiosmtpd Developers
# SPDX-License-Identifier: Apache-2.0

import asyncio
import errno
import logging
import operator
import random
import socket
import struct
import time
from base64 import b64decode
from contextlib import contextmanager, suppress
from functools import partial
from ipaddress import IPv4Address, IPv6Address
from smtplib import SMTP as SMTPClient
from smtplib import SMTPServerDisconnected
from typing import Any, Callable, Dict, List, Optional

import pytest
from pytest_mock import MockFixture

from aiosmtpd.handlers import Sink
from aiosmtpd.proxy_protocol import (
    V2_CMD,
    AF,
    PROTO,
    V2_SIGNATURE,
    AsyncReader,
    MalformedTLV,
    ProxyData,
    ProxyTLV,
    UnknownTypeTLV,
    get_proxy,
)
from aiosmtpd.smtp import SMTP as SMTPServer
from aiosmtpd.smtp import Session as SMTPSession
from aiosmtpd.smtp import Envelope as SMTPEnvelope
from aiosmtpd.tests.conftest import Global, controller_data, handler_data

DEFAULT_AUTOCANCEL = 0.1
TIMEOUT_MULTIPLIER = 2.0

param = pytest.param
parametrize = pytest.mark.parametrize
random_port = partial(random.getrandbits, 16)

TEST_TLV_DATA_1 = (
    b"\x03\x00\x04Z\xfd\xc6\xff\x02\x00\tAUTHORITI\x05\x00\tUNIKUE_ID"
    b"\x20\x00D\x01\x00\x00\x00\x00!\x00\x07TLSv1.3%\x00\x07RSA4096"
    b"$\x00\nRSA-SHA256#\x00\x1bECDHE-RSA-AES256-CBC-SHA384"
)

TEST_TLV_DATA_2 = (
    b"\x03\x00\x04Z\xfd\xc6\xff\x02\x00\tAUTHORIT2\x05\x00\tUNIQUEID2"
    b"\x20\x00D\x01\x00\x00\x00\x00!\x00\x07TLSv1.3%\x00\x07RSA4096"
    b"$\x00\nRSA-SHA256#\x00\x1bECDHE-RSA-AES256-CBC-SHA384"
    b"\x30\x00\x09something"
)

# This has a tail which is not part of PROXYv2
TEST_V2_DATA1_XTRA = b64decode(
    "DQoNCgANClFVSVQKIREAcn8AAAF/AAABsFJipQMABFT9xv8CAAlBVVRIT1JJVFkFAAlVTklRVUVf\n"
    "SUQgAEQBAAAAACEAB1RMU3YxLjIlAAdSU0E0MDk2JAAKUlNBLVNIQTI1NiMAG0VDREhFLVJTQS1B\n"
    "RVMyNTYtR0NNLVNIQTM4NFRlc3QgZGF0YSB0aGF0IGlzIG5vdCBwYXJ0IG9mIFBST1hZdjIuCg==\n"
)

# The extra part is:
# b"Test data that is not part of PROXYv2.\n"

# This same as the above but no extraneous tail
TEST_V2_DATA1_EXACT = b64decode(
    "DQoNCgANClFVSVQKIREAcn8AAAF/AAABsFJipQMABFT9xv8CAAlBVVRIT1JJVFkFAAlVTklRVUVf\n"
    "SUQgAEQBAAAAACEAB1RMU3YxLjIlAAdSU0E0MDk2JAAKUlNBLVNIQTI1NiMAG0VDREhFLVJTQS1B\n"
    "RVMyNTYtR0NNLVNIQTM4NA==\n"
)

PUBLIC_V1_PATTERNS: Dict[str, bytes] = {
    "joaoreis81": b"PROXY TCP4 222.222.22.222 111.11.11.111 33504 25",
    "haproxydoc": b"PROXY TCP4 192.168.0.1 192.168.0.11 56324 443",
    "cloudflare4": b"PROXY TCP4 192.0.2.0 192.0.2.255 42300 443",
    "cloudflare6": (
        b"PROXY TCP6 2001:db8:: 2001:db8:ffff:ffff:ffff:ffff:ffff:ffff 42300 443"
    ),
    "avinetworks": b"PROXY TCP4 12.97.16.194 136.179.21.69 31646 80",
    "googlecloud": b"PROXY TCP4 192.0.2.1 198.51.100.1 15221 110",
}

GOOD_V1_HANDSHAKE = b"PROXY TCP4 255.255.255.255 255.255.255.255 65535 65535\r\n"

HANDSHAKES = {
    "v1": GOOD_V1_HANDSHAKE,
    "v2": TEST_V2_DATA1_EXACT,
}


class ProxyPeekerHandler(Sink):
    def __init__(self, retval: bool = True):
        self.called = False
        self.sessions: List[SMTPSession] = []
        self.proxy_datas: List[ProxyData] = []
        self.retval = retval

    async def handle_PROXY(
        self,
        server: SMTPServer,
        session: SMTPSession,
        envelope: SMTPEnvelope,
        proxy_data: ProxyData,
    ) -> bool:
        self.called = True
        self.sessions.append(session)
        self.proxy_datas.append(proxy_data)
        return self.retval


@contextmanager
def does_not_raise():
    yield


@pytest.fixture
def setup_proxy_protocol(
    mocker: MockFixture, temp_event_loop: asyncio.AbstractEventLoop
) -> Callable:
    proxy_timeout = 1.0
    responses = []
    transport = mocker.Mock()
    transport.write = responses.append
    handler = ProxyPeekerHandler()
    loop = temp_event_loop

    def getter(test_obj: "_TestProxyProtocolCommon", *args, **kwargs):
        kwargs["loop"] = loop
        kwargs["proxy_protocol_timeout"] = proxy_timeout
        protocol = SMTPServer(handler, *args, **kwargs)
        protocol.connection_made(transport)

        def runner(stop_after: float = DEFAULT_AUTOCANCEL):
            loop.call_later(stop_after, protocol._handler_coroutine.cancel)
            with suppress(asyncio.CancelledError):
                loop.run_until_complete(protocol._handler_coroutine)

        test_obj.protocol = protocol
        test_obj.runner = runner
        test_obj.transport = transport

    return getter


class _TestProxyProtocolCommon:
    protocol: SMTPServer = None
    runner = None
    transport = None


class TestProxyData:
    def test_invalid_version(self):
        pd = ProxyData(version=None)
        assert not pd.valid
        assert not pd

    def test_invalid_error(self):
        pd = ProxyData(version=1)
        pd.error = "SomeError"
        assert not pd.valid
        assert not pd

    def test_invalid_protocol(self):
        pd = ProxyData(version=1)
        assert not pd.valid
        assert not pd

    def test_mismatch(self):
        pd = ProxyData(version=1)
        pd.protocol = PROTO.UNSPEC
        assert pd.valid
        assert pd
        assert not pd.same_attribs(protocol=PROTO.STREAM)

    def test_mismatch_raises(self):
        pd = ProxyData(version=1)
        pd.protocol = PROTO.UNSPEC
        assert pd.valid
        assert pd
        expectre = r"mismatch:protocol actual=.* expect=.*"
        with pytest.raises(ValueError, match=expectre):
            pd.same_attribs(_raises=True, protocol=PROTO.STREAM)

    def test_unsetkey(self):
        pd = ProxyData(version=1)
        pd.protocol = PROTO.UNSPEC
        assert pd.valid
        assert pd
        assert not pd.same_attribs(src_addr="Missing")

    def test_unknownkey(self):
        pd = ProxyData(version=1)
        pd.protocol = PROTO.UNSPEC
        assert pd.valid
        assert pd
        assert not pd.same_attribs(strange_key="Unrecognized")

    def test_unknownkey_raises(self):
        pd = ProxyData(version=1)
        pd.protocol = PROTO.UNSPEC
        assert pd.valid
        assert pd
        expectre = r"notfound:strange_key"
        with pytest.raises(KeyError, match=expectre):
            pd.same_attribs(_raises=True, strange_key="Unrecognized")

    def test_tlv_none(self):
        pd = ProxyData(version=2)
        assert pd.tlv is None

    def test_tlv_fake(self):
        pd = ProxyData(version=2)
        pd.rest = b"fake_tlv"  # Must be something that fails parsing in ProxyTLV
        assert pd.tlv is None

    def test_tlv_1(self):
        pd = ProxyData(version=2)
        pd.rest = TEST_TLV_DATA_1
        assert pd.tlv is not None
        assert pd.tlv.same_attribs(
            AUTHORITY=b"AUTHORITI",
            CRC32C=b"Z\xfd\xc6\xff",
            UNIQUE_ID=b"UNIKUE_ID",
            SSL=True,
            SSL_VERSION=b"TLSv1.3",
            SSL_CIPHER=b"ECDHE-RSA-AES256-CBC-SHA384",
            SSL_SIG_ALG=b"RSA-SHA256",
            SSL_KEY_ALG=b"RSA4096",
        )


class TestProxyTLV:
    def test_1(self):
        ptlv = ProxyTLV.from_raw(TEST_TLV_DATA_1)
        assert "ALPN" not in ptlv
        assert "NOOP" not in ptlv
        assert "SSL_CN" not in ptlv
        assert "NETNS" not in ptlv
        assert ptlv.ALPN is None
        assert ptlv.AUTHORITY == b"AUTHORITI"
        assert ptlv.same_attribs(
            AUTHORITY=b"AUTHORITI",
            CRC32C=b"Z\xfd\xc6\xff",
            UNIQUE_ID=b"UNIKUE_ID",
            SSL=True,
            SSL_VERSION=b"TLSv1.3",
            SSL_CIPHER=b"ECDHE-RSA-AES256-CBC-SHA384",
            SSL_SIG_ALG=b"RSA-SHA256",
            SSL_KEY_ALG=b"RSA4096",
        )

    def test_1_ne(self):
        ptlv = ProxyTLV.from_raw(TEST_TLV_DATA_1)
        assert not ptlv.same_attribs(SSL=False)
        assert not ptlv.same_attribs(false_attrib=None)

    def test_1_ne_raises(self):
        ptlv = ProxyTLV.from_raw(TEST_TLV_DATA_1)
        expectre = r"mismatch:AUTHORITY actual=.* expect=.*"
        with pytest.raises(ValueError, match=expectre):
            ptlv.same_attribs(_raises=True, AUTHORITY=b"whut")
        expectre = r"notfound:i_dont_even"
        with pytest.raises(KeyError, match=expectre):
            ptlv.same_attribs(_raises=True, i_dont_even=b"what is this")

    def test_2(self):
        ptlv = ProxyTLV.from_raw(TEST_TLV_DATA_2)
        assert "ALPN" not in ptlv
        assert "NOOP" not in ptlv
        assert "SSL_CN" not in ptlv
        assert "NETNS" in ptlv
        assert ptlv.ALPN is None
        assert ptlv.AUTHORITY == b"AUTHORIT2"
        assert ptlv.UNIQUE_ID == b"UNIQUEID2"
        assert ptlv.same_attribs(
            CRC32C=b"Z\xfd\xc6\xff",
            UNIQUE_ID=b"UNIQUEID2",
            SSL=True,
            SSL_VERSION=b"TLSv1.3",
            SSL_CIPHER=b"ECDHE-RSA-AES256-CBC-SHA384",
            SSL_SIG_ALG=b"RSA-SHA256",
            SSL_KEY_ALG=b"RSA4096",
            NETNS=b"something",
        )
        assert not ptlv.same_attribs(false_attrib=None)
        with pytest.raises(ValueError, match=r"mismatch:SSL"):
            ptlv.same_attribs(SSL=False, _raises=True)

    @parametrize(
        "typeint, typename",
        [
            (0x01, "ALPN"),
            (0x02, "AUTHORITY"),
            (0x03, "CRC32C"),
            (0x04, "NOOP"),
            (0x05, "UNIQUE_ID"),
            (0x20, "SSL"),
            (0x21, "SSL_VERSION"),
            (0x22, "SSL_CN"),
            (0x23, "SSL_CIPHER"),
            (0x24, "SSL_SIG_ALG"),
            (0x25, "SSL_KEY_ALG"),
            (0x30, "NETNS"),
            (None, "wrongname"),
        ],
    )
    def test_backmap(self, typename: str, typeint: int):
        assert ProxyTLV.name_to_num(typename) == typeint

    def test_parse_partial(self):
        TEST_TLV_DATA_1_TRUNCATED = (
            b"\x03\x00\x04Z\xfd\xc6\xff\x02\x00\tAUTHORITI\x05\x00\tUNIKUE_I"
        )
        rslt: Dict[str, Any]
        rslt, _ = ProxyTLV.parse(TEST_TLV_DATA_1_TRUNCATED)
        assert isinstance(rslt, dict)
        assert "CRC32C" in rslt

    def test_unknowntype_notstrict(self):
        test_data = b"\xFF\x00\x04yeah"
        rslt: Dict[str, Any]
        rslt, _ = ProxyTLV.parse(test_data, strict=False)
        assert "xFF" in rslt
        assert rslt["xFF"] == b"yeah"

    def test_unknowntype_strict(self):
        test_data = b"\xFF\x00\x04yeah"
        with pytest.raises(UnknownTypeTLV):
            _ = ProxyTLV.parse(test_data, strict=True)

    def test_malformed_ssl_partialok(self):
        test_data = (
            b"\x20\x00\x17\x01\x02\x03\x04\x05\x21\x00\x07version\x22\x00\x09trunc"
        )
        rslt: Dict[str, Any]
        rslt, _ = ProxyTLV.parse(test_data, partial_ok=True)
        assert rslt["SSL"] is False
        assert rslt["SSL_VERSION"] == b"version"
        assert "SSL_CN" not in rslt

    def test_malformed_ssl_notpartialok(self):
        test_data = (
            b"\x20\x00\x0D\x01\x02\x03\x04\x05\x21\x00\x07version\x22\x00\x09trunc"
        )
        with pytest.raises(MalformedTLV):
            _ = ProxyTLV.parse(test_data, partial_ok=False)

    def test_eq(self):
        data1 = b"\x03\x00\x04Z\xfd\xc6\xff\x02\x00\tAUTHORITI"
        ptlv1 = ProxyTLV.from_raw(data1)
        data2 = b"\x02\x00\tAUTHORITI\x03\x00\x04Z\xfd\xc6\xff"
        ptlv2 = ProxyTLV.from_raw(data2)
        assert ptlv1 is not ptlv2
        assert ptlv1 == ptlv2


class TestModule:
    class MockAsyncReader(AsyncReader):
        def __init__(self, data, timeout=0.4):
            self.data = bytearray(data)
            self.timeout = 0.4

        async def read(self, num_bytes: Optional[int] = None) -> bytes:
            emit = self.data[0:num_bytes]
            del self.data[0:num_bytes]
            return emit

        async def readexactly(self, n: int) -> bytes:
            emit = self.data[0:n]
            del self.data[0:n]
            if len(emit) < n:
                await asyncio.sleep(self.timeout)
            return emit

        async def readuntil(self, until_chars: Optional[bytes] = None) -> bytes:
            if until_chars is None:
                until_chars = b"\n"
            emit = bytearray()
            _count = 0
            for _count, char in enumerate(self.data, start=1):
                emit += char.to_bytes(1, "big")
                if char in until_chars:
                    break
            del self.data[0:_count]
            return emit

    @parametrize("handshake", HANDSHAKES.values(), ids=HANDSHAKES.keys())
    def test_get(
        self,
        caplog: pytest.LogCaptureFixture,
        temp_event_loop: asyncio.AbstractEventLoop,
        handshake: bytes,
    ):
        caplog.set_level(logging.DEBUG)
        mock_reader = self.MockAsyncReader(handshake)
        reslt = temp_event_loop.run_until_complete(get_proxy(mock_reader))
        assert isinstance(reslt, ProxyData)
        assert reslt.valid

    def test_get_cut_v1(
        self,
        caplog: pytest.LogCaptureFixture,
        temp_event_loop: asyncio.AbstractEventLoop,
    ):
        caplog.set_level(logging.DEBUG)
        mock_reader = self.MockAsyncReader(GOOD_V1_HANDSHAKE[0:20])
        reslt = temp_event_loop.run_until_complete(get_proxy(mock_reader))
        assert isinstance(reslt, ProxyData)
        assert not reslt.valid
        assert reslt.error == "PROXYv1 malformed"
        expect = ("mail.debug", 30, "PROXY error: PROXYv1 malformed")
        assert expect in caplog.record_tuples

    def test_get_cut_v2(
        self,
        caplog: pytest.LogCaptureFixture,
        temp_event_loop: asyncio.AbstractEventLoop,
    ):
        caplog.set_level(logging.DEBUG)
        mock_reader = self.MockAsyncReader(TEST_V2_DATA1_EXACT[0:20])
        reslt = temp_event_loop.run_until_complete(get_proxy(mock_reader))
        assert isinstance(reslt, ProxyData)
        assert not reslt.valid
        expect_msg = "PROXY exception: Connection lost while waiting for tail part"
        assert reslt.error == expect_msg
        expect = ("mail.debug", 30, expect_msg)
        assert expect in caplog.record_tuples

    def test_get_invalid_sig(
        self,
        caplog: pytest.LogCaptureFixture,
        temp_event_loop: asyncio.AbstractEventLoop,
    ):
        caplog.set_level(logging.DEBUG)
        mock_reader = self.MockAsyncReader(b"PROXI TCP4 1.2.3.4 5.6.7.8 9 10\r\n")
        reslt = temp_event_loop.run_until_complete(get_proxy(mock_reader))
        assert isinstance(reslt, ProxyData)
        assert not reslt.valid
        expect_msg = "PROXY unrecognized signature"
        assert reslt.error == expect_msg
        expect = ("mail.debug", 30, "PROXY error: " + expect_msg)
        assert expect in caplog.record_tuples


class TestSMTPInit:
    @parametrize("value", [int(-1), float(-1.0), int(0), float(0.0)])
    def test_value_error(self, temp_event_loop, value):
        with pytest.raises(ValueError, match=r"proxy_protocol_timeout must be > 0"):
            _ = SMTPServer(Sink(), proxy_protocol_timeout=value, loop=temp_event_loop)

    def test_lt_3(self, caplog, temp_event_loop):
        _ = SMTPServer(Sink(), proxy_protocol_timeout=1, loop=temp_event_loop)
        expect = ("mail.log", logging.WARNING, "proxy_protocol_timeout < 3.0")
        assert expect in caplog.record_tuples

    @parametrize("value", [int(3), float(3.0), int(4), float(4.0)])
    def test_ge_3(self, caplog, temp_event_loop, value):
        _ = SMTPServer(Sink(), proxy_protocol_timeout=value, loop=temp_event_loop)
        expect = ("mail.log", logging.WARNING, "proxy_protocol_timeout < 3.0")
        assert expect not in caplog.record_tuples


class TestGetV1(_TestProxyProtocolCommon):
    def test_noproxy(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        data = b"HELO example.org\r\n"
        self.protocol.data_received(data)
        self.runner()
        assert self.transport.close.called

    @parametrize("patt", PUBLIC_V1_PATTERNS.values(), ids=PUBLIC_V1_PATTERNS.keys())
    def test_valid_patterns(self, setup_proxy_protocol: Callable, patt: bytes):
        if not patt.endswith(b"\r\n"):
            patt += b"\r\n"
        setup_proxy_protocol(self)
        self.protocol.data_received(patt)
        self.runner()
        sess: SMTPSession = self.protocol.session
        assert sess.proxy_data.error == ""
        handler = self.protocol.event_handler
        assert handler.called

    def _assert_valid(self, af, proto, srcip, dstip, srcport, dstport, testline):
        self.protocol.data_received(testline.encode("ascii"))
        self.runner()
        assert self.protocol.session.proxy_data.error == ""
        handler = self.protocol.event_handler
        assert handler.called
        proxy_data = handler.proxy_datas[-1]
        ipaddr = IPv4Address if af == AF.INET else IPv6Address
        assert proxy_data.same_attribs(
            valid=True,
            version=1,
            family=af,
            protocol=proto,
            src_addr=ipaddr(srcip),
            dst_addr=ipaddr(dstip),
            src_port=srcport,
            dst_port=dstport,
        )

    def test_tcp4(self, setup_proxy_protocol):
        srcip = "1.2.3.4"
        dstip = "5.6.7.8"
        srcport = 0
        dstport = 65535
        prox_test = f"PROXY TCP4 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_valid(
            AF.INET, PROTO.STREAM, srcip, dstip, srcport, dstport, prox_test
        )

    def test_tcp4_random(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        srcip = ".".join(f"{random.getrandbits(8)}" for _ in range(0, 4))
        dstip = ".".join(f"{random.getrandbits(8)}" for _ in range(0, 4))
        srcport = random_port()
        dstport = random_port()
        prox_test = f"PROXY TCP4 {srcip} {dstip} {srcport} {dstport}\r\n"
        self._assert_valid(
            AF.INET, PROTO.STREAM, srcip, dstip, srcport, dstport, prox_test
        )

    def test_tcp6_shortened(self, setup_proxy_protocol):
        srcip = "2020:dead::0001"
        dstip = "2021:cafe::0002"
        srcport = 8000
        dstport = 65535
        prox_test = f"PROXY TCP6 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_valid(
            AF.INET6, PROTO.STREAM, srcip, dstip, srcport, dstport, prox_test
        )

    def test_tcp6_random(self, setup_proxy_protocol):
        srcip = ":".join(f"{random.getrandbits(16):04x}" for _ in range(0, 8))
        dstip = ":".join(f"{random.getrandbits(16):04x}" for _ in range(0, 8))
        srcport = random_port()
        dstport = random_port()
        prox_test = f"PROXY TCP6 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_valid(
            AF.INET6, PROTO.STREAM, srcip, dstip, srcport, dstport, prox_test
        )

    def test_unknown(self, setup_proxy_protocol):
        prox_test = "PROXY UNKNOWN whatever\r\n"
        setup_proxy_protocol(self)
        self.protocol.data_received(prox_test.encode("ascii"))
        self.runner()
        handler = self.protocol.event_handler
        assert handler.called
        proxy_data = handler.proxy_datas[0]
        assert proxy_data.same_attribs(
            valid=True,
            version=1,
            family=AF.UNSPEC,
            protocol=PROTO.UNSPEC,
            rest=b" whatever",
        )

    def test_unknown_short(self, setup_proxy_protocol):
        prox_test = "PROXY UNKNOWN\r\n"
        setup_proxy_protocol(self)
        self.protocol.data_received(prox_test.encode("ascii"))
        self.runner()
        handler = self.protocol.event_handler
        assert handler.called
        proxy_data = handler.proxy_datas[0]
        assert proxy_data.same_attribs(
            valid=True,
            version=1,
            family=AF.UNSPEC,
            protocol=PROTO.UNSPEC,
            rest=b"",
        )

    def _assert_invalid(self, testline: str, expect_err: str = ""):
        self.protocol.data_received(testline.encode("ascii"))
        self.runner()
        handler: ProxyPeekerHandler = self.protocol.event_handler
        assert not self.protocol.session.proxy_data.valid
        assert not handler.called
        assert self.transport.close.called
        assert self.protocol.session.proxy_data.error == expect_err

    def test_invalid_sig(self, setup_proxy_protocol):
        prox_test = "PROXY1 UNKNOWN whatevs\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 wrong signature")

    def test_unsupported_family(self, setup_proxy_protocol):
        prox_test = "PROXY TCP5 123.123.123.123 231.231.231.231 80 90\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 unrecognized family")
        prox_test = "PROXY TCP 123.123.123.123 231.231.231.231 80 90\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 unrecognized family")

    def test_unsupported_proto(self, setup_proxy_protocol):
        prox_test = "PROXY UDP4 123.123.123.123 231.231.231.231 80 90\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 unrecognized protocol")

    def test_too_long(self, setup_proxy_protocol):
        prox_test = "PROXY UNKNOWN " + "*" * 100 + "\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 header too long")

    def test_malformed_nocr(self, setup_proxy_protocol):
        prox_test = "PROXY UNKNOWN\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 malformed")

    def test_malformed_notproxy(self, setup_proxy_protocol):
        srcip = "1.2.3.4"
        dstip = "5.6.7.8"
        srcport = 65535
        dstport = 65535
        prox_test = f"NOTPROX TCP4 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXY unrecognized signature")

    def test_malformed_wrongtype_64(self, setup_proxy_protocol):
        srcip = "1.2.3.4"
        dstip = "5.6.7.8"
        srcport = 65535
        dstport = 65535
        prox_test = f"PROXY TCP6 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 address not IPv6")

    def test_malformed_wrongtype_46(self, setup_proxy_protocol):
        srcip = "2020:dead::0001"
        dstip = "2021:cafe::0002"
        srcport = 65535
        dstport = 65535
        prox_test = f"PROXY TCP4 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 address not IPv4")

    def test_malformed_wrongtype_6mixed(self, setup_proxy_protocol):
        srcip = "1.2.3.4"
        dstip = "2021:cafe::0002"
        srcport = 65535
        dstport = 65535
        prox_test = f"PROXY TCP6 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 address not IPv6")

    IP6_dead = "2020:dead::0001"
    IP6_cafe = "2021:cafe::0002"

    @parametrize(
        "srcip, dstip, srcport, dstport, whatwrong",
        [
            param(IP6_dead, IP6_cafe, "02501", None, "port", id="zeroleader"),
            param(" " + IP6_dead, IP6_cafe, None, None, "address", id="space1"),
            param(IP6_dead, " " + IP6_cafe, None, None, "address", id="space2"),
            param(IP6_dead, IP6_cafe, " 8080", None, "port", id="space3"),
            param(IP6_dead, IP6_cafe, None, " 0", "port", id="space4"),
            param(IP6_dead[:-1] + "g", IP6_cafe, None, None, "address", id="addr6s"),
            param(IP6_dead, IP6_cafe[:-1] + "h", None, None, "address", id="addr6d"),
        ],
    )
    def test_malformed_addr(
        self, setup_proxy_protocol, srcip, dstip, srcport, dstport, whatwrong
    ):
        if srcport is None:
            srcport = random_port()
        if dstport is None:
            dstport = random_port()
        prox_test = f"PROXY TCP6 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, f"PROXYv1 {whatwrong} malformed")

    @parametrize(
        "extra",
        [
            param(" ", id="space"),
            param(" text", id="sptext"),
        ],
    )
    def test_extra(self, setup_proxy_protocol, extra):
        prox_test = f"PROXY TCP6 {self.IP6_dead} {self.IP6_cafe} 0 25{extra}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 unrecognized extraneous data")

    def test_malformed_addr4(self, setup_proxy_protocol):
        srcip = "1.2.3.a"
        dstip = "5.6.7.8"
        srcport = 65535
        dstport = 65535
        prox_test = f"PROXY TCP6 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 address malformed")

    def test_ports_oob(self, setup_proxy_protocol):
        srcip = "1.2.3.4"
        dstip = "5.6.7.8"
        srcport = 65536
        dstport = 10200
        prox_test = f"PROXY TCP4 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 src port out of bounds")

    def test_portd_oob(self, setup_proxy_protocol):
        srcip = "2020:dead::0001"
        dstip = "2021:cafe::0002"
        srcport = 10000
        dstport = 65536
        prox_test = f"PROXY TCP6 {srcip} {dstip} {srcport} {dstport}\r\n"
        setup_proxy_protocol(self)
        self._assert_invalid(prox_test, "PROXYv1 dst port out of bounds")


class TestGetV2(_TestProxyProtocolCommon):
    def test_1(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        self.protocol.data_received(TEST_V2_DATA1_XTRA)
        self.runner()
        sess: SMTPSession = self.protocol.session
        assert sess.proxy_data.error == ""
        handler: ProxyPeekerHandler = self.protocol.event_handler
        assert handler.called
        pd: ProxyData = handler.proxy_datas[-1]
        assert pd.same_attribs(
            version=2,
            command=1,
            family=1,
            src_addr=IPv4Address("127.0.0.1"),
            dst_addr=IPv4Address("127.0.0.1"),
            src_port=45138,
            dst_port=25253,
        )
        assert b"not part of PROXYv2" not in pd.rest
        assert pd.tlv
        assert pd.tlv.same_attribs(
            AUTHORITY=b"AUTHORITY",
            CRC32C=b"T\xfd\xc6\xff",
            UNIQUE_ID=b"UNIQUE_ID",
            SSL=True,
            SSL_VERSION=b"TLSv1.2",
            SSL_CIPHER=b"ECDHE-RSA-AES256-GCM-SHA384",
            SSL_SIG_ALG=b"RSA-SHA256",
            SSL_KEY_ALG=b"RSA4096",
        )

    def _send_valid(
        self, cmd: V2_CMD, fam: AF, proto: PROTO, payload: bytes
    ) -> ProxyData:
        ver_cmd = 0x20 + cmd.value
        fam_pro = (fam << 4) + proto
        to_send = bytes(
            V2_SIGNATURE
            + ver_cmd.to_bytes(1, "big")
            + fam_pro.to_bytes(1, "big")
            + len(payload).to_bytes(2, "big")
            + payload
        )
        self.protocol.data_received(to_send)
        self.runner()
        assert self.protocol.session.proxy_data.error == ""
        assert self.protocol.session.proxy_data.whole_raw == to_send
        handler = self.protocol.event_handler
        assert handler.called
        return handler.proxy_datas[-1]

    def test_UNSPEC_empty(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        assert self._send_valid(
            V2_CMD.LOCAL, AF.UNSPEC, PROTO.UNSPEC, b""
        ).same_attribs(valid=True, version=2, command=0, family=0, protocol=0, rest=b"")

    def test_UNSPEC_notempty(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        payload = b"asdfghjkl"
        assert self._send_valid(
            V2_CMD.LOCAL, AF.UNSPEC, PROTO.UNSPEC, payload
        ).same_attribs(
            valid=True, version=2, command=0, family=0, protocol=0, rest=payload
        )

    @parametrize("ttlv", [b"", b"fake_tlv"])
    @parametrize("tproto", [PROTO.STREAM, PROTO.DGRAM])
    def test_INET4(self, setup_proxy_protocol, tproto, ttlv):
        setup_proxy_protocol(self)
        src_addr = IPv4Address("10.212.4.33")
        dst_addr = IPv4Address("10.11.12.13")
        src_port = 0
        dst_port = 65535
        payload = (
            src_addr.packed
            + dst_addr.packed
            + src_port.to_bytes(2, "big")
            + dst_port.to_bytes(2, "big")
            + ttlv
        )
        pd = self._send_valid(V2_CMD.LOCAL, AF.INET, tproto, payload)
        assert pd.same_attribs(
            valid=True,
            version=2,
            command=0,
            family=AF.INET,
            protocol=tproto,
            src_addr=src_addr,
            dst_addr=dst_addr,
            src_port=src_port,
            dst_port=dst_port,
            rest=ttlv,
        )
        assert pd.tlv is None

    @parametrize("ttlv", [b"", b"fake_tlv"])
    @parametrize("tproto", [PROTO.STREAM, PROTO.DGRAM])
    def test_INET6(self, setup_proxy_protocol, tproto, ttlv):
        setup_proxy_protocol(self)
        src_addr = IPv6Address("2020:dead::0001")
        dst_addr = IPv6Address("2021:cafe::0022")
        src_port = 65534
        dst_port = 8080
        payload = (
            src_addr.packed
            + dst_addr.packed
            + src_port.to_bytes(2, "big")
            + dst_port.to_bytes(2, "big")
            + ttlv
        )
        pd = self._send_valid(V2_CMD.LOCAL, AF.INET6, tproto, payload)
        assert pd.same_attribs(
            valid=True,
            version=2,
            command=0,
            family=2,
            protocol=tproto,
            src_addr=src_addr,
            dst_addr=dst_addr,
            src_port=src_port,
            dst_port=dst_port,
            rest=ttlv,
        )
        assert pd.tlv is None

    @parametrize("ttlv", [b"", b"fake_tlv"])
    @parametrize("tproto", [PROTO.STREAM, PROTO.DGRAM])
    def test_UNIX(self, setup_proxy_protocol, tproto, ttlv):
        setup_proxy_protocol(self)
        src_addr = struct.pack("108s", b"/proc/source")
        dst_addr = struct.pack("108s", b"/proc/dest")
        payload = src_addr + dst_addr + ttlv
        pd = self._send_valid(V2_CMD.LOCAL, AF.UNIX, tproto, payload)
        assert pd.same_attribs(
            valid=True,
            version=2,
            command=0,
            family=3,
            protocol=tproto,
            src_addr=src_addr,
            dst_addr=dst_addr,
            src_port=None,
            dst_port=None,
            rest=ttlv,
        )
        assert pd.tlv is None

    @parametrize(
        "tfam, tproto",
        [
            (AF.UNSPEC, PROTO.STREAM),
            (AF.UNSPEC, PROTO.DGRAM),
            (AF.INET, PROTO.UNSPEC),
            (AF.INET6, PROTO.UNSPEC),
            (AF.UNIX, PROTO.UNSPEC),
        ],
    )
    def test_fallback_UNSPEC(self, setup_proxy_protocol, tfam, tproto):
        setup_proxy_protocol(self)
        payload = b"whatever"
        assert self._send_valid(V2_CMD.LOCAL, tfam, tproto, payload).same_attribs(
            valid=True,
            version=2,
            command=0,
            family=tfam,
            protocol=tproto,
            src_addr=None,
            dst_addr=None,
            src_port=None,
            dst_port=None,
            rest=payload,
        )

    def _send_invalid(
        self,
        sig: bytes = V2_SIGNATURE,
        ver: int = 2,
        cmd: int = 0,
        fam: int = 0,
        proto: int = 0,
        payload: bytes = b"",
        expect: Optional[str] = None,
    ) -> ProxyData:
        ver_cmd = (ver << 4) | cmd
        fam_pro = (fam << 4) | proto
        self.protocol.data_received(
            sig
            + ver_cmd.to_bytes(1, "big")
            + fam_pro.to_bytes(1, "big")
            + len(payload).to_bytes(2, "big")
            + payload
        )
        self.runner()
        handler = self.protocol.event_handler
        assert not handler.called
        assert not self.protocol.session.proxy_data
        if expect is not None:
            assert self.protocol.session.proxy_data.error == expect
        return self.protocol.session.proxy_data

    def test_invalid_sig(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        ERRSIG = b"\r\n\r\n\x00\r\nQUIP\n"
        self._send_invalid(sig=ERRSIG, expect="PROXYv2 wrong signature")

    # Using readexactly() instead of just read() causes incomplete handshake to time-
    # out as readexactly() waits until the number of bytes is complete.
    # So, we can no longer look for this error.
    #
    # def test_incomplete(self, setup_proxy_protocol):
    #     setup_proxy_protocol(self)
    #     self.protocol.data_received(V2_SIGNATURE + b"\x20" + b"\x00" + b"\x00")
    #     self.runner()
    #     assert self.transport.close.called
    #     handler = self.protocol.event_handler
    #     assert not handler.called
    #     sess: SMTPSession = self.protocol.session
    #     assert not sess.proxy_data.valid
    #     assert sess.proxy_data.error == "PROXYv2 malformed header"

    def test_illegal_ver(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        self._send_invalid(ver=3, expect="PROXYv2 illegal version")

    def test_unsupported_cmd(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        self._send_invalid(cmd=2, expect="PROXYv2 unsupported command")

    def test_unsupported_fam(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        self._send_invalid(fam=4, expect="PROXYv2 unsupported family")

    def test_unsupported_proto(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        self._send_invalid(proto=3, expect="PROXYv2 unsupported protocol")

    def test_wrong_proto_6shouldbe4(self, setup_proxy_protocol):
        setup_proxy_protocol(self)
        src_addr = IPv4Address("192.168.0.11")
        dst_addr = IPv4Address("172.16.0.22")
        src_port = 65534
        dst_port = 8080
        payload = (
            src_addr.packed
            + dst_addr.packed
            + src_port.to_bytes(2, "big")
            + dst_port.to_bytes(2, "big")
        )
        self._send_invalid(
            fam=2, proto=1, payload=payload, expect="PROXYv2 truncated address"
        )


@controller_data(proxy_protocol_timeout=0.3)
@handler_data(class_=ProxyPeekerHandler)
class TestWithController:
    def _okay(self, handshake: bytes):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.connect(Global.SrvAddr)
            sock.sendall(handshake)
            resp = sock.makefile("rb").readline()
            assert resp.startswith(b"220 ")
            with SMTPClient() as client:
                client.sock = sock
                code, mesg = client.ehlo("example.org")
                assert code == 250
                code, mesg = client.quit()
                assert code == 221

    @parametrize("handshake", HANDSHAKES.values(), ids=HANDSHAKES.keys())
    def test_okay(self, plain_controller, handshake):
        assert plain_controller.smtpd._proxy_timeout > 0.0
        self._okay(handshake)

    @parametrize("handshake", HANDSHAKES.values(), ids=HANDSHAKES.keys())
    def test_hiccup(self, plain_controller, handshake):
        assert plain_controller.smtpd._proxy_timeout > 0.0
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.connect(Global.SrvAddr)
            sock.sendall(handshake[0:20])
            time.sleep(0.1)
            sock.sendall(handshake[20:])
            resp = sock.makefile("rb").readline()
            assert resp.startswith(b"220 ")
            with SMTPClient() as client:
                client.sock = sock
                code, mesg = client.ehlo("example.org")
                assert code == 250
                code, mesg = client.quit()
                assert code == 221

    @parametrize("handshake", HANDSHAKES.values(), ids=HANDSHAKES.keys())
    def test_timeout(self, plain_controller, handshake):
        assert plain_controller.smtpd._proxy_timeout > 0.0
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.connect(Global.SrvAddr)
            time.sleep(plain_controller.smtpd._proxy_timeout * TIMEOUT_MULTIPLIER)
            # noinspection PyTypeChecker
            with pytest.raises(ConnectionError):
                sock.send(handshake)
                resp = sock.recv(4096)
                if resp == b"":
                    raise ConnectionError
            # Try resending the handshake. Should also fail (because connection has
            # been closed by the server.
            # noinspection PyTypeChecker
            with pytest.raises(OSError) as exc_info:  # noqa: PT011
                sock.send(handshake)
                resp = sock.recv(4096)
                if resp == b"":
                    raise ConnectionAbortedError
            # MacOS sometimes raises EPROTOTYPE, which won't result in ConnectionError
            # but in OSError(errno=EPROTOTYPE). Let's check that here.
            # Refs:
            #   - https://github.com/racitup/static-ranges/issues/1
            #   - https://github.com/benoitc/gunicorn/issues/1487
            #   - http://erickt.github.io/blog/2014/11/19/adventures-in-debugging-\
            #     a-potential-osx-kernel-bug/
            exc = exc_info.value
            if isinstance(exc, ConnectionError):
                pass
            else:
                assert exc.errno in (errno.EPROTOTYPE, errno.EPIPE)
        # Assert that we can connect properly afterwards (that is, server is not
        # terminated)
        self._okay(handshake)

    @parametrize("handshake", HANDSHAKES.values(), ids=HANDSHAKES.keys())
    def test_incomplete(self, plain_controller, handshake):
        assert plain_controller.smtpd._proxy_timeout > 0.0
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.connect(Global.SrvAddr)
            sock.send(handshake[:-1])
            time.sleep(plain_controller.smtpd._proxy_timeout * TIMEOUT_MULTIPLIER)
            # noinspection PyTypeChecker
            with pytest.raises(ConnectionError):
                sock.send(b"\n")
                resp = sock.recv(4096)  # On Windows, this line raises
                if resp == b"":  # On Linux, no raise, just "EOF"
                    raise ConnectionError
            # Try resending the handshake. Should also fail (because connection has
            # been closed by the server.
            # noinspection PyTypeChecker
            with pytest.raises(OSError) as exc_info:  # noqa: PT011
                sock.send(handshake)
                resp = sock.recv(4096)
                if resp == b"":
                    raise ConnectionError
            # MacOS sometimes raises EPROTOTYPE, which won't result in ConnectionError
            # but in OSError(errno=EPROTOTYPE). Let's check that here.
            # Refs:
            #   - https://github.com/racitup/static-ranges/issues/1
            #   - https://github.com/benoitc/gunicorn/issues/1487
            #   - http://erickt.github.io/blog/2014/11/19/adventures-in-debugging-\
            #     a-potential-osx-kernel-bug/
            exc = exc_info.value
            if isinstance(exc, ConnectionError):
                pass
            else:
                assert exc.errno in (errno.EPROTOTYPE, errno.EPIPE)
        # Assert that we can connect properly afterwards (that is, server is not
        # terminated)
        self._okay(handshake)


@controller_data(proxy_protocol_timeout=0.3)
@handler_data(class_=ProxyPeekerHandler)
class TestHandlerAcceptReject:
    # We test *both* Accept *and* reject to ensure that the handshakes are valid
    @parametrize("handler_retval", [True, False])
    @parametrize(
        "handshake",
        [
            b"PROXY TCP4 255.255.255.255 255.255.255.255 65535 65535\r\n",
            TEST_V2_DATA1_EXACT,
        ],
        ids=["v1", "v2"],
    )
    def test_simple(self, plain_controller, handshake, handler_retval):
        assert plain_controller.smtpd._proxy_timeout > 0.0
        assert isinstance(plain_controller.handler, ProxyPeekerHandler)
        plain_controller.handler.retval = handler_retval
        if handler_retval:
            oper = operator.ne
            # See "Parametrizing conditional raising" in
            # https://docs.pytest.org/en/stable/example/parametrize.html
            expect = does_not_raise()
        else:
            oper = operator.eq
            expect = pytest.raises(SMTPServerDisconnected)
        oper = operator.ne if handler_retval else operator.eq
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.connect(Global.SrvAddr)
            sock.sendall(handshake)
            resp = sock.recv(4096)
            assert oper(resp, b"")
            with expect, SMTPClient() as client:
                client.sock = sock
                code, mesg = client.ehlo("example.org")
                assert code == 250