class ServiceRepositoryTest(unittest.TestCase):
def setUp(self) -> None:
from db.repositories.ServiceRepository import ServiceRepository
self.mockDbAdapter = MagicMock()
self.repository = ServiceRepository(self.mockDbAdapter)
def getServiceNamesTestCase(self, filters, expectedQuery):
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[{'name': 'service_name1'}, {'name': 'service_name2'}])
service_names = self.repository.getServiceNames(filters)
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(expectedQuery)
self.assertEqual([{'name': 'service_name1'}, {'name': 'service_name2'}], service_names)
def test_getServiceNames_InvokedWithNoFilters_FetchesAllServiceNames(self):
from app.auxiliary import Filters
expectedQuery = ("SELECT DISTINCT service.name FROM serviceObj as service "
"INNER JOIN portObj as ports "
"INNER JOIN hostObj AS hosts "
"ON hosts.id = ports.hostId AND service.id=ports.serviceId WHERE 1=1 "
"ORDER BY service.name ASC")
filters: Filters = Filters()
filters.apply(up=True, down=True, checked=True, portopen=True, portfiltered=True, portclosed=True,
tcp=True, udp=True)
self.getServiceNamesTestCase(filters=filters, expectedQuery=expectedQuery)
def test_getServiceNames_InvokedWithFewFilters_FetchesAllServiceNamesWithFiltersApplied(self):
from app.auxiliary import Filters
expectedQuery = ("SELECT DISTINCT service.name FROM serviceObj as service "
"INNER JOIN portObj as ports "
"INNER JOIN hostObj AS hosts "
"ON hosts.id = ports.hostId AND service.id=ports.serviceId WHERE 1=1 "
"AND hosts.status != 'down' AND ports.protocol != 'udp' "
"ORDER BY service.name ASC")
filters: Filters = Filters()
filters.apply(up=True, down=False, checked=True, portopen=True, portfiltered=True, portclosed=True,
tcp=True, udp=False)
self.getServiceNamesTestCase(filters=filters, expectedQuery=expectedQuery)
def test_getServiceNamesByHostIPAndPort_WhenProvidedWithHostIpAndPort_ReturnsServiceNames(self):
self.mockDbAdapter.metadata.bind.execute.return_value = mockFirstByReturnValue(
[['service-name1'], ['service-name2']])
expectedQuery = ("SELECT services.name FROM serviceObj AS services "
"INNER JOIN hostObj AS hosts ON hosts.id = ports.hostId "
"INNER JOIN portObj AS ports ON services.id=ports.serviceId "
"WHERE hosts.ip=? and ports.portId = ?")
result = self.repository.getServiceNamesByHostIPAndPort("some_host", "1234")
self.assertEqual([['service-name1'], ['service-name2']], result)
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(expectedQuery, "some_host", "1234")
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) -> None:
from db.repositories.ServiceRepository import ServiceRepository
self.mockDbAdapter = MagicMock()
self.repository = ServiceRepository(self.mockDbAdapter)