class StorageTest(TestCase): SEL_SEC_SCAN = ( "SELECT scan_id, exploit_id, exploit_app, exploit_name, node_id, node_ip, port_protocol, " "port_number, sec_scan_start, sec_scan_end FROM security_scans", ) SEL_NOD_SCAN = ( "SELECT scan_id, node_id, node_ip, time from nodes_scans", ) SEL_POR_SCAN = ( "SELECT scan_id, node_id, node_ip, port, port_protocol, time from ports_scans", ) SEL_CHANGE = ( "SELECT type, vulnerability_id, vulnerability_subid, previous_id, current_id, time FROM changes", ) SEL_SCANS = ( "SELECT ROWID, scan_start, scan_end, protocol, scanner_name FROM scans", ) SEL_VULNS = ( "SELECT scan_id, node_id, node_ip, port_protocol, port, vulnerability_id, vulnerability_subid, cve, cvss, output, time FROM vulnerabilities", ) def setUp(self): self.maxDiff = None self.storage = Storage(conn_string=getenv('AUCOTE_TEST_POSTGRES')) self.scan = Scan(start=1, end=17, protocol=TransportProtocol.UDP, scanner='test_name') self.scan_1 = Scan(rowid=56, protocol=TransportProtocol.UDP, scanner='test_name', start=13, end=19) self.scan_2 = Scan(rowid=79, protocol=TransportProtocol.UDP, scanner='test_name', start=2, end=18) self.scan_3 = Scan(rowid=80, protocol=TransportProtocol.UDP, scanner='test_name_2', start=20, end=45) self.scan_4 = Scan(rowid=78, protocol=TransportProtocol.TCP, scanner='portdetection', start=1, end=2) self.node_1 = Node(node_id=1, ip=ipaddress.ip_address('127.0.0.1')) self.node_1.name = 'test_node_1' self.node_1.scan = Scan(start=15) self.node_2 = Node(node_id=2, ip=ipaddress.ip_address('127.0.0.2')) self.node_2.name = 'test_node_2' self.node_2.scan = Scan(start=56) self.node_3 = Node(node_id=3, ip=ipaddress.ip_address('127.0.0.3')) self.node_3.name = 'test_node_3' self.node_3.scan = Scan(start=98) self.node_4 = Node(node_id=4, ip=ipaddress.ip_address('127.0.0.4')) self.node_4.name = 'test_node_4' self.node_4.scan = Scan(start=3) self.node_scan_1 = NodeScan(node=self.node_1, rowid=13, scan=self.scan_1, timestamp=15) self.node_scan_2 = NodeScan(node=self.node_2, rowid=15, scan=self.scan_1, timestamp=56) self.node_scan_3 = NodeScan(node=self.node_3, rowid=16, scan=self.scan_1, timestamp=98) self.node_scan_3 = NodeScan(node=self.node_3, rowid=17, scan=self.scan_2, timestamp=90) self.port_1 = Port(node=self.node_1, number=45, transport_protocol=TransportProtocol.UDP) self.port_1.scan = self.scan_1 self.port_scan_1 = PortScan(port=self.port_1, scan=self.scan_1, timestamp=176, rowid=124) self.port_2 = Port(node=self.node_2, transport_protocol=TransportProtocol.UDP, number=65) self.port_2.scan = Scan(start=3, end=45) self.port_scan_2 = PortScan(port=self.port_2, scan=self.scan_1, timestamp=987, rowid=15) self.port_3 = Port(node=self.node_3, transport_protocol=TransportProtocol.ICMP, number=99) self.port_3.scan = Scan(start=43, end=180) self.port_scan_3 = PortScan(port=self.port_3, scan=self.scan_1, timestamp=619, rowid=13) self.port_4 = Port(node=self.node_1, number=80, transport_protocol=TransportProtocol.UDP) self.port_4.scan = self.scan_1 self.port_scan_4 = PortScan(port=self.port_4, scan=self.scan_1, timestamp=650, rowid=480) self.port_5 = Port(node=self.node_4, number=22, transport_protocol=TransportProtocol.TCP) self.port_5.scan = self.scan_4 self.exploit_1 = Exploit(exploit_id=14, name='test_name', app='test_app') self.exploit_2 = Exploit(exploit_id=2, name='test_name_2', app='test_app_2') self.exploit_3 = Exploit(exploit_id=56, name='test_name_2', app='test_app') self.exploit_4 = Exploit(exploit_id=0, name='portdetection', app='portdetection') self.security_scan_1 = SecurityScan(exploit=self.exploit_1, port=self.port_1, scan=self.scan_1, scan_start=178, scan_end=851) self.security_scan_2 = SecurityScan(exploit=self.exploit_2, port=self.port_1, scan=self.scan_1, scan_start=109, scan_end=775) self.security_scan_3 = SecurityScan(exploit=self.exploit_3, port=self.port_1, scan=self.scan_2, scan_start=113, scan_end=353) self.security_scan_4 = SecurityScan(exploit=self.exploit_1, port=self.port_1, scan=self.scan_3, scan_start=180, scan_end=222) self.security_scan_5 = SecurityScan(exploit=self.exploit_1, port=self.port_1, scan=self.scan_2, scan_start=14, scan_end=156) self.security_scan_6 = SecurityScan(exploit=self.exploit_2, port=self.port_1, scan=self.scan_2, scan_start=56, scan_end=780) self.security_scan_7 = SecurityScan(exploit=self.exploit_4, port=self.port_5, scan=self.scan_4, scan_start=14, scan_end=890) self.vuln_change_1 = PortDetectionChange( change_time=124445, current_finding=self.port_scan_1, previous_finding=self.port_scan_2) self.vuln_change_2 = PortDetectionChange( change_time=32434, current_finding=self.port_scan_2, previous_finding=self.port_scan_3) self.vulnerability_1 = Vulnerability(port=self.port_1, output='test_output_1', exploit=self.exploit_1, cve='CVE-2017', cvss=6.7, subid=1, vuln_time=13, rowid=134, scan=self.scan_1) self.vulnerability_2 = Vulnerability(port=self.port_1, output='test_output_2', exploit=self.exploit_2, cve='CWE-14', cvss=8.9, subid=2, vuln_time=98, rowid=152, scan=self.scan_1) self.vulnerability_3 = Vulnerability(port=self.port_1, output='test_output_3', exploit=self.exploit_1, cve='CVE-2016', cvss=3.7, subid=2, vuln_time=15, rowid=153, scan=self.scan_2) self.vulnerability_4 = Vulnerability(port=self.port_1, output='test_output_4', exploit=self.exploit_2, cve='CWE-15', cvss=2.9, subid=1, vuln_time=124, rowid=169, scan=self.scan_2) self.vulnerability_5 = Vulnerability(port=self.port_5, output='', exploit=self.exploit_4, cve=None, cvss=None, subid=0, vuln_time=124, rowid=200, scan=self.scan_4, expiration_time=400) self.vulnerability_6 = Vulnerability(port=self.port_5, output='tftp', exploit=self.exploit_4, cve=None, cvss=None, subid=1, vuln_time=124, rowid=201, scan=self.scan_4, expiration_time=400) self.vulnerability_7 = Vulnerability(port=self.port_5, output='tftp server name', exploit=self.exploit_4, cve=None, cvss=None, subid=2, vuln_time=124, rowid=202, scan=self.scan_4, expiration_time=400) self.vulnerability_8 = Vulnerability(port=self.port_5, output='6.7.8', exploit=self.exploit_4, cve=None, cvss=None, subid=3, vuln_time=124, rowid=203, scan=self.scan_4, expiration_time=400) self.vulnerability_9 = Vulnerability(port=self.port_5, output='HERE IS TFTP\n\n\n > ', exploit=self.exploit_4, cve=None, cvss=None, subid=4, vuln_time=124, rowid=204, scan=self.scan_4, expiration_time=400) self.vulnerability_10 = Vulnerability(port=self.port_5, output='test:cpe', exploit=self.exploit_4, cve=None, cvss=None, subid=5, vuln_time=124, rowid=205, scan=self.scan_4, expiration_time=400) self.vulnerability_11 = Vulnerability(port=self.port_5, output='os name', exploit=self.exploit_4, cve=None, cvss=None, subid=6, vuln_time=124, rowid=206, scan=self.scan_4, expiration_time=400) self.vulnerability_12 = Vulnerability(port=self.port_5, output='os version', exploit=self.exploit_4, cve=None, cvss=None, subid=7, vuln_time=124, rowid=207, scan=self.scan_4, expiration_time=400) self.vulnerability_13 = Vulnerability(port=self.port_5, output='test:os:cpe', exploit=self.exploit_4, cve=None, cvss=None, subid=8, vuln_time=124, rowid=208, scan=self.scan_4, expiration_time=400) def prepare_tables(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() self.prepare_scans() self.prepare_nodes_scans() self.prepare_ports_scans() self.prepare_vulnerabilities() self.prepare_security_scans() def prepare_scans(self): self.storage.execute(( "INSERT INTO scans(ROWID, protocol, scanner_name, scan_start, scan_end) " "VALUES " "(56, 17, 'test_name', 13, 19), " "(80, 17, 'test_name_2', 20, 45), " "(79, 17, 'test_name', 2, 18)," "(78, 6, 'portdetection', 1, 2)", )) def prepare_nodes_scans(self): self.storage.execute(( "INSERT INTO nodes_scans (ROWID, scan_id, node_id, node_ip, time) VALUES " "(13, 56, 1, '127.0.0.1', 15), " "(15, 56, 2, '127.0.0.2', 56), " "(16, 56, 3, '127.0.0.3', 98), " "(17, 79, 3, '127.0.0.3', 90) " "", )) def prepare_ports_scans(self): self.storage.execute(( 'INSERT INTO ports_scans(ROWID, scan_id, node_id, node_ip, port, port_protocol, time) ' 'VALUES ' "(124, 56, 1, '127.0.0.1', 45, 17, 176), " "(480, 56, 1, '127.0.0.1', 80, 17, 650), " "(13, 56, 3, '127.0.0.3', 99, 1, 619), " "(15, 56, 2, '127.0.0.2', 65, 17, 987)", )) def prepare_security_scans(self): self.storage.execute(( "INSERT INTO security_scans(exploit_id, exploit_name, exploit_app, scan_id, node_id, " "node_ip, port_number, port_protocol, sec_scan_start, sec_scan_end) VALUES " "(14, 'test_name', 'test_app', 56, 1, '127.0.0.1', 45, 17, 178, 851), " "(56, 'test_name_2', 'test_app', 79, 1, '127.0.0.1', 45, 17, 113, 353), " "(14, 'test_name', 'test_app', 80, 1, '127.0.0.1', 45, 17, 180, 222), " "(2, 'test_name_2', 'test_app_2', 56, 1, '127.0.0.1', 45, 17, 109, 775), " "(14, 'test_name', 'test_app', 79, 1, '127.0.0.1', 45, 17, 14, 156), " "(2, 'test_name_2', 'test_app_2', 79, 1, '127.0.0.1', 45, 17, 56, 780), " "(0, 'portdetection', 'portdetection', 78, 4, '127.0.0.4', 22, 6, 14, 890)", )) def prepare_vulnerabilities(self): self.storage.execute(( "INSERT INTO vulnerabilities (ROWID, scan_id, node_id, node_ip, port_protocol, port, " "vulnerability_id, vulnerability_subid, cve, cvss, output, time, expiration_time) VALUES " "(134, 56, 1, '127.0.0.1', 17, 45, 14, 1, 'CVE-2017', 6.7, 'test_output_1', 13, 600)," "(152, 56, 1, '127.0.0.1', 17, 45, 2, 2, 'CWE-14', 8.9, 'test_output_2', 98, 600)," "(153, 79, 1, '127.0.0.1', 17, 45, 14, 2, 'CVE-2016', 3.7, 'test_output_3', 15, NULL)," "(169, 79, 1, '127.0.0.1', 17, 45, 2, 1, 'CWE-15', 2.9, 'test_output_4', 124, 600), " "(200, 78, 4, '127.0.0.4', 6, 22, 0, 0, NULL, NULL, '', 124, 400), " "(201, 78, 4, '127.0.0.4', 6, 22, 0, 1, NULL, NULL, 'tftp', 124, 400), " "(202, 78, 4, '127.0.0.4', 6, 22, 0, 2, NULL, NULL, 'tftp server name', 124, 400), " "(203, 78, 4, '127.0.0.4', 6, 22, 0, 3, NULL, NULL, '6.7.8', 124, 400), " "(204, 78, 4, '127.0.0.4', 6, 22, 0, 4, NULL, NULL, 'HERE IS TFTP\n\n\n > ', 124, 400), " "(205, 78, 4, '127.0.0.4', 6, 22, 0, 5, NULL, NULL, 'test:cpe', 124, 400), " "(206, 78, 4, '127.0.0.4', 6, 22, 0, 6, NULL, NULL, 'os name', 124, 400), " "(207, 78, 4, '127.0.0.4', 6, 22, 0, 7, NULL, NULL, 'os version', 124, 400), " "(208, 78, 4, '127.0.0.4', 6, 22, 0, 8, NULL, NULL, 'test:os:cpe', 124, 400)" "", )) def test_context_manager(self): with self.storage: self.assertTrue(True) def test_connect(self): self.storage.connect() self.assertIsInstance(self.storage.conn, connection) def test_close(self): self.storage.conn = connect(":memory:") self.storage.close() self.assertEqual(self.storage.conn, None) def test_cursor_property(self): self.assertEqual(self.storage.cursor, self.storage._cursor) @patch('utils.storage.time.time', MagicMock(return_value=134)) def test_save_nodes(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() expected = [(56, 1, '127.0.0.1', 134), (56, 2, '127.0.0.2', 134), (56, 3, '127.0.0.3', 134)] self.storage.save_nodes(nodes=[self.node_1, self.node_2, self.node_3], scan=self.scan_1) result = self.storage.execute(self.SEL_NOD_SCAN) self.assertEqual(result, expected) @patch('utils.storage.time.time', MagicMock(return_value=60)) def test_get_nodes(self): self.prepare_tables() expected = [self.node_2, self.node_3, self.node_3] result = self.storage.get_nodes(pasttime=10, scan=self.scan_1) self.assertCountEqual(result, expected) def test_get_vulnerabilities(self): self.prepare_tables() expected = [self.vulnerability_1] result = self.storage.get_vulnerabilities(port=self.port_1, scan=self.scan_1, exploit=self.exploit_1) self.assertCountEqual(result, expected) def test_get_nodes_by_scan(self): self.prepare_tables() expected = [self.node_1, self.node_2, self.node_3] result = self.storage.get_nodes_by_scan(scan=self.scan_1) self.assertCountEqual(result, expected) @patch('utils.storage.time.time', MagicMock(return_value=167)) def test_save_port(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() expected = [(56, 1, '127.0.0.1', 45, 17, 167)] self.storage.save_port(port=self.port_1, scan=self.scan_1) result = self.storage.execute(self.SEL_POR_SCAN) self.assertEqual(result, expected) @patch('utils.storage.time.time', MagicMock(return_value=167)) def test_save_ports(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() expected = [(56, 1, '127.0.0.1', 45, 17, 167), (56, 2, '127.0.0.2', 65, 17, 167), (56, 3, '127.0.0.3', 99, 1, 167)] self.storage.save_ports(ports=[self.port_1, self.port_2, self.port_3], scan=self.scan_1) result = self.storage.execute(self.SEL_POR_SCAN) self.assertEqual(result, expected) @patch('utils.storage.time.time', MagicMock(return_value=1000)) def test_get_ports(self): self.prepare_tables() expected = [ self.port_scan_2.port, self.port_scan_3.port, self.port_scan_4.port ] result = self.storage.get_ports(pasttime=400, scan=self.scan_1) self.assertCountEqual(result, expected) def test_get_ports_by_scan_and_node(self): self.prepare_tables() expected = [self.port_scan_1, self.port_scan_4] result = self.storage.get_ports_by_scan_and_node(node=self.node_1, scan=self.scan_1) self.assertCountEqual(result, expected) def test_save_security_scan(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() expected = [(56, 14, 'test_app', 'test_name', 1, '127.0.0.1', 17, 45, 13.0, 19.0)] self.storage.save_security_scan(exploit=self.exploit_1, port=self.port_1, scan=self.scan_1) result = self.storage.execute(self.SEL_SEC_SCAN) self.assertEqual(result, expected) def test_save_security_scan_twice(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() self.port_1.scan.start = None expected = [(56, 14, 'test_app', 'test_name', 1, '127.0.0.1', 17, 45, None, 19.0)] self.storage.save_security_scan(exploit=self.exploit_1, port=self.port_1, scan=self.scan_1) result = self.storage.execute(self.SEL_SEC_SCAN) self.assertEqual(result, expected) self.port_1.scan.start = 13 self.port_1.scan.end = None expected = [(56, 14, 'test_app', 'test_name', 1, '127.0.0.1', 17, 45, 13.0, 19.0)] self.storage.save_security_scan(exploit=self.exploit_1, port=self.port_1, scan=self.scan_1) result = self.storage.execute(self.SEL_SEC_SCAN) self.assertEqual(result, expected) def test_save_security_scans(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() expected = [(56, 14, 'test_app', 'test_name', 1, '127.0.0.1', 17, 45, 13.0, 19.0), (56, 2, 'test_app_2', 'test_name_2', 1, '127.0.0.1', 17, 45, 13.0, 19.0), (56, 56, 'test_app', 'test_name_2', 1, '127.0.0.1', 17, 45, 13.0, 19.0)] self.storage.save_security_scans( exploits=[self.exploit_1, self.exploit_2, self.exploit_3], port=self.port_1, scan=self.scan_1) result = self.storage.execute(self.SEL_SEC_SCAN) self.assertEqual(result, expected) def test_save_changes(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() expected = [(1, 0, 0, 15, 124, 124445), (1, 0, 0, 13, 15, 32434)] self.storage.save_changes( changes=[self.vuln_change_1, self.vuln_change_2]) result = self.storage.execute(self.SEL_CHANGE) self.assertEqual(result, expected) def test_get_ports_by_nodes_without_nodes(self): expected = [] result = self.storage.get_ports_by_nodes([]) self.assertEqual(result, expected) @patch('utils.storage.time.time', MagicMock(return_value=1000)) def test_get_ports_by_nodes(self): self.storage.nodes_limit = 2 self.prepare_tables() expected = [self.port_scan_2.port, self.port_scan_4.port] result = self.storage.get_ports_by_nodes( nodes=[self.node_1, self.node_3, self.node_2], pasttime=400, protocol=TransportProtocol.UDP) self.assertCountEqual(result, expected) def test_execute_query(self): query = "part_1", "arg_1", "arg_2" self.storage._cursor = MagicMock() self.storage._cursor.rowcount = 2 self.storage.conn = MagicMock() result = self.storage.execute(query) self.storage.cursor.execute.assert_called_once_with( "part_1", "arg_1", "arg_2") self.storage.cursor.fetchall.assert_called_once_with() self.assertEqual(result, self.storage.cursor.fetchall()) def test_execute_query_list(self): queries = [("part_1", "arg_1", "arg_2"), ("part_2", "arg_3")] self.storage._cursor = MagicMock() self.storage._cursor.rowcount = 2 self.storage.conn = MagicMock() self.storage.execute(queries) self.storage.cursor.execute.assert_has_calls( (call("part_1", "arg_1", "arg_2"), call("part_2", "arg_3"))) def test_transport_protocol_none(self): result = self.storage._transport_protocol(None) self.assertIsNone(result) @patch('utils.storage.TransportProtocol') def test_transport_protocol(self, protocol): result = self.storage._transport_protocol(6) protocol.from_iana.assert_called_once_with(6) self.assertEqual(result, protocol.from_iana()) def test_protocol_to_iana_none(self): protocol = None result = self.storage._protocol_to_iana(protocol) self.assertIsNone(result) def test_save_scan(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() expected = [(1, 13, 19, 17, 'test_name')] self.storage.save_scan(scan=self.scan_1) result = self.storage.execute(self.SEL_SCANS) self.assertEqual(result, expected) def test_update_scan(self): self.prepare_tables() self.scan_1.end = 456 self.storage.update_scan(self.scan_1) expected = [(456, )] result = self.storage.execute( ("SELECT scan_end FROM scans WHERE ROWID=56", )) self.assertEqual(result, expected) def test_get_scan_id(self): expected = 56 result = self.storage.get_scan_id(self.scan_1) self.assertEqual(result, expected) def test_get_scan_id(self): self.prepare_tables() expected = self.scan_1.rowid self.scan_1.rowid = None result = self.storage.get_scan_id(self.scan_1) self.assertEqual(result, expected) def test_get_scan_id_without_results(self): self.prepare_tables() expected = None self.scan_1.rowid = None self.scan_1._scanner = None result = self.storage.get_scan_id(self.scan_1) self.assertEqual(result, expected) def test_get_scans(self): self.prepare_tables() result = self.storage.get_scans(scanner_name='test_name', protocol=TransportProtocol.UDP, amount=2) self.assertEqual([self.scan_1, self.scan_2], result) for obj in self.scan_1, self.scan_2: self.assertEqual(result[result.index(obj)].rowid, obj.rowid) def test_get_scans_by_node(self): self.prepare_tables() expected = [self.scan_1, self.scan_2] result = self.storage.get_scans_by_node(node=self.node_3, scan=self.scan_1) self.assertEqual(result, expected) def test_get_scans_by_sec_scan(self): self.prepare_tables() expected = [self.scan_3, self.scan_1, self.scan_2] result = self.storage.get_scans_by_security_scan( port=self.port_1, exploit=self.exploit_1) self.assertEqual(result, expected) def test_get_scans_by_sec_scan_second(self): self.prepare_tables() expected = [self.scan_2] result = self.storage.get_scans_by_security_scan( port=self.port_1, exploit=self.exploit_3) self.assertEqual(result, expected) def test_get_scan_by_id(self): self.prepare_tables() expected = self.scan_1 result = self.storage.get_scan_by_id(56) self.assertEqual(result, expected) def test_get_scan_by_id_no_results(self): self.storage.connect() self.storage.init_schema() result = self.storage.get_scan_by_id(15) self.assertEqual(result, None) def test_save_vulnerabilities(self): self.storage.connect() self.storage.remove_all() self.storage.init_schema() expected = [(56, 1, '127.0.0.1', 17, 45, 14, 1, 'CVE-2017', '6.7', 'test_output_1', 13), (56, 1, '127.0.0.1', 17, 45, 14, 2, 'CVE-2016', '3.7', 'test_output_3', 15)] self.storage.save_vulnerabilities( vulnerabilities=[self.vulnerability_1, self.vulnerability_3], scan=self.scan_1) result = self.storage.execute(self.SEL_VULNS) self.assertCountEqual(result, expected) def test_scans(self): self.prepare_tables() self.assertEqual(self.storage.scans(2, 0), [self.scan_3, self.scan_1]) self.assertEqual(self.storage.scans(2, 1), [self.scan_2, self.scan_4]) def ports_scans_by_scan(self): self.prepare_tables() expected = [ self.port_scan_1, self.port_scan_2, self.port_scan_3, self.port_scan_4 ] result = self.storage.ports_scans_by_scan(self.scan_1) def test_expire_vulnerability(self): self.prepare_tables() result = self.storage.expire_vulnerability(self.vulnerability_3) self.assertEqual(result.expiration_time, 178) def test_active_vulnerabilities(self): self.prepare_tables() expected = [self.vulnerability_3] result = self.storage.active_vulnerabilities() self.assertEqual(result, expected) def test_expire_vulnerabilities(self): self.prepare_tables() expected = [] self.storage.expire_vulnerabilities() result = self.storage.active_vulnerabilities(timestamp=400) self.assertEqual(result, expected) def test_active_vulnerabilities_different_timestamps(self): self.prepare_tables() expected = [self.vulnerability_3] result = self.storage.active_vulnerabilities(timestamp=0) self.assertEqual(result, expected) expected = [] result = self.storage.active_vulnerabilities(timestamp=40) self.assertEqual(result, expected) def test_portdetection_vulns(self): self.prepare_tables() expected = { 'name': 'tftp server name', 'version': '6.7.8', 'banner': 'HERE IS TFTP\n\n\n > ', 'cpe': 'test:cpe', 'protocol': 'tftp', 'os_name': 'os name', 'os_version': 'os version', 'os_cpe': 'test:os:cpe' } result = self.storage.portdetection_vulns(self.vulnerability_5) self.assertEqual(result, expected)
class APITest(AsyncHTTPTestCase): def setUp(self, *args, **kwargs): self.maxDiff = None super(APITest, self).setUp(*args, **kwargs) def get_app(self): self.aucote = MagicMock() self.storage = Storage(getenv('AUCOTE_TEST_POSTGRES')) self.aucote.storage = self.storage self.storage.connect() self.storage.remove_all() self.storage.init_schema() self.scan_1 = Scan(start=123, end=446, protocol=TransportProtocol.TCP, scanner='tcp') self.scan_2 = Scan(start=230, end=447, protocol=TransportProtocol.UDP, scanner='udp') for scan in (self.scan_1, self.scan_2): self.storage.save_scan(scan) self.node_1 = Node(node_id=13, ip=ipaddress.ip_address("10.156.67.18")) self.node_2 = Node(node_id=75, ip=ipaddress.ip_address("10.156.67.34")) self.node_scan_1 = NodeScan(node=self.node_1, scan=self.scan_1, timestamp=45) self.node_scan_2 = NodeScan(node=self.node_2, scan=self.scan_2, timestamp=88) self.node_scan_3 = NodeScan(node=self.node_1, scan=self.scan_2, timestamp=67) for node_scan in (self.node_scan_1, self.node_scan_2, self.node_scan_3): self.storage.save_node_scan(node_scan) self.port_1 = Port(node=self.node_1, number=34, transport_protocol=TransportProtocol.UDP) self.port_2 = Port(node=self.node_2, number=78, transport_protocol=TransportProtocol.TCP) self.port_scan_1 = PortScan(port=self.port_1, timestamp=1234, scan=self.scan_1, rowid=13) self.port_scan_2 = PortScan(port=self.port_2, timestamp=2345, scan=self.scan_1, rowid=15) for port_scan in (self.port_scan_1, self.port_scan_2): self.storage.save_port_scan(port_scan) self.exploit_1 = Exploit(exploit_id=14, name='test_name', app='test_app') self.exploit_2 = Exploit(exploit_id=2, name='test_name_2', app='test_app_2') self.security_scan_1 = SecurityScan(exploit=self.exploit_1, port=self.port_1, scan=self.scan_1, scan_start=178, scan_end=851) self.security_scan_2 = SecurityScan(exploit=self.exploit_2, port=self.port_1, scan=self.scan_1, scan_start=109, scan_end=775) self.security_scan_3 = SecurityScan(exploit=self.exploit_1, port=self.port_1, scan=self.scan_2, scan_start=114, scan_end=981) for scan in (self.security_scan_1, self.security_scan_2, self.security_scan_3): self.storage.save_sec_scan(scan) self.vulnerability_1 = Vulnerability(exploit=self.exploit_1, port=self.port_1, cvss="6.8", cve="CVE-2017-1231", scan=self.scan_1, output="Vulnerable stuff", vuln_time=134, subid=34) self.vulnerability_2 = Vulnerability(exploit=self.exploit_1, port=self.port_1, cvss="6.8", cve="CVE-2017-1232", scan=self.scan_2, output=None, vuln_time=718, subid=34) for vulnerability in (self.vulnerability_1, self.vulnerability_2): self.storage.save_vulnerability(vulnerability) self.scanner = TCPScanner(aucote=self.aucote, host='localhost', port=1339) self.scanner.NAME = 'test_name' self.scanner.scan.start = 1290 self.scanner.nodes = [ Node(node_id=1, ip=ipaddress.ip_address('127.0.0.1')) ] self.aucote.scanners = [ self.scanner, ToolsScanner(name='tools', aucote=self.aucote) ] endpoints = [(url, handler, { 'aucote': self.aucote }) for url, handler in [ (r"/api/v1/kill", KillHandler), (r"/api/v1/scanners", ScannersHandler), (r"/api/v1/scanners/([\w_]+)", ScannersHandler), (r"/api/v1/tasks", TasksHandler), (r"/api/v1/scans", ScansHandler), (r"/api/v1/scans/([\d]+)", ScansHandler), (r"/api/v1/nodes", NodesHandler), (r"/api/v1/nodes/([\d]+)", NodesHandler), (r"/api/v1/ports", PortsHandler), (r"/api/v1/ports/([\d]+)", PortsHandler), (r"/api/v1/security_scans", SecurityScansHandler), (r"/api/v1/security_scans/([\d]+)", SecurityScansHandler), (r"/api/v1/vulnerabilities", VulnerabilitiesHandler), ]] endpoints.append( (r"/api/v1/vulnerabilities/([\d]+)", VulnerabilitiesHandler, { 'aucote': self.aucote, 'path': '/aucote/' })) self.app = Application(endpoints) return self.app def tearDown(self): self.storage.remove_all()