"""Unit tests for preseem.snmp_probe (the SNMP v1/v2c probe engine).

Covers the BER codec, response decoding, version selection / v1-fallback, ICMP
error parsing, and the async probe loop against in-process fake agents. The
subnet/CLI layers that wrap this engine (scanner, snmpscan) live in
tools/olt-scan and are tested there.
"""
import asyncio
import errno
import socket
import struct
import sys
import time
import unittest

from preseem import snmp_probe as s


def _cfg(**kw):
    """A ProbeConfig with fast test defaults; override via kwargs."""
    kw.setdefault('rate', 100000.0)
    kw.setdefault('burst', 50)
    kw.setdefault('max_inflight', 500)
    kw.setdefault('timeout', 1.0)
    return s.ProbeConfig(**kw)


def _build_response(request_id,
                    community,
                    varbinds,
                    error_status=0,
                    version=s.SNMP_V2C):
    """Build an SNMP Response PDU (mirror of build_get) for decode tests.

    varbinds: list of (oid, tag, value_bytes).
    """
    vbs = bytearray()
    for oid, tag, val in varbinds:
        vb = s.ber_tlv(s.T_OID, s.ber_oid(oid)) + s.ber_tlv(tag, val)
        vbs += s.ber_tlv(s.T_SEQUENCE, vb)
    pdu = (s.ber_tlv(s.T_INTEGER, s.ber_uint(request_id)) +
           s.ber_tlv(s.T_INTEGER, s.ber_uint(error_status)) +
           s.ber_tlv(s.T_INTEGER, s.ber_uint(0)) + s.ber_tlv(s.T_SEQUENCE, bytes(vbs)))
    msg = (s.ber_tlv(s.T_INTEGER, s.ber_uint(version)) +
           s.ber_tlv(s.T_OCTETSTR, community.encode()) +
           s.ber_tlv(s.T_GET_RESPONSE, pdu))
    return s.ber_tlv(s.T_SEQUENCE, msg)


def _parse_get(data):
    """Extract (version, request_id) from a GET we received (fake-agent side)."""
    _, msg, _ = s._read_tlv(data, 0)
    _, ver, i = s._read_tlv(msg, 0)
    _, _comm, i = s._read_tlv(msg, i)
    _, pdu, _ = s._read_tlv(msg, i)
    _, rid, _ = s._read_tlv(pdu, 0)
    return int.from_bytes(ver, 'big'), int.from_bytes(rid, 'big')


class TestBer(unittest.TestCase):

    def test_ber_len(self):
        self.assertEqual(s.ber_len(0), b'\x00')
        self.assertEqual(s.ber_len(127), b'\x7f')
        self.assertEqual(s.ber_len(128), b'\x81\x80')
        self.assertEqual(s.ber_len(256), b'\x82\x01\x00')

    def test_ber_uint(self):
        self.assertEqual(s.ber_uint(0), b'\x00')
        self.assertEqual(s.ber_uint(1), b'\x01')
        self.assertEqual(s.ber_uint(127), b'\x7f')
        # 128 has the top bit set -> needs a leading 0x00 to stay positive.
        self.assertEqual(s.ber_uint(128), b'\x00\x80')
        self.assertEqual(s.ber_uint(255), b'\x00\xff')
        self.assertEqual(s.ber_uint(256), b'\x01\x00')

    def test_ber_oid_known_vector(self):
        # 1.3.6.1.2.1.1.1.0 -> 0x2b then each sub-id (all < 128).
        self.assertEqual(s.ber_oid('1.3.6.1.2.1.1.1.0'),
                         bytes([0x2b, 0x06, 0x01, 0x02, 0x01, 0x01, 0x01, 0x00]))
        # Multi-byte sub-id: Huawei enterprise 2011 -> 0x8f 0x5b (base-128).
        self.assertEqual(s.ber_oid('1.3.6.1.4.1.2011'),
                         bytes([0x2b, 0x06, 0x01, 0x04, 0x01, 0x8f, 0x5b]))

    def test_oid_roundtrip(self):
        for oid in (
                '1.3.6.1.2.1.1.5.0',
                '1.3.6.1.4.1.2011.2.1.1',
                '0.0',
                '2.5.4.3',
                '1.3.6.1.4.1.99999999',
                '2.100',  # first sub-id 180 (>=120): arc1 stays 2, arc2 = 100
                '2.999',  # first sub-id 1079: multi-byte first sub-identifier
                '2.40',  # first sub-id 120: boundary just past the /40 split
        ):
            self.assertEqual(s._decode_oid(s.ber_oid(oid)), oid, oid)

    def test_assemble_get_matches_build_get(self):
        # The precompute path (encode_varbinds once + assemble_get) must produce
        # byte-identical packets to build_get, for both versions.
        vbs = s.encode_varbinds(s.SCAN_OIDS)
        self.assertEqual(vbs, s._SCAN_VARBINDS)
        for version in (s.SNMP_V1, s.SNMP_V2C):
            self.assertEqual(s.assemble_get(42, 'public', vbs, version),
                             s.build_get(42, 'public', s.SCAN_OIDS, version))


