def setUp(self, get_logger) -> None:
from db.repositories.ProcessRepository import ProcessRepository
self.mockProcess = MagicMock()
self.mockDbSession = MagicMock()
self.mockDbAdapter = MagicMock()
self.mockLogger = MagicMock()
self.mockFilters = MagicMock()
self.mockDbAdapter.session.return_value = self.mockDbSession
self.processRepository = ProcessRepository(self.mockDbAdapter, self.mockLogger)
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)
class ProcessRepositoryTest(unittest.TestCase):
@patch('utilities.stenoLogging.get_logger')
def setUp(self, get_logger) -> None:
from db.repositories.ProcessRepository import ProcessRepository
self.mockProcess = MagicMock()
self.mockDbSession = MagicMock()
self.mockDbAdapter = MagicMock()
self.mockLogger = MagicMock()
self.mockFilters = MagicMock()
self.mockDbAdapter.session.return_value = self.mockDbSession
self.processRepository = ProcessRepository(self.mockDbAdapter,
self.mockLogger)
def test_getProcesses_WhenProvidedShowProcessesWithNoNmapFlag_ReturnsProcesses(
self):
expectedQuery = (
'SELECT "0", "0", "0", process.name, "0", "0", "0", "0", "0", "0", "0", "0", "0", "0", "0" '
'FROM process AS process WHERE process.closed = "False" AND process.name != "nmap" '
'GROUP BY process.name')
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[['some-process'], ['some-process2']])
processes = self.processRepository.getProcesses(self.mockFilters,
showProcesses='noNmap')
self.assertEqual(processes, [['some-process'], ['some-process2']])
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery)
def test_getProcesses_WhenProvidedShowProcessesWithFlagFalse_ReturnsProcesses(
self):
expectedQuery = (
'SELECT process.id, process.hostIp, process.tabTitle, process.outputfile, output.output '
'FROM process AS process INNER JOIN process_output AS output ON process.id = output.processId '
'WHERE process.display = ? AND process.closed = "False" order by process.id desc'
)
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[['some-process'], ['some-process2']])
processes = self.processRepository.getProcesses(self.mockFilters,
showProcesses=False)
self.assertEqual(processes, [['some-process'], ['some-process2']])
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, 'False')
def test_getProcesses_WhenProvidedShowProcessesWithNoFlag_ReturnsProcesses(
self):
expectedQuery = "SELECT * FROM process AS process WHERE process.display=? order by id asc"
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[['some-process'], ['some-process2']])
processes = self.processRepository.getProcesses(self.mockFilters,
"True",
sort='asc',
ncol='id')
self.assertEqual(processes, [['some-process'], ['some-process2']])
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, 'True')
def test_storeProcess_WhenProvidedAProcess_StoreProcess(self):
self.processRepository.storeProcess(self.mockProcess)
self.mockDbSession.add.assert_called_once()
self.mockDbAdapter.commit.assert_called_once()
def test_storeProcessOutput_WhenProvidedExistingProcessIdAndOutput_StoresProcessOutput(
self):
from db.entities.process import process
from db.entities.processOutput import process_output
expected_process: process = MagicMock()
process.status = 'Running'
expected_process_output: process_output = MagicMock()
mock_query = MagicMock()
mock_query.filter_by.return_value = mockFirstBySideEffect(
[expected_process, expected_process_output])
self.mockDbSession.query.return_value = mock_query
self.processRepository.storeProcessOutput("some_process_id",
"this is some cool output")
self.mockDbSession.add.assert_has_calls(
[mock.call(expected_process_output),
mock.call(expected_process)])
self.mockDbAdapter.commit.assert_called_once()
def test_storeProcessOutput_WhenProvidedProcessIdDoesNotExist_DoesNotPerformAnyUpdate(
self):
self.mockDbAdapter.session.return_value = self.mockDbSession
self.mockDbSession.query.return_value = mockQueryWithFilterBy(
mockFirstByReturnValue(False))
self.processRepository.storeProcessOutput(
"some_non_existent_process_id", "this is some cool output")
self.mockDbSession.add.assert_not_called()
self.mockDbAdapter.commit.assert_not_called()
def test_storeProcessOutput_WhenProvidedExistingProcessIdAndOutputButProcKilled_StoresOutputButStatusNotUpdated(
self):
self.whenProcessDoesNotFinishGracefully("Killed")
def test_storeProcessOutput_WhenProvidedExistingProcessIdAndOutputButProcCancelled_StoresOutputButStatusNotUpdated(
self):
self.whenProcessDoesNotFinishGracefully("Cancelled")
def test_storeProcessOutput_WhenProvidedExistingProcessIdAndOutputButProcCrashed_StoresOutputButStatusNotUpdated(
self):
self.whenProcessDoesNotFinishGracefully("Crashed")
def test_getStatusByProcessId_WhenGivenProcId_FetchesProcessStatus(self):
expectedQuery = 'SELECT process.status FROM process AS process WHERE process.id=?'
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[['Running']])
actual_status = self.processRepository.getStatusByProcessId(
"some_process_id")
self.assertEqual(actual_status, 'Running')
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some_process_id")
def test_getStatusByProcessId_WhenProcIdDoesNotExist_ReturnsNegativeOne(
self):
expectedQuery = 'SELECT process.status FROM process AS process WHERE process.id=?'
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
False)
actual_status = self.processRepository.getStatusByProcessId(
"some_process_id")
self.assertEqual(actual_status, -1)
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some_process_id")
def test_getPIDByProcessId_WhenGivenProcId_FetchesProcessId(self):
expectedQuery = 'SELECT process.pid FROM process AS process WHERE process.id=?'
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[['1234']])
actual_status = self.processRepository.getPIDByProcessId(
"some_process_id")
self.assertEqual(actual_status, '1234')
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some_process_id")
def test_getPIDByProcessId_WhenProcIdDoesNotExist_ReturnsNegativeOne(self):
expectedQuery = 'SELECT process.pid FROM process AS process WHERE process.id=?'
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
False)
actual_status = self.processRepository.getPIDByProcessId(
"some_process_id")
self.assertEqual(actual_status, -1)
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some_process_id")
def test_isKilledProcess_WhenProvidedKilledProcessId_ReturnsTrue(self):
expectedQuery = "SELECT process.status FROM process AS process WHERE process.id=?"
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[["Killed"]])
self.assertTrue(
self.processRepository.isKilledProcess("some_process_id"))
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some_process_id")
def test_isKilledProcess_WhenProvidedNonKilledProcessId_ReturnsFalse(self):
expectedQuery = "SELECT process.status FROM process AS process WHERE process.id=?"
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[["Running"]])
self.assertFalse(
self.processRepository.isKilledProcess("some_process_id"))
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some_process_id")
def test_isCancelledProcess_WhenProvidedCancelledProcessId_ReturnsTrue(
self):
expectedQuery = "SELECT process.status FROM process AS process WHERE process.id=?"
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[["Cancelled"]])
self.assertTrue(
self.processRepository.isCancelledProcess("some_process_id"))
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some_process_id")
def test_isCancelledProcess_WhenProvidedNonCancelledProcessId_ReturnsFalse(
self):
expectedQuery = "SELECT process.status FROM process AS process WHERE process.id=?"
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[["Running"]])
self.assertFalse(
self.processRepository.isCancelledProcess("some_process_id"))
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some_process_id")
def test_storeProcessCrashStatus_WhenProvidedProcessId_StoresProcessCrashStatus(
self):
self.mockProcessStatusAndReturnSingle("Running")
self.processRepository.storeProcessCrashStatus("some-process-id")
self.assertProcessStatusUpdatedTo("Crashed")
def test_storeProcessCancelledStatus_WhenProvidedProcessId_StoresProcessCancelledStatus(
self):
self.mockProcessStatusAndReturnSingle("Running")
self.processRepository.storeProcessCancelStatus("some-process-id")
self.assertProcessStatusUpdatedTo("Cancelled")
def test_storeProcessRunningStatus_WhenProvidedProcessId_StoresProcessRunningStatus(
self):
self.mockProcessStatusAndReturnSingle("Waiting")
self.processRepository.storeProcessRunningStatus(
"some-process-id", "3123")
self.assertProcessStatusUpdatedTo("Running")
def test_storeProcessKillStatus_WhenProvidedProcessId_StoresProcessKillStatus(
self):
self.mockProcessStatusAndReturnSingle("Running")
self.processRepository.storeProcessKillStatus("some-process-id")
self.assertProcessStatusUpdatedTo("Killed")
def test_storeProcessRunningElapsedTime_WhenProvidedProcessId_StoresProcessRunningElapsedTime(
self):
self.mockProcess.elapsed = "some-time"
self.mockDbSession.query.return_value = mockQueryWithFilterBy(
mockFirstByReturnValue(self.mockProcess))
self.processRepository.storeProcessRunningElapsedTime(
"some-process-id", "another-time")
self.assertEqual("another-time", self.mockProcess.elapsed)
self.mockDbSession.add.assert_called_once_with(self.mockProcess)
self.mockDbAdapter.commit.assert_called_once()
def test_getHostsByToolName_WhenProvidedToolNameAndClosedFalse_StoresProcessRunningElapsedTime(
self):
expectedQuery = (
'SELECT process.id, "0", "0", "0", "0", "0", "0", process.hostIp, process.port, '
'process.protocol, "0", "0", process.outputfile, "0", "0", "0" FROM process AS process '
'WHERE process.name=? and process.closed="False"')
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[["some-host1"], ["some-host2"]])
hosts = self.processRepository.getHostsByToolName(
"some-toolname", "False")
self.assertEqual([["some-host1"], ["some-host2"]], hosts)
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some-toolname")
def test_getHostsByToolName_WhenProvidedToolNameAndClosedAsFetchAll_StoresProcessRunningElapsedTime(
self):
expectedQuery = (
'SELECT "0", "0", "0", "0", "0", process.hostIp, process.port, process.protocol, "0", "0", '
'process.outputfile, "0", "0", "0" FROM process AS process WHERE process.name=?'
)
self.mockDbAdapter.metadata.bind.execute.return_value = mockExecuteFetchAll(
[["some-host1"], ["some-host2"]])
hosts = self.processRepository.getHostsByToolName(
"some-toolname", "FetchAll")
self.assertEqual([["some-host1"], ["some-host2"]], hosts)
self.mockDbAdapter.metadata.bind.execute.assert_called_once_with(
expectedQuery, "some-toolname")
def test_storeCloseStatus_WhenProvidedProcessId_StoresCloseStatus(self):
self.mockProcess.closed = 'False'
self.mockDbSession.query.return_value = mockQueryWithFilterBy(
mockFirstByReturnValue(self.mockProcess))
self.processRepository.storeCloseStatus("some-process-id")
self.assertEqual('True', self.mockProcess.closed)
self.mockDbSession.add.assert_called_once_with(self.mockProcess)
self.mockDbAdapter.commit.assert_called_once()
def test_storeScreenshot_WhenProvidedIPAndPortAndFileName_StoresScreenshot(
self):
processId = self.processRepository.storeScreenshot(
"some-ip", "some-port", "some-filename")
self.mockDbSession.add.assert_called_once()
self.mockDbSession.commit.assert_called_once()
def test_toggleProcessDisplayStatus_whenResetAllIsTrue_setDisplayToFalseForAllProcessesThatAreNotRunning(
self):
process1 = build_mock_process(status="Waiting", display="True")
process2 = build_mock_process(status="Waiting", display="True")
mock_query_response = MagicMock()
mock_filtered_response = MagicMock()
mock_filtered_response.all.return_value = [process1, process2]
mock_query_response.filter_by.return_value = mock_filtered_response
self.mockDbSession.query.return_value = mock_query_response
self.processRepository.toggleProcessDisplayStatus(resetAll=True)
self.assertEqual("False", process1.display)
self.assertEqual("False", process2.display)
self.mockDbSession.add.assert_has_calls([
mock.call(process1),
mock.call(process2),
])
self.mockDbAdapter.commit.assert_called_once()
def test_toggleProcessDisplayStatus_whenResetAllIFalse_setDisplayToFalseForAllProcessesThatAreNotRunningOrWaiting(
self):
process1 = build_mock_process(status="Random Status", display="True")
process2 = build_mock_process(status="Another Random Status",
display="True")
process3 = build_mock_process(status="Running", display="True")
mock_query_response = MagicMock()
mock_filtered_response = MagicMock()
mock_filtered_response.all.return_value = [process1, process2]
mock_query_response.filter_by.return_value = mock_filtered_response
self.mockDbSession.query.return_value = mock_query_response
self.processRepository.toggleProcessDisplayStatus()
self.assertEqual("False", process1.display)
self.assertEqual("False", process2.display)
self.assertEqual("True", process3.display)
self.mockDbSession.add.assert_has_calls([
mock.call(process1),
mock.call(process2),
])
self.mockDbAdapter.commit.assert_called_once()
def mockProcessStatusAndReturnSingle(self, processStatus: str):
self.mockProcess.status = processStatus
self.mockDbSession.query.return_value = mockQueryWithFilterBy(
mockFirstByReturnValue(self.mockProcess))
def assertProcessStatusUpdatedTo(self, expected_status: str):
self.assertEqual(expected_status, self.mockProcess.status)
self.mockDbSession.add.assert_called_once_with(self.mockProcess)
self.mockDbAdapter.commit.assert_called_once()
def whenProcessDoesNotFinishGracefully(self, process_status: str):
from db.entities.process import process
from db.entities.processOutput import process_output
expected_process: process = MagicMock()
expected_process.status = process_status
expected_process_output: process_output = MagicMock()
self.mockDbSession.query.return_value = mockQueryWithFilterBy(
mockFirstBySideEffect([expected_process, expected_process_output]))
self.processRepository.storeProcessOutput("some_process_id",
"this is some cool output")
self.mockDbSession.add.assert_called_once_with(expected_process_output)
self.mockDbAdapter.commit.assert_called_once()