Python model_managers.WorkspaceManager类代码示例

    def testOpenWorkspaceChangesAndReportManagerWatch(self):
        reportManager = mock()

        dbManager = mock()
        mappersManager = mock()
        dbConnector = mock()
        mappers = mock()
        changesController = mock()
        workspaceMapper = mock()

        workspace = Workspace('test_workspace', 'a desc')

        when(dbManager).getAllDbNames().thenReturn(['test_workspace'])
        when(dbManager).getConnector('test_workspace').thenReturn(dbConnector)
        when(mappersManager).createMappers(dbConnector).thenReturn(True)
        when(mappersManager).getMapper(Workspace.__name__).thenReturn(workspaceMapper)
        when(workspaceMapper).find('test_workspace').thenReturn(workspace)

        workspace_manager = WorkspaceManager(dbManager,
                                mappersManager,
                                changesController,
                                reportManager)


        opened_workspace = workspace_manager.openWorkspace('test_workspace')

        verify(reportManager).watch('test_workspace')
        verify(changesController).watch(mappersManager, dbConnector)
        self.assertEquals(opened_workspace.getName(), 'test_workspace')

    def testCreateWorkspaceDBManagerInteract(self):
        dbManager = mock()
        dbConnector = mock()
        changesController = mock()

        when(dbManager).createDb('test_workspace', DBTYPE.FS).thenReturn(dbConnector)
        workspace_manager = WorkspaceManager(dbManager, mock(), changesController, mock())
        workspace_manager.createWorkspace('test_workspace', 'a test workspace',
                                        DBTYPE.FS)
        verify(dbManager).createDb('test_workspace', DBTYPE.FS)

    def testGetWorkspaceTypeFS(self):
        work = Workspace('testname')
        dbManager = mock()
        mappersManager = mock()
        changesController = mock()
        when(dbManager).getDbType('testname').thenReturn(DBTYPE.FS)
        workspace_manager = WorkspaceManager(dbManager, mappersManager, changesController, mock())

        wtype = workspace_manager.getWorkspaceType(work.getName())
        self.assertEquals(wtype, 'FS', 'Workspace type not returning correct value')

    def testGetAvailableWorkspaceTypes(self): 
        dbManager = mock()
        workspace_manager = WorkspaceManager(dbManager,
                                                mock(),
                                                mock(),
                                                mock())
        when(dbManager).getAvailableDBs().thenReturn([DBTYPE.COUCHDB, DBTYPE.FS])
        retrievedTypes = workspace_manager.getAvailableWorkspaceTypes()

        self.assertListEqual(['CouchDB', 'FS'], retrievedTypes, 
                                    "Workspaces available Types not set")

    def testSetActiveWorkspace(self):
        work = Workspace('testname')
        dbManager = mock()
        mappersManager = mock()
        changesController = mock()
        workspace_manager = WorkspaceManager(dbManager, mappersManager, changesController, mock())

        workspace_manager.setActiveWorkspace(work)

        self.assertEquals(workspace_manager.active_workspace, work,
                'active workspace not set')
        self.assertTrue(workspace_manager.isActive(work.getName()),
                'could not retrive as active workspace')

    def testCloseWorkspace(self):
        dbManager = mock()
        mappersManager = mock()
        changesController = mock()
        reportManager = mock()


        workspace_manager = WorkspaceManager(dbManager,
                                                mappersManager,
                                                changesController,
                                                reportManager)

        workspace_manager.closeWorkspace()
        verify(changesController).unwatch()

    def testResourceManager(self):
        dbManager = mock()
        mappersManager = mock()
        changesController = mock()
        reportManager = mock()


        workspace_manager = WorkspaceManager(dbManager,
                                                mappersManager,
                                                changesController,
                                                reportManager)

        workspace_manager.resource()

        verify(dbManager).reloadConfig()

    def testRemoveWorkspace(self):
        dbManager = mock()
        mappersManager = mock()
        dbConnector = mock()
        mappers = mock()
        changesController = mock()

        workspace = Workspace('test_workspace', 'a desc')
        when(dbManager).removeDb('test_workspace').thenReturn(True)
        when(dbManager).getAllDbNames().thenReturn(['test_workspace'])

        workspace_manager = WorkspaceManager(dbManager, mappersManager, changesController, mock())
        remove_ret = workspace_manager.removeWorkspace('test_workspace')

        verify(dbManager).removeDb('test_workspace')
        self.assertTrue(remove_ret, 'bbdd not removed')

    def testCreateWorkspaceSetsChangesCallback(self):
        dbManager = mock()
        mappersManager = mock()
        dbConnector = mock()
        mappers = mock()
        changesController = mock()

        when(mappersManager).save(any()).thenReturn(True) 
        when(dbManager).createDb('test_workspace', DBTYPE.FS).thenReturn(dbConnector)
        when(mappersManager).createMappers(dbConnector).thenReturn(True) 

        workspace_manager = WorkspaceManager(dbManager, mappersManager, changesController, mock())
        workspace = workspace_manager.createWorkspace('test_workspace', 'a test workspace',
                                        DBTYPE.FS)

        verify(changesController).watch(mappersManager, dbConnector)

    def __init__(self, args):
        self._original_excepthook = sys.excepthook

        self._configuration = CONF

        self._security_manager = SecurityManager()
        self._mappers_manager = MapperManager()
        self._changes_controller = ChangeController()
        self._db_manager = DbManager()

        self._model_controller = ModelController(self._security_manager, self._mappers_manager)

        self._plugin_manager = PluginManager(os.path.join(CONF.getConfigPath(), "plugins"), self._mappers_manager)

        self._reports_manager = ReportManager(10, self._plugin_manager.createController("ReportManager"))

        self._workspace_manager = WorkspaceManager(
            self._db_manager, self._mappers_manager, self._changes_controller, self._reports_manager
        )

        self.gui_app = UiFactory.create(self._model_controller, self._plugin_manager, self._workspace_manager, args.gui)

        self.gui_app.setSplashImage(os.path.join(CONF.getImagePath(), "splash2.png"))

        self.timer = TimerClass()
        self.timer.start()

    def testOpenWorkspaceNoneExisting(self):
        dbManager = mock()
        mappersManager = mock()
        dbConnector = mock()
        mappers = mock()
        changesController = mock()

        workspace = Workspace('test_workspace', 'a desc') 
        when(dbManager).getAllDbNames().thenReturn([])

        workspace_manager = WorkspaceManager(dbManager, mappersManager, changesController, mock())
        opened_workspace = workspace_manager.openWorkspace('test_workspace')


        verify(mappersManager, times=0).createMappers(dbConnector)
        verify(mappersManager, times=0).find('test_workspace')
        self.assertFalse(opened_workspace, 'Workspace retrieved but non existing')

    def testCreateExistingWorkspaceReturnsFalse(self):
        dbManager = mock()
        mappersManager = mock()
        dbConnector = mock()
        mappers = mock()
        changesController = mock()

        when(mappersManager).save(any()).thenReturn(True) 
        when(dbManager).createDb('test_workspace', DBTYPE.FS).thenReturn(False)
        when(mappersManager).createMappers(dbConnector).thenReturn(True) 

        workspace_manager = WorkspaceManager(dbManager, mappersManager, changesController, mock())
        workspace = workspace_manager.createWorkspace('test_workspace', 'a test workspace',
                                        DBTYPE.FS)

        verify(dbManager).createDb('test_workspace', DBTYPE.FS)
        verify(mappersManager, times=0).createMappers(dbConnector)
        verify(mappersManager, times=0).save(any())

    def testOpenWorkspaceSetsChangesCallback(self):
        dbManager = mock()
        mappersManager = mock()
        dbConnector = mock()
        mappers = mock()
        changesController = mock()
        workspaceMapper = mock()

        workspace = Workspace('test_workspace', 'a desc')

        when(dbManager).getConnector('test_workspace').thenReturn(dbConnector)
        when(mappersManager).getMapper(Workspace.__name__).thenReturn(workspaceMapper)
        when(dbManager).getAllDbNames().thenReturn(['test_workspace'])
        when(mappersManager).createMappers(dbConnector).thenReturn(True)
        when(workspaceMapper).find('test_workspace').thenReturn(workspace)

        workspace_manager = WorkspaceManager(dbManager, mappersManager, changesController, mock())

        opened_workspace = workspace_manager.openWorkspace('test_workspace')

        verify(changesController).watch(mappersManager, dbConnector)

