145 lines
8.8 KiB
Python
145 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
|
|
|
"""
|
|
This file contains tests for the softflowd UDP collector saved in main.py
|
|
The test packets (defined below as hex streams) were extracted from a "real" softflowd export
|
|
based on a sample PCAP capture file. They consist of one export with the templates and three without.
|
|
|
|
Two tests are defined, one slow, one fast. During some runs exceptions occured which might hint
|
|
to race conditions during reading and writing to the JSON output file.
|
|
For now, both tests run successfully.
|
|
|
|
Copyright 2017-2019 Dominik Pataky <dev@bitkeks.eu>
|
|
Licensed under MIT License. See LICENSE.
|
|
"""
|
|
|
|
import ipaddress
|
|
import json
|
|
import logging
|
|
from pprint import pprint
|
|
import random
|
|
import socket
|
|
import socketserver
|
|
import subprocess
|
|
import tempfile
|
|
from time import sleep
|
|
import threading
|
|
import unittest
|
|
|
|
from main import SoftflowUDPHandler
|
|
|
|
|
|
logging.getLogger().setLevel(logging.DEBUG)
|
|
|
|
# The flowset with 2 templates and 8 flows
|
|
template_packet = '0009000a000000035c9f55980000000100000000000000400400000e00080004000c000400150004001600040001000400020004000a0004000e000400070002000b00020004000100060001003c000100050001000000400800000e001b0010001c001000150004001600040001000400020004000a0004000e000400070002000b00020004000100060001003c000100050001040001447f0000017f000001fb3c1aaafb3c18fd000190100000004b00000000000000000050942c061b04007f0000017f000001fb3c1aaafb3c18fd00000f94000000360000000000000000942c0050061f04007f0000017f000001fb3c1cfcfb3c1a9b0000d3fc0000002a000000000000000000509434061b04007f0000017f000001fb3c1cfcfb3c1a9b00000a490000001e000000000000000094340050061f04007f0000017f000001fb3bb82cfb3ba48b000002960000000300000000000000000050942a061904007f0000017f000001fb3bb82cfb3ba48b00000068000000020000000000000000942a0050061104007f0000017f000001fb3c1900fb3c18fe0000004c0000000100000000000000000035b3c9110004007f0000017f000001fb3c1900fb3c18fe0000003c000000010000000000000000b3c9003511000400'
|
|
|
|
# Three packets without templates, each with 12 flows, anonymized
|
|
packets = [
|
|
'0009000c000000035c9f55980000000200000000040001e47f0000017f000001fb3c1a17fb3c19fd000001480000000200000000000000000035ea82110004007f0000017f000001fb3c1a17fb3c19fd0000007a000000020000000000000000ea820035110004007f0000017f000001fb3c1a17fb3c19fd000000f80000000200000000000000000035c6e2110004007f0000017f000001fb3c1a17fb3c19fd0000007a000000020000000000000000c6e20035110004007f0000017f000001fb3c1a9efb3c1a9c0000004c0000000100000000000000000035adc1110004007f0000017f000001fb3c1a9efb3c1a9c0000003c000000010000000000000000adc10035110004007f0000017f000001fb3c1b74fb3c1b720000004c0000000100000000000000000035d0b3110004007f0000017f000001fb3c1b74fb3c1b720000003c000000010000000000000000d0b30035110004007f0000017f000001fb3c2f59fb3c1b7100001a350000000a000000000000000000509436061b04007f0000017f000001fb3c2f59fb3c1b710000038a0000000a000000000000000094360050061b04007f0000017f000001fb3c913bfb3c91380000004c0000000100000000000000000035e262110004007f0000017f000001fb3c913bfb3c91380000003c000000010000000000000000e262003511000400',
|
|
'0009000c000000035c9f55980000000300000000040001e47f0000017f000001fb3ca523fb3c913b0000030700000005000000000000000000509438061b04007f0000017f000001fb3ca523fb3c913b000002a200000005000000000000000094380050061b04007f0000017f000001fb3f7fe1fb3dbc970002d52800000097000000000000000001bb8730061b04007f0000017f000001fb3f7fe1fb3dbc970000146c000000520000000000000000873001bb061f04007f0000017f000001fb3d066ffb3d066c0000004c0000000100000000000000000035e5bd110004007f0000017f000001fb3d066ffb3d066c0000003c000000010000000000000000e5bd0035110004007f0000017f000001fb3d1a61fb3d066b000003060000000500000000000000000050943a061b04007f0000017f000001fb3d1a61fb3d066b000002a2000000050000000000000000943a0050061b04007f0000017f000001fb3fed00fb3f002c0000344000000016000000000000000001bbae50061f04007f0000017f000001fb3fed00fb3f002c00000a47000000120000000000000000ae5001bb061b04007f0000017f000001fb402f17fb402a750003524c000000a5000000000000000001bbc48c061b04007f0000017f000001fb402f17fb402a75000020a60000007e0000000000000000c48c01bb061f0400',
|
|
'0009000c000000035c9f55980000000400000000040001e47f0000017f000001fb3d7ba2fb3d7ba00000004c0000000100000000000000000035a399110004007f0000017f000001fb3d7ba2fb3d7ba00000003c000000010000000000000000a3990035110004007f0000017f000001fb3d8f85fb3d7b9f000003070000000500000000000000000050943c061b04007f0000017f000001fb3d8f85fb3d7b9f000002a2000000050000000000000000943c0050061b04007f0000017f000001fb3d9165fb3d7f6d0000c97b0000002a000000000000000001bbae48061b04007f0000017f000001fb3d9165fb3d7f6d000007f40000001a0000000000000000ae4801bb061b04007f0000017f000001fb3dbc96fb3dbc7e0000011e0000000200000000000000000035bd4f110004007f0000017f000001fb3dbc96fb3dbc7e0000008e000000020000000000000000bd4f0035110004007f0000017f000001fb3ddbb3fb3c1a180000bfee0000002f00000000000000000050ae56061b04007f0000017f000001fb3ddbb3fb3c1a1800000982000000270000000000000000ae560050061b04007f0000017f000001fb3ddbb3fb3c1a180000130e0000001200000000000000000050e820061b04007f0000017f000001fb3ddbb3fb3c1a180000059c000000140000000000000000e8200050061b0400'
|
|
]
|
|
|
|
|
|
class ThreadedUDPServer(socketserver.ThreadingMixIn, socketserver.UDPServer):
|
|
pass
|
|
|
|
class TestSoftflowExport(unittest.TestCase):
|
|
CONNECTION = ('127.0.0.1', 1337)
|
|
COUNT_PACKETS_TO_TEST = 5
|
|
SLEEP_TIME = 0.3
|
|
RUN_ANALYZER = False
|
|
|
|
def setUp(self):
|
|
logging.debug("Creating temporary JSON output file.")
|
|
self.temp_output_file = tempfile.NamedTemporaryFile(prefix="softflowd_")
|
|
|
|
# FIXME: templates are saved between test runs, because they are stored with the class
|
|
# Maybe the templates should be stored with an instance?
|
|
logging.debug("Resetting SoftflowUDPHandler templates.")
|
|
SoftflowUDPHandler.templates = {}
|
|
|
|
logging.debug("Setting temporary file {} as output for SoftflowUDPHandler".format(self.temp_output_file.name))
|
|
SoftflowUDPHandler.set_output_file(self.temp_output_file.name)
|
|
|
|
logging.debug("Writing empty dict to output file.")
|
|
with open(self.temp_output_file.name, "w") as fh:
|
|
json.dump({}, fh)
|
|
|
|
logging.debug("Creating and running the Softflow collector in another thread.")
|
|
self.server = ThreadedUDPServer(self.CONNECTION, SoftflowUDPHandler)
|
|
self.server_thread = threading.Thread(target=self.server.serve_forever)
|
|
self.server_thread.daemon = True
|
|
self.server_thread.start()
|
|
|
|
logging.debug("Creating UDP socket for client packets.")
|
|
self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
|
|
def tearDown(self):
|
|
logging.debug("Running tear down procedure.")
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.server_thread.join()
|
|
self.sock.close()
|
|
self.temp_output_file.close()
|
|
|
|
def _test_export(self):
|
|
logging.info("Running UDP client sending raw hex packets with flows.")
|
|
|
|
# Get a random index on which the template is sent
|
|
template_idx = random.randint(1, self.COUNT_PACKETS_TO_TEST - 1) # 1 for enhanced testing, -1 because randint
|
|
|
|
# Save the order of lengths for later check
|
|
lens = []
|
|
|
|
for idx in range(self.COUNT_PACKETS_TO_TEST):
|
|
# Choose a random packet payload
|
|
p = random.choice(packets)
|
|
|
|
logging.info("Sending packet {}.".format(idx))
|
|
self.sock.sendto(bytes.fromhex(p), self.CONNECTION)
|
|
lens.append(12)
|
|
sleep(self.SLEEP_TIME)
|
|
|
|
# Randomly inject the template packet
|
|
if idx == template_idx:
|
|
logging.info("Sending template packet.")
|
|
self.sock.sendto(bytes.fromhex(template_packet), self.CONNECTION)
|
|
lens.append(8)
|
|
sleep(self.SLEEP_TIME)
|
|
|
|
with open(self.temp_output_file.name, "r") as fh:
|
|
exported = json.load(fh)
|
|
|
|
# We got four exports
|
|
logging.info("Testing the existence of all exports, including the ones with formerly unknown templates: {} of {}".format(
|
|
len(exported.keys()), self.COUNT_PACKETS_TO_TEST + 1))
|
|
self.assertEqual(len(exported.keys()), self.COUNT_PACKETS_TO_TEST + 1) # +1 including the template packet
|
|
|
|
# Test lengths of exports
|
|
logging.info("Testing the correct lengths of all exports.")
|
|
for idx, val in enumerate(exported.values()):
|
|
self.assertEqual(len(val), lens[idx])
|
|
|
|
if self.RUN_ANALYZER:
|
|
logging.info("Running analyze_json.py")
|
|
analyzer = subprocess.run(['python3', 'analyze_json.py', self.temp_output_file.name], stdout=subprocess.PIPE)
|
|
for line in analyzer.stdout.split(b"\n"):
|
|
print(line.decode())
|
|
|
|
def test_slow(self):
|
|
logging.info("Running slow test")
|
|
self.SLEEP_TIME = 0.5
|
|
self.COUNT_PACKETS_TO_TEST = 3
|
|
self._test_export()
|
|
|
|
def test_fast(self):
|
|
logging.info("Running fast test")
|
|
self.SLEEP_TIME = 0.1
|
|
self.COUNT_PACKETS_TO_TEST = 30
|
|
self._test_export()
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|