"""Unit tests for preseem.snmp_bulkwalk (native GETBULK subtree walk).

Covers the GETBULK PDU builder and OID helpers, and drives bulkwalk() against a
configurable in-process fake agent that serves a small MIB subtree and can
simulate the failure modes the walk must survive: MTU black holes (large
responses silently dropped), tooBig, dead hosts, out-of-order OIDs, and runaway
caps.
"""
import asyncio
import unittest

from preseem import snmp_bulkwalk as b
from preseem import snmp_engine as eng
from preseem import snmp_wire as w
from test import _snmp_fakes as fakes

ROOT = '1.3.6.1.2.1.2.2'  # ifTable
ROOT_B = '1.3.6.1.2.1.31.1.1'  # ifXTable


def make_mib(root, nrows, ncols=2):
    mib = []
    for col in range(1, ncols + 1):
        for idx in range(1, nrows + 1):
            mib.append((f'{root}.1.{col}.{idx}', w.T_INTEGER,
                        w.ber_uint(col * 100 + idx)))
    return mib


class FakeAgent:
    """A UDP SNMP agent serving a sorted MIB, with failure-mode knobs."""

    def __init__(self, mib, *, mtu_reps=None, toobig_reps=None, error_status=0,
                 drop_getbulk=False, drop_get=False, out_of_order=False):
        self.mib = sorted(mib, key=lambda e: w.oid_to_ints(e[0]))
        self.mtu_reps = mtu_reps  # drop GETBULK whose max-reps exceeds this
        self.toobig_reps = toobig_reps  # answer tooBig when max-reps exceeds this
        self.error_status = error_status  # non-zero -> return this error-status
        self.drop_getbulk = drop_getbulk
        self.drop_get = drop_get
        self.out_of_order = out_of_order
        self.getbulk_count = 0
        self.getbulk_drops = 0
        self.get_count = 0

    def handle(self, data):
        req = fakes.parse_request(data)
        ptag, rid, maxrep, oids = (req.pdu_tag, req.request_id,
                                   req.field2, req.oids)
        if ptag == w.T_GET_BULK:
            self.getbulk_count += 1
            if self.drop_getbulk or (self.mtu_reps is not None
                                     and maxrep > self.mtu_reps):
                self.getbulk_drops += 1
                return None
            if self.toobig_reps is not None and maxrep > self.toobig_reps:
                return fakes.build_response(rid, [], error_status=1)
            if self.error_status:
                return fakes.build_response(rid, [], error_status=self.error_status)
            start = w.oid_to_ints(oids[0])
            entries = [e for e in self.mib if w.oid_to_ints(e[0]) > start]
            if self.out_of_order and entries:
                return fakes.build_response(rid, [entries[0], entries[0]])
            n = max(1, maxrep)
            chosen = entries[:n]
            varbinds = list(chosen)
            if len(chosen) < n:  # ran out -> signal end of view
                varbinds.append(('1.3.6.1.2.1.99', w.T_ENDOFMIBVIEW, b''))
            return fakes.build_response(rid, varbinds)
        # GET (liveness probe or scalar read)
        self.get_count += 1
        if self.drop_get:
            return None
        return fakes.build_response(rid,
                               [(o, w.T_TIMETICKS, w.ber_uint(1)) for o in oids])


async def _run_walk(agent, root, **kw):
    transport, port = await fakes.start_agent(agent.handle)
    try:
        async with eng.SnmpEngine() as e:
            return await b.bulkwalk(e, '127.0.0.1', 'public', root, port=port,
                                    **kw)
    finally:
        transport.close()