class PluginsToModelControllerIntegration(unittest.TestCase):

    def setUp(self):
        """
        Generic test to verify that the object exists and can be
        instantiated without problems.
        """
        self.dbManager = mock()
        self.changesController = mock()
        self.reportManager = mock()

        self.dbManager = DbManager()
        self.mappersManager = MapperManager()

        self.model_controller = controller.ModelController(mock(), self.mappersManager)
        self.workspace_manager = WorkspaceManager(self.dbManager,
                                             self.mappersManager,
                                             self.changesController,
                                             self.reportManager)
        self.workspace_manager.createWorkspace('temp_workspace', 'desc', DBTYPE.FS)
        self.workspace_manager.openWorkspace('temp_workspace')

        self._plugin_controller = PluginControllerForApi("test", {"nmap": nmap_plugin.NmapPlugin(),
                                                                    "nessus": nessus_plugin.NessusPlugin()}, mock())

        api.setUpAPIs(self.model_controller, self.workspace_manager)

    def tearDown(self): 
        self.workspace_manager.removeWorkspace('temp_workspace')

    def test_nmap_scan_saves_host(self):
        output_file = open(os.path.join(os.getcwd(), 'test_cases/data/nmap_plugin_with_api.xml'))
        output = output_file.read()
        self._plugin_controller.processCommandInput("nmap localhost")
        self._plugin_controller.onCommandFinished("nmap localhost", output)
        self.model_controller.processAllPendingActions()
        self.assertEquals(len(self.model_controller.getAllHosts()), 1,
                "Not all hosts added to model")

        host = self.model_controller.getAllHosts()[0]
        self.assertEquals(len(host.getAllInterfaces()), 1,
            "Not all interfaces added to model")

        interface = host.getAllInterfaces()[0]
        self.assertEquals(len(interface.getAllServices()), 3,
            "Not all services added to model")

        services = interface.getAllServices()
        self.assertTrue(all( [ s.getStatus() == 'open' for s in services]),
                "Port status not saved correctly")


    def test_nessus_scan_saves_host(self):
        output_file = open(os.path.join(os.getcwd(), 'test_cases/data/nessus_plugin_with_api.nessus'))
        output = output_file.read() 
        self._plugin_controller.processCommandInput("./nessus report")
        self._plugin_controller.onCommandFinished("./nessus report", output)
        self.model_controller.processAllPendingActions()
        self.assertEquals(len(self.model_controller.getAllHosts()), 7,
                "Not all hosts added to model")

    def testCreateWorkspaceCreateMappersAndWorkspace(self):
        dbManager = mock()
        mappersManager = mock()
        dbConnector = mock()
        mappers = mock()
        changesController = mock()
        workspaceMapper = mock()

        when(mappersManager).getMapper(Workspace.__name__).thenReturn(workspaceMapper)
        when(mappersManager).save(any()).thenReturn(True) 
        when(dbManager).createDb('test_workspace', DBTYPE.FS).thenReturn(dbConnector)
        when(mappersManager).createMappers(dbConnector).thenReturn(True) 

        workspace_manager = WorkspaceManager(dbManager, mappersManager, changesController, mock())
        workspace = workspace_manager.createWorkspace('test_workspace', 'a test workspace',
                                        DBTYPE.FS)

        verify(mappersManager).createMappers(dbConnector)
        verify(mappersManager).save(any())

        self.assertTrue(workspace, 'workspace not instantiated')
        self.assertEquals(workspace.name, 'test_workspace',
                            'Workspace name not set, is it valid?')