class TestDecode(unittest.TestCase):

    def test_decode_response(self):
        # Long OCTET STRING (>127) exercises long-form length on the wire.
        descr = 'X' * 200
        pkt = _build_response(12345, 'public', [
            ('1.3.6.1.2.1.1.1.0', s.T_OCTETSTR, descr.encode()),
            ('1.3.6.1.2.1.1.2.0', s.T_OID, s.ber_oid('1.3.6.1.4.1.2011.2.240')),
            ('1.3.6.1.2.1.1.3.0', s.T_TIMETICKS, s.ber_uint(987654)),
            ('1.3.6.1.2.1.1.5.0', s.T_OCTETSTR, b'olt-01'),
        ])
        resp = s.decode_response(pkt)
        self.assertIsNotNone(resp)
        self.assertEqual(resp.version, s.SNMP_V2C)
        self.assertEqual(resp.request_id, 12345)
        self.assertEqual(resp.error_status, 0)
        vals = dict(resp.varbinds)
        self.assertEqual(vals['1.3.6.1.2.1.1.1.0'], descr)
        self.assertEqual(vals['1.3.6.1.2.1.1.2.0'], '1.3.6.1.4.1.2011.2.240')
        self.assertEqual(vals['1.3.6.1.2.1.1.3.0'], 987654)
        self.assertEqual(vals['1.3.6.1.2.1.1.5.0'], 'olt-01')

    def test_decode_rejects_non_response(self):
        # A GET request must not decode as a response.
        getpkt = s.build_get(7, 'public', s.SCAN_OIDS)
        self.assertIsNone(s.decode_response(getpkt))

    def test_decode_exception_varbind(self):
        pkt = _build_response(1, 'public', [
            ('1.3.6.1.2.1.1.5.0', s.T_NOSUCHINSTANCE, b''),
        ])
        resp = s.decode_response(pkt)
        self.assertEqual(dict(resp.varbinds)['1.3.6.1.2.1.1.5.0'], 'noSuchInstance')


class TestVersion(unittest.TestCase):

    def test_build_get_version_byte(self):
        # The version INTEGER is the 3rd TLV inside the message SEQUENCE.
        for version in (s.SNMP_V1, s.SNMP_V2C):
            pkt = s.build_get(1, 'public', s.SCAN_OIDS, version)
            _, msg, _ = s._read_tlv(pkt, 0)
            tag, ver_body, _ = s._read_tlv(msg, 0)
            self.assertEqual(tag, s.T_INTEGER)
            self.assertEqual(int.from_bytes(ver_body, 'big'), version)
        # Default is v2c.
        _, msg, _ = s._read_tlv(s.build_get(1, 'public', s.SCAN_OIDS), 0)
        _, ver_body, _ = s._read_tlv(msg, 0)
        self.assertEqual(int.from_bytes(ver_body, 'big'), s.SNMP_V2C)

    def test_decode_v1_response_version(self):
        pkt = _build_response(9,
                              'public', [('1.3.6.1.2.1.1.5.0', s.T_OCTETSTR, b'olt')],
                              version=s.SNMP_V1)
        resp = s.decode_response(pkt)
        self.assertEqual(resp.version, s.SNMP_V1)
        self.assertEqual(dict(resp.varbinds)['1.3.6.1.2.1.1.5.0'], 'olt')

    def _version_seq(self, attempts, **cfg):
        prober = s.Prober(_cfg(**cfg), lambda r: None)
        return [prober._version_for(a) for a in attempts]

    def test_version_for_single_v2c(self):
        self.assertEqual(
            self._version_seq((1, 2), retries=1, primary_version=s.SNMP_V2C),
            [s.SNMP_V2C, s.SNMP_V2C])

    def test_version_for_single_v1(self):
        self.assertEqual(
            self._version_seq((1, 2, 3), retries=2, primary_version=s.SNMP_V1),
            [s.SNMP_V1] * 3)

    def test_version_for_fallback_last_attempt_only(self):
        # retries=2 -> attempts 1,2 are v2c (primary), attempt 3 (last) is v1.
        self.assertEqual(
            self._version_seq((1, 2, 3),
                              retries=2,
                              primary_version=s.SNMP_V2C,
                              fallback_version=s.SNMP_V1),
            [s.SNMP_V2C, s.SNMP_V2C, s.SNMP_V1])
        # retries=1 -> v2c then v1.
        self.assertEqual(
            self._version_seq((1, 2),
                              retries=1,
                              primary_version=s.SNMP_V2C,
                              fallback_version=s.SNMP_V1), [s.SNMP_V2C, s.SNMP_V1])