class TestWire(unittest.TestCase):

    def test_getbulk_pdu_shape(self):
        pkt = w.build_getbulk(9, 'public', [ROOT], non_repeaters=0,
                              max_repetitions=7)
        _, msg, _ = w._read_tlv(pkt, 0)
        _, ver, i = w._read_tlv(msg, 0)
        self.assertEqual(int.from_bytes(ver, 'big'), w.SNMP_V2C)
        _, _comm, i = w._read_tlv(msg, i)
        ptag, pdu, _ = w._read_tlv(msg, i)
        self.assertEqual(ptag, w.T_GET_BULK)
        _, rid, j = w._read_tlv(pdu, 0)
        _, nr, j = w._read_tlv(pdu, j)
        _, mr, j = w._read_tlv(pdu, j)
        self.assertEqual(int.from_bytes(rid, 'big'), 9)
        self.assertEqual(int.from_bytes(nr, 'big'), 0)
        self.assertEqual(int.from_bytes(mr, 'big'), 7)

    def test_oid_helpers(self):
        self.assertTrue(w.oid_to_ints('1.3.9') < w.oid_to_ints('1.3.10'))
        self.assertTrue(
            w.oid_in_subtree(w.oid_to_ints(ROOT + '.1.1.1'),
                             w.oid_to_ints(ROOT)))
        self.assertFalse(
            w.oid_in_subtree(w.oid_to_ints('1.3.6.1.2.1.3.0'),
                             w.oid_to_ints(ROOT)))

    def test_decode_oid_as_ints(self):
        # oid_as_ints yields the numeric sub-id tuple; default yields the dotted
        # string. Both describe the same OID.
        pkt = fakes.build_response(1, [(ROOT + '.1.7.3', w.T_COUNTER32, w.ber_uint(9))])
        as_str = w.decode_response(pkt)
        as_ints = w.decode_response(pkt, oid_as_ints=True)
        self.assertEqual(as_str.varbinds[0][0], ROOT + '.1.7.3')
        self.assertEqual(as_ints.varbinds[0][0], w.oid_to_ints(ROOT + '.1.7.3'))
        self.assertEqual(as_ints.varbinds[0][1], 9)

    def test_decode_octets_are_bytes(self):
        # OCTET STRING and Opaque decode to raw bytes; text/formatting is the
        # caller's choice.
        pkt = fakes.build_response(1, [
            (ROOT + '.1.1.1', w.T_OCTETSTR, b'\x00\x1a\x2b'),
            (ROOT + '.1.1.2', w.T_OPAQUE, b'\x9f\x78'),
        ])
        vals = dict(w.decode_response(pkt).varbinds)
        self.assertEqual(vals[ROOT + '.1.1.1'], b'\x00\x1a\x2b')
        self.assertEqual(vals[ROOT + '.1.1.2'], b'\x9f\x78')

    def test_decode_over_memoryview_is_stale_byte_safe(self):
        # The engine decodes a memoryview slice of a reused receive buffer that
        # may hold stale bytes past the datagram; decode must read only the
        # datagram (bounded by TLV lengths) and return correct varbinds.
        mib = make_mib(ROOT, nrows=5, ncols=2)  # 10 varbinds
        pkt = fakes.build_response(55, [(o, t, v) for (o, t, v) in mib])
        buf = bytearray(65535)
        buf[:len(pkt)] = pkt
        for k in range(len(pkt), len(pkt) + 400):
            buf[k] = 0xEE  # stale garbage past the datagram
        mv = memoryview(buf)
        resp = w.decode_response(mv[:len(pkt)])
        self.assertEqual(resp.request_id, 55)
        self.assertEqual(len(resp.varbinds), 10)
        self.assertEqual(resp.varbinds[0][0], ROOT + '.1.1.1')

    def test_octet_string_endofmibview_not_terminal(self):
        # A real OCTET STRING whose bytes spell "endOfMibView" must NOT be taken
        # for the SnmpException end-of-view marker -- the walk keeps going. Over
        # the engine, OCTET STRING decodes to bytes, so there is no collision.
        mib = [(ROOT + '.1.1.1', w.T_OCTETSTR, b'endOfMibView'),
               (ROOT + '.1.1.2', w.T_INTEGER, w.ber_uint(7))]
        res = asyncio.run(
            _run_walk(FakeAgent(mib), ROOT, timeout=1.0, retries=0,
                      max_repetitions=4))
        self.assertEqual(res.status, 'ok')
        self.assertEqual(len(res.rows), 2)
        self.assertEqual(res.rows[0][1], b'endOfMibView')
        self.assertNotIsInstance(res.rows[0][1], w.SnmpException)


