def test_waitForResponse(self): pkts = [AP.PingPDU(True, 1, 2, 3), AP.ResponsePDU(True, 10, 20, 30, 0, 0, 0), AP.ResponsePDU(True, 5, 23, 100, 0, 0, 0), AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0), AP.ResponsePDU(True, 100, 200, 300, 0, 0, 0)] opkt = AP.PingPDU(True, 42, 23, 100) def packetEater_jig(self): self.receivedPackets.append(pkts.pop(0)) p = AX.PacketControl(None, None) p.packetEater = (lambda : packetEater_jig(p)) p.spinGap = 0 # Test with session id match required ret = p.waitForResponse(opkt) self.assertEqual(pkts, [AP.ResponsePDU(True, 100, 200, 300, 0, 0, 0)]) self.assertEqual(ret, AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0)) # Test ignoring session id pkts = [AP.PingPDU(True, 1, 2, 3), AP.ResponsePDU(True, 10, 20, 30, 0, 0, 0), AP.ResponsePDU(True, 5, 23, 100, 0, 0, 0), AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0)] ret = p.waitForResponse(opkt, True) self.assertEqual(pkts, [AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0)]) self.assertEqual(ret, AP.ResponsePDU(True, 5, 23, 100, 0, 0, 0))
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_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_packetEater(self): pkt = AP.ResponsePDU(True, 42, 23, 100, 0, 0, 0) enc = pkt.encode() part1, rest = ntp.util.slicedata(enc, 15) part2, part3 = ntp.util.slicedata(rest, 6) pollReturns = [part1, part2, part3 + b"blah"] def pollSocket_jig(self): self.receivedData += pollReturns.pop(0) p = AX.PacketControl(None, None) p.sessionID = 42 p.pollSocket = (lambda: pollSocket_jig(p)) # Test incomplete header p.packetEater() self.assertEqual(p.receivedPackets, []) self.assertEqual(p.receivedData, part1) # Test incomplete packet p.packetEater() self.assertEqual(p.receivedPackets, []) self.assertEqual(p.receivedData, part1 + part2) # Test packet eaten p.packetEater() self.assertEqual(p.receivedPackets, [pkt]) self.assertEqual(p.receivedData, b"blah")
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_handle_CleanupSetPDU(self): handlerCalls = [] handler = (lambda a, vb: handlerCalls.append((a, vb))) mjig = AX.MIBControl() mjig.setVarbinds = ["foo", "bar"] mjig.setHandlers = [handler, handler] p = AX.PacketControl(None, mjig) p.sessionID = 42 pkt = AP.CleanupSetPDU(True, 42, 23, 100) p.handle_CleanupSetPDU(pkt) self.assertEqual(p.database.inSetP, False) self.assertEqual(handlerCalls, [("clean", "foo"), ("clean", "bar")])
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_checkResponses(self): returns = [] callback = (lambda resp, pkt: returns.append((resp, pkt))) faketimemod = jigs.TimeModuleJig() p = AX.PacketControl(None, None) p.packetLog = {(42, 23, 100): (25, "foo", callback), (42, 24, 200): (50, "bar", callback)} faketimemod.time_returns = [40] try: timetemp = AX.time AX.time = faketimemod p.checkResponses() self.assertEqual(returns, [(None, "foo")]) self.assertEqual(p.packetLog, {(42, 24, 200): (50, "bar", callback)}) finally: AX.time = timetemp
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_mainloop(self): stillCons = [] loop_calls = [0] cb_calls = [0] def loop_jig(self): loop_calls[0] += 1 self.stillConnected = stillCons.pop(0) def callback_jig(p): cb_calls[0] += 1 p = AX.PacketControl(None, None) p.spinGap = 0 p._doloop = (lambda: loop_jig(p)) p.stillConnected = False # Test not connected self.assertEqual(p.mainloop(False), False) self.assertEqual(loop_calls, [0]) # Test single shot p.stillConnected = True stillCons = [True] * 2 self.assertEqual(p.mainloop(False), True) self.assertEqual(stillCons, [True]) self.assertEqual(loop_calls, [1]) # Test loop p.stillConnected = True loop_calls[0] = 0 stillCons = [True] * 5 + [False] self.assertEqual(p.mainloop(True), False) self.assertEqual(stillCons, []) self.assertEqual(loop_calls, [6]) # Test loop, with callback p.stillConnected = True loop_calls[0] = 0 p.loopCallback = callback_jig stillCons = [True] * 5 + [False] self.assertEqual(p.mainloop(True), False) self.assertEqual(stillCons, []) self.assertEqual(loop_calls, [6]) self.assertEqual(cb_calls, [6])
def test___init__(self): p = AX.PacketControl("sock", "base") self.assertEqual(p.socket, "sock") self.assertEqual(0.0005 < p.spinGap < 0.0015, True) self.assertEqual(p.packetLog, {}) self.assertEqual(p.loopCallback, None) self.assertEqual(p.database, "base") self.assertEqual(p.receivedData, b"") self.assertEqual(p.timeout, 30) self.assertEqual(p.sessionID, None) self.assertEqual(p.highestTransactionID, 0) self.assertEqual(p.lastReception, None) self.assertEqual(p.stillConnected, False) self.assertEqual(p.pduHandlers, {AP.PDU_GET: p.handle_GetPDU, AP.PDU_GET_NEXT: p.handle_GetNextPDU, AP.PDU_GET_BULK: p.handle_GetBulkPDU, AP.PDU_TEST_SET: p.handle_TestSetPDU, AP.PDU_COMMIT_SET: p.handle_CommitSetPDU, AP.PDU_UNDO_SET: p.handle_UndoSetPDU, AP.PDU_CLEANUP_SET: p.handle_CleanupSetPDU, AP.PDU_RESPONSE: p.handle_ResponsePDU})
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