class TestIcmpParse(unittest.TestCase):

    def test_classify_icmp(self):
        # ICMP origin, dest-unreach type 3.
        self.assertEqual(s.classify_icmp(s.SO_EE_ORIGIN_ICMP, 3, 3, errno.ECONNREFUSED),
                         'port-closed')
        self.assertEqual(s.classify_icmp(s.SO_EE_ORIGIN_ICMP, 3, 1, 0),
                         'host-unreachable')
        self.assertEqual(s.classify_icmp(s.SO_EE_ORIGIN_ICMP, 3, 13, 0),
                         'admin-prohibited')
        self.assertEqual(s.classify_icmp(s.SO_EE_ORIGIN_ICMP, 11, 0, 0), 'ttl-exceeded')
        # Local origin -> errno driven.
        self.assertEqual(
            s.classify_icmp(s.SO_EE_ORIGIN_LOCAL, 0, 0, errno.EHOSTUNREACH),
            'host-unreachable')

    @staticmethod
    def _make_ee(ee_errno, origin, etype, code, offender_ip=None):
        """Build a sock_extended_err (+ optional SO_EE_OFFENDER sockaddr_in)."""
        ee = struct.pack('=IBBBBII', ee_errno, origin, etype, code, 0, 0, 0)
        if offender_ip is not None:
            parts = bytes(int(x) for x in offender_ip.split('.'))
            ee += struct.pack('=HH', socket.AF_INET, 0) + parts + b'\x00' * 8
        return ee

    def test_parse_icmp_error_uses_msg_name(self):
        cdata = self._make_ee(errno.ECONNREFUSED, s.SO_EE_ORIGIN_ICMP, 3, 3)
        anc = [(socket.IPPROTO_IP, s.IP_RECVERR, cdata)]
        ip, status = s.parse_icmp_error(anc, ('10.9.8.7', 0))
        self.assertEqual(ip, '10.9.8.7')
        self.assertEqual(status, 'port-closed')

    def test_parse_icmp_error_offender_fallback(self):
        # No msg_name -> fall back to the SO_EE_OFFENDER sockaddr_in.
        cdata = self._make_ee(errno.ECONNREFUSED,
                              s.SO_EE_ORIGIN_ICMP,
                              3,
                              3,
                              offender_ip='203.0.113.5')
        anc = [(socket.IPPROTO_IP, s.IP_RECVERR, cdata)]
        ip, status = s.parse_icmp_error(anc, None)
        self.assertEqual(ip, '203.0.113.5')
        self.assertEqual(status, 'port-closed')

    def test_parse_icmp_error_ignores_unrelated_cmsg(self):
        ip, status = s.parse_icmp_error([(socket.SOL_SOCKET, 1, b'\x00')], None)
        self.assertIsNone(ip)
        self.assertIsNone(status)

    def test_record_icmp_only_resolves_terminal_statuses(self):
        prober = s.Prober(_cfg(), lambda r: None)
        # Transient status must NOT mark the host unreachable (keep probing it).
        prober._record_icmp('10.0.0.1', 'ttl-exceeded')
        self.assertNotIn('10.0.0.1', prober.unreachable)
        # A terminal status does resolve it.
        prober._record_icmp('10.0.0.2', 'port-closed')
        self.assertEqual(prober.unreachable.get('10.0.0.2'), 'port-closed')


def _respond_any(data):
    """Answer any GET with a v2c sysinfo response."""
    _ver, rid = _parse_get(data)
    return _build_response(rid, 'public', [
        ('1.3.6.1.2.1.1.1.0', s.T_OCTETSTR, b'FakeOLT 1.0'),
        ('1.3.6.1.2.1.1.5.0', s.T_OCTETSTR, b'fake-olt'),
    ])


def _respond_v1_only(data):
    """Answer only v1 GETs (ignore the v2c probe), version 0."""
    version, rid = _parse_get(data)
    if version != s.SNMP_V1:
        return None
    return _build_response(rid,
                           'public',
                           [('1.3.6.1.2.1.1.5.0', s.T_OCTETSTR, b'v1-olt')],
                           version=s.SNMP_V1)