class TestWalk(unittest.TestCase):

    def test_full_walk_ok(self):
        mib = make_mib(ROOT, nrows=3, ncols=2)  # 6 rows
        res = asyncio.run(
            _run_walk(FakeAgent(mib), ROOT, timeout=1.0, retries=0,
                      max_repetitions=4))
        self.assertEqual(res.status, 'ok')
        self.assertEqual(res.version, '2c')
        self.assertEqual(len(res.rows), 6)
        # OIDs come back as integer sub-identifier tuples, strictly ordered.
        oids = [oid for oid, _ in res.rows]
        self.assertEqual(oids[0], w.oid_to_ints(ROOT + '.1.1.1'))
        self.assertEqual(oids, sorted(oids))

    def test_ends_at_subtree_exit(self):
        mib = make_mib(ROOT, nrows=2, ncols=2)  # 4 rows in-subtree
        mib.append(('1.3.6.1.2.1.3.0', w.T_INTEGER, w.ber_uint(1)))  # outside
        res = asyncio.run(
            _run_walk(FakeAgent(mib), ROOT, timeout=1.0, retries=0,
                      max_repetitions=10))
        self.assertEqual(res.status, 'ok')
        self.assertEqual(len(res.rows), 4)
        root_ints = w.oid_to_ints(ROOT)
        for oid, _ in res.rows:
            self.assertTrue(w.oid_in_subtree(oid, root_ints))

    def test_loop_guard_out_of_order(self):
        mib = make_mib(ROOT, nrows=3)
        res = asyncio.run(
            _run_walk(FakeAgent(mib, out_of_order=True), ROOT, timeout=1.0,
                      retries=0, max_repetitions=4))
        self.assertEqual(res.status, 'oid-not-increasing')
        self.assertEqual(len(res.rows), 1)  # partial rows returned, no hang

    def test_mtu_blackhole_reduces_by_one(self):
        mib = make_mib(ROOT, nrows=3, ncols=1)  # 3 rows
        agent = FakeAgent(mib, mtu_reps=2)  # drop any GETBULK with reps > 2
        res = asyncio.run(
            _run_walk(agent, ROOT, timeout=0.1, retries=0, max_repetitions=5,
                      min_repetitions=1))
        self.assertEqual(res.status, 'ok')
        self.assertEqual(len(res.rows), 3)
        # reps 5,4,3 each dropped once before 2 works: proves by-one descent
        # (halving 5//2 would drop only once).
        self.assertEqual(agent.getbulk_drops, 3)
        self.assertEqual(agent.get_count, 3)  # one liveness GET per drop

    def test_dead_host_times_out(self):
        agent = FakeAgent(make_mib(ROOT, 3), drop_getbulk=True, drop_get=True)
        res = asyncio.run(
            _run_walk(agent, ROOT, timeout=0.05, retries=0, max_repetitions=3))
        self.assertEqual(res.status, 'timeout')

    def test_alive_but_bulk_blocked(self):
        # GETBULK always dropped, but GET answered -> alive, can't bulk.
        agent = FakeAgent(make_mib(ROOT, 3), drop_getbulk=True)
        res = asyncio.run(
            _run_walk(agent, ROOT, timeout=0.05, retries=0, max_repetitions=3,
                      min_repetitions=1))
        self.assertEqual(res.status, 'bulk-failed')
        self.assertGreaterEqual(agent.get_count, 1)

    def test_toobig_reduces_then_completes(self):
        mib = make_mib(ROOT, nrows=3, ncols=1)
        agent = FakeAgent(mib, toobig_reps=2)  # tooBig unless reps <= 2
        res = asyncio.run(
            _run_walk(agent, ROOT, timeout=1.0, retries=0, max_repetitions=5))
        self.assertEqual(res.status, 'ok')
        self.assertEqual(len(res.rows), 3)
        self.assertEqual(agent.get_count, 0)  # tooBig is answered; no liveness

    def test_toobig_at_min(self):
        agent = FakeAgent(make_mib(ROOT, 3), toobig_reps=0)  # always tooBig
        res = asyncio.run(
            _run_walk(agent, ROOT, timeout=1.0, retries=0, max_repetitions=4,
                      min_repetitions=1))
        self.assertEqual(res.status, 'too-big')

    def test_snmp_error(self):
        agent = FakeAgent(make_mib(ROOT, 3), error_status=5)  # genErr
        res = asyncio.run(
            _run_walk(agent, ROOT, timeout=1.0, retries=0, max_repetitions=4))
        self.assertEqual(res.status, 'snmp-error')

    def test_truncated_max_rows(self):
        mib = make_mib(ROOT, nrows=10, ncols=1)
        res = asyncio.run(
            _run_walk(FakeAgent(mib), ROOT, timeout=1.0, retries=0,
                      max_repetitions=10, max_rows=3))
        self.assertEqual(res.status, 'truncated')
        self.assertEqual(len(res.rows), 3)

    def test_truncated_max_requests(self):
        mib = make_mib(ROOT, nrows=10, ncols=1)
        res = asyncio.run(
            _run_walk(FakeAgent(mib), ROOT, timeout=1.0, retries=0,
                      max_repetitions=1, max_requests=2))
        self.assertEqual(res.status, 'truncated')
        self.assertLessEqual(len(res.rows), 2)

    def test_concurrent_walks_one_engine(self):
        mib = make_mib(ROOT, 2, 1) + make_mib(ROOT_B, 2, 1)
        agent = FakeAgent(mib)

        async def go():
            transport, port = await fakes.start_agent(agent.handle)
            try:
                async with eng.SnmpEngine() as e:
                    return await asyncio.gather(
                        b.bulkwalk(e, '127.0.0.1', 'public', ROOT, port=port,
                                   timeout=1.0, retries=0, max_repetitions=3),
                        b.bulkwalk(e, '127.0.0.1', 'public', ROOT_B, port=port,
                                   timeout=1.0, retries=0, max_repetitions=3))
            finally:
                transport.close()

        a, bb = asyncio.run(go())
        self.assertEqual((a.status, len(a.rows)), ('ok', 2))
        self.assertEqual((bb.status, len(bb.rows)), ('ok', 2))
        self.assertEqual(a.root, ROOT)
        self.assertEqual(bb.root, ROOT_B)

    def test_rejects_bad_repetitions(self):
        with self.assertRaises(ValueError):
            asyncio.run(b.bulkwalk(None, 'h', 'public', ROOT, min_repetitions=0))
        with self.assertRaises(ValueError):
            asyncio.run(
                b.bulkwalk(None, 'h', 'public', ROOT, max_repetitions=2,
                           min_repetitions=5))


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