def test_initNewSession(self): resp_returns = [] resp_calls = [] def waitForResponse_jig(opkt, ignoreSID=False): resp_calls.append((opkt, ignoreSID)) return resp_returns.pop(0) mjig = AX.MIBControl(mibRoot=(0, 1, 2), rangeSubid=0, upperBound=None, mibContext="foo") sock = jigs.SocketJig() p = AX.PacketControl(sock, mjig) p.waitForResponse = waitForResponse_jig resp_returns = [AP.ResponsePDU(True, 42, 1, 1, 0, 0, 0), AP.ResponsePDU(True, 42, 1, 1, 0, 0, 0)] p.initNewSession() self.assertEqual(resp_returns, []) self.assertEqual(len(sock.data), 2) openpkt = AP.decode_packet(sock.data[0])[0] registerpkt = AP.decode_packet(sock.data[1])[0] self.assertEqual(openpkt, AP.OpenPDU(True, 23, 0, 0, 30, (), "NTPsec SNMP subagent")) self.assertEqual(registerpkt, AP.RegisterPDU(True, 42, 1, 1, 30, 1, mjig.mib_rootOID(), mjig.mib_rangeSubid(), mjig.mib_upperBound(), mjig.mib_context())) self.assertEqual(p.sessionID, 42) self.assertEqual(p.stillConnected, True)
def test_handle_UndoSetPDU(self): handlerCalls = [] handlerReturns = [AP.ERR_NOERROR, AP.ERR_NOERROR] handler = (lambda a, vb, ud: (handlerCalls.append((a, vb, ud)), handlerReturns.pop(0))[1]) mjig = AX.MIBControl() mjig.setVarbinds = ["foo", "bar"] mjig.setHandlers = [handler, handler] mjig.setUndoData = ["fey", "yar"] sock = jigs.SocketJig() p = AX.PacketControl(sock, mjig) p.sessionID = 42 # Test success pkt = AP.UndoSetPDU(True, 42, 23, 100) p.handle_UndoSetPDU(pkt) self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0) self.assertEqual(sentPkt, respPkt) self.assertEqual(handlerCalls, [("undo", "foo", "fey"), ("undo", "bar", "yar")]) # Test failure sock.data = [] handlerCalls = [] handlerReturns = [AP.ERR_NOERROR, AP.ERR_GENERR] p.handle_UndoSetPDU(pkt) self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU(True, 42, 23, 100, 0, 5, 1) self.assertEqual(sentPkt, respPkt) self.assertEqual(handlerCalls, [("undo", "foo", "fey"), ('undo', 'bar', 'yar')])
def test_handleGetNextPDU(self): readReturns = [1] read = (lambda oid: AP.Varbind(AP.VALUE_INTEGER, oid, readReturns.pop(0))) mjig = AX.MIBControl() mjig.addNode((0, 0), read) mjig.addNode((0, 1), read) sock = jigs.SocketJig() p = AX.PacketControl(sock, mjig) p.sessionID = 42 # Test successful pkt = AP.GetNextPDU(True, 42, 23, 100, [AP.SearchRange((0, 0), ())]) p.handle_GetNextPDU(pkt) self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0, (AP.Varbind(AP.VALUE_INTEGER, (0, 1), 1),)) self.assertEqual(sentPkt, respPkt) # Test non-existent sock.data = [] pkt = AP.GetNextPDU(True, 42, 23, 100, [AP.SearchRange((1, 0), ())]) p.handle_GetNextPDU(pkt) self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0, (AP.Varbind(AP.VALUE_END_OF_MIB_VIEW, (1, 0)),)) self.assertEqual(sentPkt, respPkt)
def test_sendPing(self): sock = jigs.SocketJig() p = AX.PacketControl(sock, None) p.sessionID = 42 p.sendPing() self.assertEqual(sock.data, [b"\x01\r\x10\x00" b"\x00\x00\x00*\x00\x00\x00\x05" b"\x00\x00\x00\x01\x00\x00\x00\x00"]) self.assertEqual(list(p.packetLog.keys()), [(42, 5, 1)])
def test_sendErrorResponse(self): sock = jigs.SocketJig() p = AX.PacketControl(sock, None) p.sessionID = 42 header = {"flags": {"bigEndian": True}, "session_id": 42, "transaction_id": 32, "packet_id": 100} p.sendErrorResponse(header, AP.ERR_GENERR, 12) self.assertEqual(sock.data, [b"\x01\x12\x10\x00" b"\x00\x00\x00*\x00\x00\x00 " b"\x00\x00\x00d\x00\x00\x00\x08" b"\x00\x00\x00\x00\x00\x05\x00\x0c"]) self.assertEqual(p.packetLog, {})
def test_sendNotify(self): sock = jigs.SocketJig() p = AX.PacketControl(sock, None) p.sessionID = 42 vb = AP.Varbind(AP.VALUE_INTEGER, (0, 1, 2), 23) p.sendNotify([vb]) self.assertEqual(sock.data, [b"\x01\x0c\x10\x00" b"\x00\x00\x00*\x00\x00\x00\x05" b"\x00\x00\x00\x01\x00\x00\x00\x18" b"\x00\x02\x00\x00" b"\x03\x00\x00\x00\x00\x00\x00\x00" b"\x00\x00\x00\x01\x00\x00\x00\x02" b"\x00\x00\x00\x17"]) self.assertEqual(list(p.packetLog.keys()), [(42, 5, 1)])
def test_handle_ResponsePDU(self): cbreturn = [] def callback(pkt, origpkt): cbreturn.append((pkt, origpkt)) opkt = AP.PingPDU(True, 42, 4, 50) incorrect_rpkt = AP.ResponsePDU(True, 23, 1, 99, 0, 0, 0) correct_rpkt = AP.ResponsePDU(True, 42, 4, 50, 0, 0, 0) sock = jigs.SocketJig() p = AX.PacketControl(sock, None) p.sessionID = 42 p.packetLog = {(42, 4, 50): (30.0, opkt, callback)} # Test non-expected response (likely to happen in a timeout) p.handle_ResponsePDU(incorrect_rpkt) self.assertEqual(p.packetLog, {(42, 4, 50): (30.0, opkt, callback)}) # Test expected response p.handle_ResponsePDU(correct_rpkt) self.assertEqual(p.packetLog, {})
def test_pollSocket(self): fakeselectmod = jigs.SelectModuleJig() sock = jigs.SocketJig() p = AX.PacketControl(sock, None) p.sessionID = 42 # Test pkt = AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0) enc = pkt.encode() part1, part2 = ntp.util.slicedata(enc, 15) sock.return_data = [part1, part2, b"blah"] fakeselectmod.do_return = [True, True, True, False] try: selecttemp = AX.select AX.select = fakeselectmod p.pollSocket() self.assertEqual(p.receivedData, enc + b"blah") finally: AX.select = selecttemp
def test_sendPacket(self): sock = jigs.SocketJig() p = AX.PacketControl(sock, None) testpkt = AP.PingPDU(True, 0, 1, 2, "fake") # Test not expecting reply p.sendPacket(testpkt, False) self.assertEqual(sock.data, [b"\x01\r\x18\x00" b"\x00\x00\x00\x00\x00\x00\x00\x01" b"\x00\x00\x00\x02\x00\x00\x00\x08" b"\x00\x00\x00\x04fake"]) self.assertEqual(p.packetLog, {}) sock.data = [] # Test expecting reply p.sendPacket(testpkt, True, callback="foo") self.assertEqual(sock.data, [b"\x01\r\x18\x00" b"\x00\x00\x00\x00\x00\x00\x00\x01" b"\x00\x00\x00\x02\x00\x00\x00\x08" b"\x00\x00\x00\x04fake"]) self.assertEqual(p.packetLog, {(0, 1, 2): (30.0, testpkt, "foo")})
def test_handle_GetBulkPDU(self): readReturns = [1, 2, 3, 4, 5, 6] read = ( lambda oid: AP.Varbind(AP.VALUE_INTEGER, oid, readReturns.pop(0))) mjig = AX.MIBControl() # non-repeated mjig.addNode((0, 0), read) mjig.addNode((0, 1), read) # shouldn't see mjig.addNode((1, 0), read) mjig.addNode((1, 1), read) # shouldn't see # repeated mjig.addNode((2, 0), read) mjig.addNode((2, 1), read) mjig.addNode((2, 2), read) # shouldn't see mjig.addNode((3, 0), read) mjig.addNode((3, 1), read) mjig.addNode((3, 2), read) # shouldn't see mjig.addNode((4, 0), read) # shouldn't see sock = jigs.SocketJig() p = AX.PacketControl(sock, mjig) p.sessionID = 42 pkt = AP.GetBulkPDU(True, 42, 23, 100, 2, 2, [ AP.SearchRange((0, ), (1, ), True), AP.SearchRange((1, ), (2, ), True), AP.SearchRange((2, 0), (3, ), True), AP.SearchRange((3, 0), (), True) ]) p.handle_GetBulkPDU(pkt) self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU( True, 42, 23, 100, 0, 0, 0, (AP.Varbind(AP.VALUE_INTEGER, (0, 0), 1), AP.Varbind(AP.VALUE_INTEGER, (1, 0), 2), AP.Varbind(AP.VALUE_INTEGER, (2, 0), 3), AP.Varbind(AP.VALUE_INTEGER, (2, 1), 4), AP.Varbind(AP.VALUE_INTEGER, (3, 0), 5), AP.Varbind(AP.VALUE_INTEGER, (3, 1), 6))) self.assertEqual(sentPkt, respPkt)
def test_handle_TestSetPDU(self): mjig = AX.MIBControl() readReturns = [] read = (lambda oid: AP.Varbind(AP.VALUE_INTEGER, oid, readReturns.pop(0))) writeCalls = [] writeReturns = [] write = (lambda a, vb: (writeCalls.append((a, vb)), writeReturns.pop(0))[1]) mjig.addNode((0,), read, write) mjig.addNode((1,), read) sock = jigs.SocketJig() p = AX.PacketControl(sock, mjig) p.sessionID = 42 # Test success readReturns = [10] writeReturns = [AP.ERR_NOERROR] pkt = AP.TestSetPDU(True, 42, 23, 100, [AP.Varbind(AP.VALUE_INTEGER, (0,), 16)]) p.handle_TestSetPDU(pkt) self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0) self.assertEqual(sentPkt, respPkt) self.assertEqual(writeCalls, [("test", AP.Varbind(AP.VALUE_INTEGER, (0,), 16))]) self.assertEqual(mjig.setVarbinds, [AP.Varbind(AP.VALUE_INTEGER, (0,), 16)]) self.assertEqual(mjig.setHandlers, [write]) self.assertEqual(mjig.setUndoData, [AP.Varbind(AP.VALUE_INTEGER, (0,), 10)]) # Test failure, no access mjig.setVarbinds = [] mjig.setHandlers = [] mjig.setUndoData = [] sock.data = [] readReturns = [10] writeReturns = [AP.ERR_NOERROR] writeCalls = [] pkt = AP.TestSetPDU(True, 42, 23, 100, [AP.Varbind(AP.VALUE_INTEGER, (0, 1), 16)]) p.handle_TestSetPDU(pkt) self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU(True, 42, 23, 100, 0, 6, 0) self.assertEqual(sentPkt, respPkt) self.assertEqual(writeCalls, []) self.assertEqual(mjig.setVarbinds, []) self.assertEqual(mjig.setHandlers, []) self.assertEqual(mjig.setUndoData, []) # Test failure, not writable mjig.setVarbinds = [] mjig.setHandlers = [] mjig.setUndoData = [] sock.data = [] readReturns = [10] writeReturns = [AP.ERR_NOERROR] writeCalls = [] pkt = AP.TestSetPDU(True, 42, 23, 100, [AP.Varbind(AP.VALUE_INTEGER, (1,), 16)]) p.handle_TestSetPDU(pkt) self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU(True, 42, 23, 100, 0, 17, 0) self.assertEqual(sentPkt, respPkt) self.assertEqual(writeCalls, []) self.assertEqual(mjig.setVarbinds, []) self.assertEqual(mjig.setHandlers, []) self.assertEqual(mjig.setUndoData, [])
def test__doloop(self): # need jigs for # time pkts = [] # recpTimes = [] handler_calls = [0] def packetEater_jig(self): self.receivedPackets = pkts self.lastReception = recpTime def pinghandler_jig(self): handler_calls[0] += 1 faketimemod = jigs.TimeModuleJig() sock = jigs.SocketJig() p = AX.PacketControl(sock, None) p.packetEater = (lambda: packetEater_jig(p)) p.checkResponses = (lambda: None) p.pduHandlers[AP.PDU_PING] = pinghandler_jig p.sessionID = 42 try: timetemp = AX.time AX.time = faketimemod # Test first time, no reception recpTime = None faketimemod.time_returns = [] p._doloop() self.assertEqual(sock.data, []) # Test wrong session ID pkts = [AP.PingPDU(True, 13, 0, 0)] recpTime = 50 faketimemod.time_returns = [55] p._doloop() self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.ResponsePDU(True, 13, 0, 0, 0, AP.RSPERR_NOT_OPEN, 0) self.assertEqual(sentPkt, respPkt) # Test unimplemented packet type faketimemod.time_returns = [55] sock.data = [] pkts = [AP.PingPDU(True, 42, 0, 0)] pkts[0].pduType = 1337 # spoof the type p._doloop() self.assertEqual(sock.data, []) self.assertEqual(handler_calls, [0]) # Test successful faketimemod.time_returns = [55] sock.data = [] pkts = [AP.PingPDU(True, 42, 0, 0)] p._doloop() self.assertEqual(sock.data, []) self.assertEqual(handler_calls, [1]) # Test pinging handler_calls = [0] faketimemod.time_returns = [120] sock.data = [] p._doloop() self.assertEqual(len(sock.data), 1) sentPkt = AP.decode_packet(sock.data[0])[0] respPkt = AP.PingPDU(True, 42, 5, 1) self.assertEqual(sentPkt, respPkt) self.assertEqual(handler_calls, [0]) finally: AX.time = timetemp