class _AgentProtocol(asyncio.DatagramProtocol):
    """A UDP fake agent driven by a responder(data) -> reply_bytes|None."""

    def __init__(self, responder):
        self._responder = responder

    def connection_made(self, transport):
        self._transport = transport

    def datagram_received(self, data, addr):
        reply = self._responder(data)
        if reply is not None:
            self._transport.sendto(reply, addr)


async def _start_agent(responder):
    """Bind a fake agent on loopback; return (transport, port).

    Uses create_datagram_endpoint (available since 3.4) rather than
    loop.sock_recvfrom (3.11+), so the tests run on the SDK's full range.
    """
    loop = asyncio.get_running_loop()
    transport, _proto = await loop.create_datagram_endpoint(
        lambda: _AgentProtocol(responder), local_addr=('127.0.0.1', 0))
    return transport, transport.get_extra_info('sockname')[1]


class TestConfig(unittest.TestCase):

    def test_close_tolerates_unopened_prober(self):
        # If open() fails before the socket exists, close() (run's finally)
        # must be a no-op rather than raising on a None socket.
        prober = s.Prober(_cfg(), lambda r: None)
        prober.close(loop=None)  # must not raise

    def test_rejects_nonpositive_rate_or_burst(self):
        with self.assertRaises(ValueError):
            s.ProbeConfig(rate=0)
        with self.assertRaises(ValueError):
            s.ProbeConfig(burst=0)

    def test_fallback_requires_a_retry(self):
        # retries=0 + fallback would make the sole attempt use the fallback
        # version, never sending the primary -- reject it.
        with self.assertRaises(ValueError):
            s.ProbeConfig(retries=0, fallback_version=s.SNMP_V1)
        # retries >= 1 with a fallback is fine, as is retries=0 without one.
        s.ProbeConfig(retries=1, fallback_version=s.SNMP_V1)
        s.ProbeConfig(retries=0)


class _OneShotFailSock:
    """Fake socket whose sendto raises OSError the first `fail` times."""

    def __init__(self, fail=1):
        self.fail = fail
        self.calls = 0

    def sendto(self, packet, addr):
        self.calls += 1
        if self.calls <= self.fail:
            raise ConnectionRefusedError('queued ICMP error')
        return len(packet)


class TestSend(unittest.TestCase):

    def test_send_drains_and_retries_once_in_icmp_mode(self):
        # A sendto OSError in icmp mode is treated as a possibly-unrelated
        # queued ICMP error: drain and retry once rather than dropping the host.
        prober = s.Prober(_cfg(), lambda r: None)
        prober.icmp_enabled = True
        prober._drain_errqueue = lambda: None  # queue mechanics not under test
        prober.sock = _OneShotFailSock(fail=1)
        self.assertTrue(asyncio.run(prober._send(b'pkt', '10.0.0.1')))
        self.assertEqual(prober.sock.calls, 2)   # failed once, retried, sent
        self.assertEqual(prober.send_errors, 0)

    def test_send_gives_up_without_icmp(self):
        # Without icmp there is no async-error ambiguity: one OSError is fatal.
        prober = s.Prober(_cfg(), lambda r: None)
        prober.sock = _OneShotFailSock(fail=1)
        self.assertFalse(asyncio.run(prober._send(b'pkt', '10.0.0.1')))
        self.assertEqual(prober.sock.calls, 1)
        self.assertEqual(prober.send_errors, 1)


