def reinitialize(self, db: Database, hostRepository: HostRepository):
self.serviceRepository: ServiceRepository = ServiceRepository(db)
self.processRepository: ProcessRepository = ProcessRepository(db, log)
self.hostRepository: HostRepository = hostRepository
self.portRepository: PortRepository = PortRepository(db)
self.cveRepository: CVERepository = CVERepository(db)
self.noteRepository: NoteRepository = NoteRepository(db, log)
self.scriptRepository: ScriptRepository = ScriptRepository(db)
def buildRepositories(self, database: Database) -> RepositoryContainer:
hostRepository = HostRepository(database)
processRepository = ProcessRepository(database, self.logger)
serviceRepository = ServiceRepository(database)
portRepository: PortRepository = PortRepository(database)
cveRepository: CVERepository = CVERepository(database)
noteRepository: NoteRepository = NoteRepository(database, self.logger)
scriptRepository: ScriptRepository = ScriptRepository(database)
return RepositoryContainer(serviceRepository, processRepository,
hostRepository, portRepository,
cveRepository, noteRepository,
scriptRepository)
def setUp(self, get_logger) -> None:
from db.repositories.PortRepository import PortRepository
self.mockDbAdapter = MagicMock()
self.mockDbSession = MagicMock()
self.mockDbAdapter.session.return_value = self.mockDbSession
self.repository = PortRepository(self.mockDbAdapter)
class PortRepositoryTest(unittest.TestCase):
@patch('utilities.stenoLogging.get_logger')
def setUp(self, get_logger) -> None:
from db.repositories.PortRepository import PortRepository
self.mockDbAdapter = MagicMock()
self.mockDbSession = MagicMock()
self.mockDbAdapter.session.return_value = self.mockDbSession
self.repository = PortRepository(self.mockDbAdapter)
def test_getPortsByIPAndProtocol_ReturnsPorts(self):
expected_query = ("SELECT ports.portId FROM portObj AS ports "
"INNER JOIN hostObj AS hosts ON hosts.id = ports.hostId "
"WHERE hosts.ip = ? and ports.protocol = ?")
self.mockDbAdapter.metadata.bind.execute.return_value = mockFirstByReturnValue(
[['port-id1'], ['port-id2']])
ports = self.repository.getPortsByIPAndProtocol("some_host_ip", "tcp")
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(expected_query, "some_host_ip", "tcp")
self.assertEqual([['port-id1'], ['port-id2']], ports)
def test_getPortStatesByHostId_ReturnsPortsStates(self):
expected_query = 'SELECT port.state FROM portObj as port WHERE port.hostId = ?'
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[['port-state1'], ['port-state2']])
port_states = self.repository.getPortStatesByHostId("some_host_id")
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(expected_query, "some_host_id")
self.assertEqual([['port-state1'], ['port-state2']], port_states)
def test_getPortsAndServicesByHostIP_InvokedWithNoFilters_ReturnsPortsAndServices(self):
from app.auxiliary import Filters
expected_query = ("SELECT hosts.ip, ports.portId, ports.protocol, ports.state, ports.hostId, ports.serviceId, "
"services.name, services.product, services.version, services.extrainfo, services.fingerprint "
"FROM portObj AS ports INNER JOIN hostObj AS hosts ON hosts.id = ports.hostId "
"LEFT OUTER JOIN serviceObj AS services ON services.id = ports.serviceId "
"WHERE hosts.ip = ?")
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll([['ip1'], ['ip2']])
filters: Filters = Filters()
filters.apply(up=True, down=True, checked=True, portopen=True, portfiltered=True, portclosed=True,
tcp=True, udp=True)
results = self.repository.getPortsAndServicesByHostIP("some_host_ip", filters)
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(expected_query, "some_host_ip")
self.assertEqual([['ip1'], ['ip2']], results)
def test_getPortsAndServicesByHostIP_InvokedWithFewFilters_ReturnsPortsAndServices(self):
from app.auxiliary import Filters
expected_query = ("SELECT hosts.ip, ports.portId, ports.protocol, ports.state, ports.hostId, ports.serviceId, "
"services.name, services.product, services.version, services.extrainfo, services.fingerprint "
"FROM portObj AS ports INNER JOIN hostObj AS hosts ON hosts.id = ports.hostId "
"LEFT OUTER JOIN serviceObj AS services ON services.id = ports.serviceId "
"WHERE hosts.ip = ? AND ports.protocol != 'tcp' AND ports.protocol != 'udp'")
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll([['ip1'], ['ip2']])
filters: Filters = Filters()
filters.apply(up=True, down=True, checked=True, portopen=True, portfiltered=True, portclosed=True,
tcp=False, udp=False)
results = self.repository.getPortsAndServicesByHostIP("some_host_ip", filters)
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(expected_query, "some_host_ip")
self.assertEqual([['ip1'], ['ip2']], results)
def test_deleteAllPortsAndScriptsByHostId_WhenProvidedByHostIDAndProtocol_DeletesAllPortsAndScripts(self):
mockFilterHost = mockProtocolFilter = mockReturnAll = MagicMock()
mockPort1 = mockPort2 = MagicMock()
mockReturnAll.all.return_value = [mockPort1, mockPort2]
mockProtocolFilter.filter.return_value = mockReturnAll
mockFilterHost.filter.return_value = mockProtocolFilter
mockFilterScript = mockReturnAllScripts = MagicMock()
mockReturnAllScripts.all.return_value = ['some-script1', 'some-script2']
mockFilterScript.filter.return_value = mockReturnAllScripts
self.mockDbSession.query.side_effect = [mockFilterHost, mockFilterScript, mockFilterScript]
self.repository.deleteAllPortsAndScriptsByHostId("some-host-id", "some-protocol")
self.mockDbSession.delete.assert_has_calls([
mock.call('some-script1'), mock.call('some-script2'),
mock.call('some-script1'), mock.call('some-script2'),
mock.call(mockPort1), mock.call(mockPort2)
])
self.mockDbSession.commit.assert_called_once()