def setUp(self):
self._model_controller = controller.ModelController(mock())
self.cm = mock(CommandManager)
when(self.cm).saveCommand().thenReturn(True)
self._plugin_controller = PluginController("test", {}, self.cm)
class PluginTest(PluginBase):
def __init__(self):
PluginBase.__init__(self)
self.id = "Test"
self.name = "Test"
def parseOutputString(self, output, debug=False):
pass
self.workspace = mock(Workspace)
when(self.workspace).getContainee().thenReturn(ModelObjectContainer())
self._model_controller.setWorkspace(self.workspace)
self.plugin = PluginTest()
api.setUpAPIs(self._model_controller)
self._plugin_controller.setActivePlugin(self.plugin)
self.cmdinfo = CommandRunInformation(
**{
'workspace': 'test',
'itime': time(),
'command': 'test',
'params': 'test'
})
def genericEdit(self, obj, params, callback, process_pending=False):
mappersManager = self.createMapperMock()
dataMapper = mock()
objId = obj.getID()
when(mappersManager).getMapper(obj.class_signature).thenReturn(dataMapper)
when(dataMapper).save(obj).thenReturn(True)
when(mappersManager).find(objId).thenReturn(obj)
when(mappersManager).save(obj).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
callback(model_controller, obj, *params)
if process_pending:
model_controller.processAllPendingActions()
def testDeleteHostObjectDispatchRemoveSYNC(self):
host = Host("coquito")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).find(host.getID()).thenReturn(host)
when(mappersManager).remove(host.getID()).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.delHostSYNC(host.getID())
verify(mappersManager).remove(host.getID())
verify(mappersManager).find(host.getID())
def testAddHostGetsMapperDispatchSaveSYNC(self):
host = Host('coco')
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(host.class_signature).thenReturn(objectMapper)
when(objectMapper).save(host).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addHostSYNC(host)
verify(mappersManager).getMapper(host.class_signature)
verify(objectMapper).save(host)
def testAddInterfaceGetsMapperDispatchSaveSYNC(self):
host = Host('coco')
interface = Interface("int_mock0")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(interface.class_signature).thenReturn(objectMapper)
when(objectMapper).save(interface).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addInterfaceSYNC(host.getID(), interface)
verify(mappersManager).getMapper(interface.class_signature)
verify(objectMapper).save(interface)
def testGetAllHosts(self):
hosts = [ Host("coquito%i" % i ) for i in range(10)]
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(Host.__name__).thenReturn(objectMapper)
when(objectMapper).getAll().thenReturn(hosts)
model_controller = controller.ModelController(mock(), mappersManager)
hosts_obt = model_controller.getAllHosts()
verify(objectMapper).getAll()
verify(mappersManager).getMapper(Host.__name__)
self.assertListEqual(hosts, hosts_obt)
def testAddServiceGetsMapperDispatchSaveSYNC(self):
interface = Interface("int_mock0")
service = Service("servi")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(service.class_signature).thenReturn(objectMapper)
when(objectMapper).save(service).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addServiceToInterfaceSYNC(None, interface.getID(), service)
verify(mappersManager).getMapper(service.class_signature)
verify(objectMapper).save(service)
def testAddVulnToInterfaceGetsMapperDispatchSaveSYNC(self):
interface = Interface("int0")
vuln = ModelObjectVuln("a_vuln")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(vuln.class_signature).thenReturn(objectMapper)
when(objectMapper).save(vuln).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addVulnToInterfaceSYNC(None, interface.getID(), vuln)
verify(mappersManager).getMapper(vuln.class_signature)
verify(objectMapper).save(vuln)
def testAddVulnToHostGetsMapperDispatchSaveSYNC(self):
host = Host("pepito")
vuln = ModelObjectVuln("a_vuln")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(vuln.class_signature).thenReturn(objectMapper)
when(objectMapper).save(vuln).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addVulnToHostSYNC(host.getID(), vuln)
verify(mappersManager).getMapper(vuln.class_signature)
verify(objectMapper).save(vuln)
def testAddNoteToModelObjectSYNC(self):
host = Host("pepito")
note = ModelObjectNote("a_note")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(note.class_signature).thenReturn(objectMapper)
when(objectMapper).save(note).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addNoteSYNC(host.getID(), note)
verify(mappersManager).getMapper(note.class_signature)
verify(objectMapper).save(note)
def testAddCredToServiceGetsMapperDispatchSaveSYNC(self):
service = Service("pepito")
cred = ModelObjectCred("usr", "pass")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(cred.class_signature).thenReturn(objectMapper)
when(objectMapper).save(cred).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addCredToServiceSYNC(None, cred.getID(), cred)
verify(mappersManager).getMapper(cred.class_signature)
verify(objectMapper).save(cred)
def testAddNoteToInterfaceGetsMapperDispatchSaveSYNC(self):
interface = Interface("pepito")
note = ModelObjectNote("a_note")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(note.class_signature).thenReturn(objectMapper)
when(objectMapper).save(note).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addNoteToInterfaceSYNC(None, interface.getID(), note)
verify(mappersManager).getMapper(note.class_signature)
verify(objectMapper).save(note)
def testAddNoteToVulnGetsMapperDispatchSave(self):
vuln = ModelObjectVuln('a vuln')
note = ModelObjectNote("a_note")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(note.class_signature).thenReturn(objectMapper)
when(objectMapper).save(note).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addNoteToServiceSYNC(None, vuln.getID(), note)
verify(mappersManager).getMapper(note.class_signature)
verify(objectMapper).save(note)
def testGetHost(self):
host = Host("coquito")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(host.__class__.__name__).thenReturn(objectMapper)
when(objectMapper).find(host.getName()).thenReturn(host)
model_controller = controller.ModelController(mock(), mappersManager)
host_obt = model_controller.getHost('coquito')
verify(objectMapper).find(host.getName())
verify(mappersManager).getMapper(host.__class__.__name__)
self.assertEqual(host, host_obt)
def testDelRemovesObjectFromTrie(self):
host = Host("coquito")
mappersManager = self.createMapperMock()
objectMapper = mock()
triemock = mock()
when(mappersManager).getMapper(host.class_signature).thenReturn(objectMapper)
when(mappersManager).find(host.getID()).thenReturn(host)
when(triemock).addWord(host.getName()).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.treeWordsTries = triemock
model_controller.delHostSYNC(host.getID())
verify(mappersManager).remove(host.getID())
verify(triemock).removeWord(host.getName())
def testAddVulnToServiceGetsMapperDispatchSaveASYNC(self):
service = Service("servi")
vuln = ModelObjectVuln("a_vuln")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(vuln.class_signature).thenReturn(objectMapper)
when(objectMapper).save(vuln).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addVulnToServiceASYNC(None, service.getID(), vuln)
model_controller.processAllPendingActions()
verify(mappersManager).getMapper(vuln.class_signature)
verify(objectMapper).save(vuln)
def testAddNoteToNoteGetsMapperDispatchSaveASYNC(self):
host = Host("pepito")
note = ModelObjectNote("a_note")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(note.class_signature).thenReturn(objectMapper)
when(objectMapper).save(note).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addNoteToNoteASYNC(None, None, note.getID(), note)
model_controller.processAllPendingActions()
verify(mappersManager).getMapper(note.class_signature)
verify(objectMapper).save(note)
def testAddObjectSavesChildInParent(self):
host = Host('coco')
interface = Interface("int_mock0")
mappersManager = self.createMapperMock()
objectMapper = mock()
when(mappersManager).getMapper(interface.class_signature).thenReturn(objectMapper)
when(objectMapper).save(interface).thenReturn(True)
when(mappersManager).find(host.getID()).thenReturn(host)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.addInterfaceSYNC(host.getID(), interface)
verify(mappersManager).getMapper(interface.class_signature)
verify(objectMapper).save(interface)
self.assertEquals(interface, host.findChild(interface.getID()),
"Orphan child, what happen papi?")
def testAddSavesObjectNameInTrie(self):
host = Host('coco')
mappersManager = self.createMapperMock()
objectMapper = mock()
triemock = mock()
when(mappersManager).getMapper(host.class_signature).thenReturn(objectMapper)
when(objectMapper).save(host).thenReturn(True)
when(triemock).addWord(host.getName()).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.treeWordsTries = triemock
model_controller.addHostSYNC(host)
verify(mappersManager).getMapper(host.class_signature)
verify(objectMapper).save(host)
verify(triemock).addWord(host.getName())
def testEditHostSyncGetsMapperDispatchedSYNC(self):
host = Host("coquito")
mappersManager = self.createMapperMock()
dataMapper = mock()
objectMapper = mock()
triemock = mock()
when(mappersManager).getMapper(host.class_signature).thenReturn(dataMapper)
when(dataMapper).save(host).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.editHostSYNC(host, 'new_name', 'new_desc', 'new_os', True)
verify(dataMapper).save(host)
self.assertEquals(host.getName(), 'new_name', "Name not updated")
self.assertEquals(host.getDescription(), 'new_desc', "Description not updated")
self.assertEquals(host.getOS(), 'new_os', "OS not updated")
self.assertEquals(host.isOwned(), True, "Owned status not updated")
def genericDelTest(self, obj1, obj2, test_method, process_pending=False):
mappersManager = self.createMapperMock()
objectMapper = mock()
triemock = mock()
when(mappersManager).find(obj2.getID()).thenReturn(obj2)
when(objectMapper).delObject(obj2.getID()).thenReturn(True)
model_controller = controller.ModelController(mock(), mappersManager)
model_controller.treeWordsTries = triemock
try:
test_method(model_controller, None, obj2.getID())
except:
test_method(model_controller, None, None, obj2.getID())
if process_pending:
model_controller.processAllPendingActions()
verify(mappersManager).find(obj2.getID())
verify(mappersManager).remove(obj2.getID())