class TestProbeLoop(unittest.TestCase):
    """Drive the engine (snmp_probe.run) against in-process fake agents."""

    def test_label_for_fault_is_isolated(self):
        # A raising label_for must not abort the scan; the record is still
        # delivered, just without a label.
        records = []

        def raising_label(ip):
            raise ValueError('bad labeler')

        prober = s.Prober(_cfg(icmp=True), records.append, label_for=raising_label)
        prober._record_icmp('10.0.0.1', 'port-closed')  # must not raise
        self.assertEqual(len(records), 1)
        self.assertIsNone(records[0]['subnet'])
        self.assertEqual(records[0]['status'], 'port-closed')

    def test_reply_from_wrong_source_keeps_pending(self):
        # A reply whose request-id matches a pending entry but whose source IP
        # does not must be ignored WITHOUT dropping that entry (else a stray or
        # spoofed packet would silently cancel a real in-flight probe).
        async def go():
            loop = asyncio.get_running_loop()
            prober = s.Prober(_cfg(), lambda r: None)
            # Bind directly (no add_reader) so the packet waits for our manual
            # _drain_data call instead of being drained by the loop.
            prober.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            prober.sock.setblocking(False)
            prober.sock.bind(('127.0.0.1', 0))
            port = prober.sock.getsockname()[1]
            try:
                rid = 4242
                prober.pending[rid] = s.Pending(
                    s.Target('grp', '10.0.0.1', 'public'),  # not loopback
                    time.monotonic() + 5,
                    1)
                pkt = _build_response(rid, 'public',
                                      [('1.3.6.1.2.1.1.5.0', s.T_OCTETSTR, b'x')])
                sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                sender.sendto(pkt, ('127.0.0.1', port))  # source is 127.0.0.1
                sender.close()
                await asyncio.sleep(0.05)
                prober._drain_data()
                return rid in prober.pending, dict(prober.responders)
            finally:
                prober.sock.close()

        still_pending, responders = asyncio.run(go())
        self.assertTrue(still_pending)  # entry survived the mismatch
        self.assertNotIn('10.0.0.1', responders)  # and was not recorded

    def test_scan_against_fake_agent(self):

        async def go():
            transport, port = await _start_agent(_respond_any)
            cfg = _cfg(port=port)
            records = []
            targets = [s.Target('loopback', '127.0.0.1', 'public')]
            prober, _ = await s.run(cfg, records.append, targets)
            transport.close()
            return prober, records

        prober, records = asyncio.run(go())
        self.assertEqual(len(records), 1, records)
        rec = records[0]
        self.assertEqual(rec['ip'], '127.0.0.1')
        self.assertEqual(rec['subnet'], 'loopback')
        self.assertEqual(rec['status'], 'responder')
        self.assertEqual(rec['version'], '2c')
        self.assertEqual(rec['sysDescr'], 'FakeOLT 1.0')
        self.assertEqual(rec['sysName'], 'fake-olt')

    def test_callback_exception_is_isolated(self):

        async def go():
            transport, port = await _start_agent(_respond_any)
            cfg = _cfg(port=port)

            def boom(rec):
                raise ValueError('consumer bug')

            targets = [s.Target('lo', '127.0.0.1', 'public')]
            prober, _ = await s.run(cfg, boom, targets)
            transport.close()
            return prober

        prober = asyncio.run(go())
        # A raising callback must not abort the scan; it is counted instead.
        self.assertEqual(prober.callback_errors, 1)
        self.assertIn('127.0.0.1', prober.responders)  # result still recorded

    def test_v1_fallback_discovers_v1_only_host(self):

        async def go():
            transport, port = await _start_agent(_respond_v1_only)
            cfg = _cfg(port=port,
                       timeout=0.3,
                       retries=1,
                       primary_version=s.SNMP_V2C,
                       fallback_version=s.SNMP_V1)
            records = []
            targets = [s.Target('lo', '127.0.0.1', 'public')]
            prober, _ = await s.run(cfg, records.append, targets)
            transport.close()
            return prober, records

        prober, records = asyncio.run(go())
        # v2c attempt got nothing; the final-retry v1 probe found it, and we
        # recorded the version the reply actually carried.
        self.assertEqual(len(records), 1, records)
        self.assertEqual(records[0]['version'], '1')
        self.assertEqual(records[0]['sysName'], 'v1-olt')
        self.assertGreaterEqual(prober.retransmits, 1)  # v1 fallback = retransmit

    @unittest.skipUnless(sys.platform.startswith('linux'),
                         'IP_RECVERR error queue is Linux-only')
    def test_icmp_port_closed_on_loopback(self):

        async def go():
            # A loopback UDP port with nothing bound draws an ICMP
            # port-unreachable back through the error queue.
            probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            probe.bind(('127.0.0.1', 0))
            port = probe.getsockname()[1]
            probe.close()
            cfg = _cfg(port=port, timeout=0.5, retries=0, icmp=True)
            records = []
            targets = [s.Target('lo', '127.0.0.1', 'public')]
            prober, _ = await s.run(cfg,
                                    records.append,
                                    targets,
                                    label_for=lambda ip: 'lo')
            return prober, records

        prober, records = asyncio.run(go())
        if not prober.icmp_enabled:
            self.skipTest('IP_RECVERR not supported in this environment')
        self.assertEqual(prober.unreachable.get('127.0.0.1'), 'port-closed',
                         prober.unreachable)
        self.assertFalse(prober.responders)
        self.assertTrue(records)
        self.assertEqual(records[0]['status'], 'port-closed')
        self.assertEqual(records[0]['subnet'], 'lo')


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