class TestActionManager(unittest.TestCase): def setUp(self): output = StringIO() self.logger = Logger.get('test', output, OutputStream) store = MongoStore.get_default() self.collection = store.get_collection(name='default_test') self.mongo_collection = MongoCollection(collection=self.collection, logger=self.logger) # Cleanup self.tearDown() self.manager = ActionManager(logger=self.logger, mongo_collection=self.mongo_collection) self.id_ = 'testid' self.action = { "_id": self.id_, "type": "pbehavior", "fields": ["Resource"], "regex": ".*wine.*", "parameters": { "author": "Matho", "name": "Salammbo", "reason": "Madness", "type": "Mercenary War", "rrule": "" } } def tearDown(self): """Teardown""" self.mongo_collection.remove({}) def test_crud(self): res = self.manager.create(action=self.action) self.assertTrue(res) res = self.manager.get_id(self.id_) self.assertIsNotNone(res) self.assertDictEqual(res.to_dict(), self.action) action2 = self.action.copy() action2[Action.FIELDS] = ['Component'] res = self.manager.update_id(id_=self.id_, action=action2) self.assertTrue(res) res = self.manager.get_id(self.id_) self.assertIsNotNone(res) self.assertDictEqual(res.to_dict(), action2) res = self.manager.delete_id(id_=self.id_) self.assertTrue(res) res = self.manager.get_id(self.id_) self.assertIsNone(res)
class SessionManagerTest(TestCase): def setUp(self): self.storage = Middleware.get_middleware_by_uri( 'mongodb-default-testsession://' ) self.collection = MongoCollection(self.storage._backend) self.manager = Session(collection=self.collection) self.user = '******' def tearDown(self): self.collection.remove() def test_keep_alive(self): self.manager.session_start(self.user) sleep(1) got = self.manager.keep_alive(self.user) session = self.collection.find_one({'_id': self.user}) self.assertTrue(isinstance(session, dict)) self.assertEqual(got, session['last_check']) def test_session_start(self): got = self.manager.session_start(self.user) session = self.collection.find_one({'_id': self.user}) self.assertTrue(isinstance(session, dict)) self.assertTrue(session['active']) self.assertEqual(got, session['session_start']) def test_session_start_already_started(self): self.test_session_start() got = self.manager.session_start(self.user) self.assertTrue(got is None) def test_is_session_active(self): self.assertFalse(self.manager.is_session_active(self.user)) self.manager.session_start(self.user) self.assertTrue(self.manager.is_session_active(self.user)) def test_sessions_close(self): got = self.manager.session_start(self.user) self.manager.alive_session_duration = 0 self.assertTrue(got is not None) sessions = self.manager.sessions_close() self.assertTrue(len(sessions) > 0) self.assertEqual(got, sessions[0]['last_check'])
class SessionManagerTest(TestCase): def setUp(self): self.storage = Middleware.get_middleware_by_uri( 'mongodb-default-testsession://') self.collection = MongoCollection(self.storage._backend) self.manager = Session(collection=self.collection) self.user = '******' def tearDown(self): self.collection.remove() def test_keep_alive(self): self.manager.session_start(self.user) sleep(1) got = self.manager.keep_alive(self.user) session = self.collection.find_one({'_id': self.user}) self.assertTrue(isinstance(session, dict)) self.assertEqual(got, session['last_check']) def test_session_start(self): got = self.manager.session_start(self.user) session = self.collection.find_one({'_id': self.user}) self.assertTrue(isinstance(session, dict)) self.assertTrue(session['active']) self.assertEqual(got, session['session_start']) def test_session_start_already_started(self): self.test_session_start() got = self.manager.session_start(self.user) self.assertTrue(got is None) def test_is_session_active(self): self.assertFalse(self.manager.is_session_active(self.user)) self.manager.session_start(self.user) self.assertTrue(self.manager.is_session_active(self.user)) def test_sessions_close(self): got = self.manager.session_start(self.user) self.manager.alive_session_duration = 0 self.assertTrue(got is not None) sessions = self.manager.sessions_close() self.assertTrue(len(sessions) > 0) self.assertEqual(got, sessions[0]['last_check'])
class TestMongoCollection(unittest.TestCase): def setUp(self): output = StringIO() self.logger = Logger.get('test', output, OutputStream) self.storage = Middleware.get_middleware_by_uri( 'storage-default-testmongocollection://') self.collection = MongoCollection(collection=self.storage._backend, logger=self.logger) self.id_ = 'testid' def tearDown(self): """Teardown""" self.collection.remove() def test_insert(self): res = self.collection.insert(document={'_id': self.id_}) self.assertEqual(res, self.id_) res2 = self.collection.insert(document={'up': 'down'}) self.assertTrue(isinstance(res, string_types)) self.assertNotEqual(res, res2) def test_update(self): res = self.collection.update(query={'_id': self.id_}, document={'strange': 'charm'}) self.assertTrue(MongoCollection.is_successfull(res)) self.assertEqual(res['n'], 0) res = self.collection.update(query={'_id': self.id_}, document={'yin': 'yang'}, upsert=True) self.assertTrue(MongoCollection.is_successfull(res)) self.assertEqual(res['n'], 1) res = self.collection.find_one(self.id_) self.assertEqual(res['yin'], 'yang') self.assertTrue('strange' not in res) def test_remove(self): res = self.collection.insert(document={ '_id': self.id_, 'top': 'bottom' }) self.assertIsNotNone(res) res = self.collection.remove(query={'_id': self.id_}) self.assertTrue(MongoCollection.is_successfull(res)) self.assertEqual(res['n'], 1) # Deleting non-existing object doesn't throw error res = self.collection.remove(query={}) self.assertTrue(MongoCollection.is_successfull(res)) self.assertEqual(res['n'], 0) def test_find(self): res = self.collection.insert(document={'_id': self.id_, 'yin': 'yang'}) self.assertIsNotNone(res) res = self.collection.find(query={'_id': self.id_}) self.assertTrue(isinstance(res, Cursor)) res = list(res) self.assertTrue(isinstance(res, list)) self.assertEqual(res[0]['yin'], 'yang') def test_find_one(self): res = self.collection.insert(document={'_id': self.id_, 'up': 'down'}) self.assertIsNotNone(res) res = self.collection.insert(document={'strange': 'charm'}) self.assertIsNotNone(res) res = self.collection.find_one(query={'_id': self.id_}) self.assertTrue(isinstance(res, dict)) self.assertEqual(res['up'], 'down') def test_is_successfull(self): dico = {'ok': 1.0, 'n': 2} self.assertTrue(MongoCollection.is_successfull(dico)) dico = {'ok': 666.667, 'n': 1} self.assertFalse(MongoCollection.is_successfull(dico)) dico = {'n': 2} self.assertFalse(MongoCollection.is_successfull(dico))
class RuleManager(object): """ Manager for event filter rules. """ def __init__(self, logger): self.logger = logger self.rule_collection = MongoCollection( MongoStore.get_default().get_collection(RULE_COLLECTION)) def get_by_id(self, rule_id): """ Get an event filter rule given its id. :param str rule_id: the id of the rule. :rtype: Dict[str, Any] """ return self.rule_collection.find_one({ RuleField.id: rule_id }) def create(self, rule): """ Create a new rule and return its id. :param Dict[str, Any] rule: :rtype: str :raises: InvalidRuleError if the rule is invalid. CollectionError if the creation fails. """ rule_id = str(uuid4()) rule[RuleField.id] = rule_id self.validate(rule_id, rule) self.rule_collection.insert(rule) return rule_id def remove_with_id(self, rule_id): """ Remove a rule given its id. :param str rule_id: the id of the rule. CollectionError if the creation fails. """ self.rule_collection.remove({ RuleField.id: rule_id }) def list(self): """ Return a list of all the rules. :rtype: List[Dict[str, Any]] """ return list(self.rule_collection.find({})) def update(self, rule_id, rule): """ Update a rule given its id. :param str rule_id: the id of the rule. :param Dict[str, Any] rule: :raises: InvalidRuleError if the rule is invalid. CollectionError if the creation fails. """ self.validate(rule_id, rule) self.rule_collection.update({ RuleField.id: rule_id }, rule, upsert=False) def validate(self, rule_id, rule): """ Check that the rule is valid. The pattern and external_data fields are not validated by this method. :param Dict[str, Any] view: :raises: InvalidRuleError if it is invalid. """ # Validate id field if rule.get(RuleField.id, rule_id) != rule_id: raise InvalidRuleError( 'The {0} field should not be modified.'.format(RuleField.id)) # Check that there are no unexpected fields in the rule unexpected_fields = set(rule.keys()).difference(RuleField.values) if unexpected_fields: raise InvalidRuleError( 'Unexpected fields: {0}.'.format(', '.join(unexpected_fields))) # Validate the type field if RuleField.type not in rule: raise InvalidRuleError( 'The {0} field is required.'.format(RuleField.type)) if rule.get(RuleField.type) not in RuleType.values: raise InvalidRuleError( 'The {0} field should be one of: {1}.'.format( RuleField.type, ', '.join(RuleType.values))) # Validate the priority field if not isinstance(rule.get(RuleField.priority, 0), int): raise InvalidRuleError( 'The {0} field should be an integer.'.format( RuleField.priority)) # Validate the enabled field if not isinstance(rule.get(RuleField.enabled, True), bool): raise InvalidRuleError( 'The {0} field should be a boolean.'.format( RuleField.enabled)) if rule.get(RuleField.type) != RuleType.enrichment: # Check that the enrichment fields are not defined for # non-enrichment rules. unexpected_fields = set(rule.keys()).intersection( ENRICHMENT_FIELDS) if unexpected_fields: raise InvalidRuleError( 'The following fields should only be defined for ' 'enrichment rules: {0}.'.format( ', '.join(unexpected_fields))) else: # Validate the actions field of the enrichment rules. if RuleField.actions not in rule: raise InvalidRuleError( 'The {0} field is required for enrichment rules.'.format( RuleField.actions)) if not isinstance(rule.get(RuleField.actions), list): raise InvalidRuleError( 'The {0} field should be a list.'.format( RuleField.actions)) if not rule.get(RuleField.actions): raise InvalidRuleError( 'The {0} field should contain at least one action.'.format( RuleField.actions)) # Validate the on_success field of the enrichment rules. outcome = rule.get(RuleField.on_success) if outcome and outcome not in RuleOutcome.values: raise InvalidRuleError( 'The {0} field should be one of: {1}.'.format( RuleField.on_success, ', '.join(RuleOutcome.values))) # Validate the on_failure field of the enrichment rules. outcome = rule.get(RuleField.on_failure) if outcome and outcome not in RuleOutcome.values: raise InvalidRuleError( 'The {0} field should be one of: {1}.'.format( RuleField.on_failure, ', '.join(RuleOutcome.values)))
class TestActionManager(unittest.TestCase): def setUp(self): output = StringIO() self.logger = Logger.get('test', output, OutputStream) store = MongoStore.get_default() self.collection = store.get_collection(name='default_test') self.mongo_collection = MongoCollection( collection=self.collection, logger=self.logger ) # Cleanup self.tearDown() self.manager = ActionManager( logger=self.logger, mongo_collection=self.mongo_collection ) self.id_ = 'testid' self.action = { "_id": self.id_, "type": "pbehavior", "fields": ["Resource"], "regex": ".*wine.*", "parameters": { "author": "Matho", "name": "Salammbo", "reason": "Madness", "type": "Mercenary War", "rrule": "" } } def tearDown(self): """Teardown""" self.mongo_collection.remove({}) def test_crud(self): res = self.manager.create(action=self.action) self.assertTrue(res) res = self.manager.get_id(self.id_) self.assertIsNotNone(res) self.assertDictEqual(res.to_dict(), self.action) action2 = self.action.copy() action2[Action.FIELDS] = ['Component'] res = self.manager.update_id(id_=self.id_, action=action2) self.assertTrue(res) res = self.manager.get_id(self.id_) self.assertIsNotNone(res) self.assertDictEqual(res.to_dict(), action2) res = self.manager.delete_id(id_=self.id_) self.assertTrue(res) res = self.manager.get_id(self.id_) self.assertIsNone(res)
class HeartbeatManager(object): """ Heartbeat service manager abstraction. """ COLLECTION = 'heartbeat' LOG_PATH = 'var/log/heartbeat.log' @classmethod def provide_default_basics(cls): """ Provide logger, config, storages... ! Do not use in tests ! :rtype: Union[logging.Logger, canopsis.common.collection.MongoCollection] """ store = MongoStore.get_default() collection = store.get_collection(name=cls.COLLECTION) return (Logger.get('action', cls.LOG_PATH), MongoCollection(collection)) def __init__(self, logger, collection): """ :param `~.logger.Logger` logger: object. :param `~.common.collection.MongoCollection` collection: object. """ self.__logger = logger self.__collection = MongoCollection(collection) def create(self, heartbeat): """ Create a new Heartbeat. :param `HeartBeat` heartbeat: a Heartbeat model. :returns: a created Heartbeat ID. :rtype: `str`. :raises: (`.HeartbeatPatternExistsError`, `pymongo.errors.PyMongoError`, `~.common.collection.CollectionError`, ). """ if self.get(heartbeat.id): raise HeartbeatPatternExistsError() return self.__collection.insert(heartbeat.to_dict()) def get(self, heartbeat_id=None): """ Get Heartbeat by ID or a list of Heartbeats when calling with default arguments. :param `Optional[str]` heartbeat_id: :returns: list of Heartbeat documents if **heartbeat_id** is None else single Heartbeat document or None if not found. :raises: (`pymongo.errors.PyMongoError`, ). """ if heartbeat_id: return self.__collection.find_one({"_id": heartbeat_id}) return [x for x in self.__collection.find({})] def delete(self, heartbeat_id): """ Delete Heartbeat by ID. :param `str` heartbeat_id: Heartbeat ID. :return: :raises: (`~.common.collection.CollectionError`, ). """ return self.__collection.remove({"_id": heartbeat_id})
class SessionManagerTest(TestCase): def setUp(self): self.storage = Middleware.get_middleware_by_uri( 'mongodb-default-testsession://') self.collection = MongoCollection(self.storage._backend) self.manager = Session(collection=self.collection) self.user = '******' self.id_beaker_session = 'cm9vdF8xNTc2MDY1MzY2' self.path = [ "view/da7ac9b9-db1c-4435-a1f2-edb4d6be4db8", "view-tab_edd5855b-54f1-4c51-9550-d88c2da60768" ] self.path_bis = [ "view/da7ac9b9-db1c-4435-a1f2-edb4d6be4db8", "view-tab_edd5855b-54f1-4c51-azerty" ] def tearDown(self): self.collection.remove() def test_keep_alive(self): self.manager.session_start(self.id_beaker_session, self.user) sleep(1) got = self.manager.keep_alive(self.id_beaker_session, self.user, True, self.path) session = self.collection.find_one({ "id_beaker_session": self.id_beaker_session, 'username': self.user }) self.assertTrue(isinstance(session, dict)) self.assertEqual(got, session['last_ping']) self.assertEqual(self.path, session['last_visible_path']) got = self.manager.keep_alive(self.id_beaker_session, self.user, False, self.path_bis) session = self.collection.find_one({ "id_beaker_session": self.id_beaker_session, 'username': self.user }) self.assertTrue(isinstance(session, dict)) self.assertEqual(got, session['last_ping']) self.assertEqual(self.path, session['last_visible_path']) got = self.manager.keep_alive(self.id_beaker_session, self.user, True, self.path_bis) session = self.collection.find_one({ "id_beaker_session": self.id_beaker_session, 'username': self.user }) self.assertTrue(isinstance(session, dict)) self.assertEqual(got, session['last_ping']) self.assertEqual(self.path_bis, session['last_visible_path']) def test_session_start(self): got = self.manager.session_start(self.id_beaker_session, self.user) session = self.collection.find_one({ "id_beaker_session": self.id_beaker_session, 'username': self.user }) self.assertTrue(isinstance(session, dict)) self.assertTrue(self.manager.is_session_active(self.id_beaker_session)) self.assertEqual(got, session['start']) def test_session_start_already_started(self): self.test_session_start() got = self.manager.session_start(self.id_beaker_session, self.user) self.assertTrue(got is None) def test_is_session_active(self): self.assertFalse(self.manager.is_session_active( self.id_beaker_session)) self.manager.session_start(self.id_beaker_session, self.user) self.assertTrue(self.manager.is_session_active(self.id_beaker_session)) def test_session_tracepath(self): self.manager.session_start(self.id_beaker_session, self.user) sleep(1) got = self.manager.session_tracepath(self.id_beaker_session, self.user, self.path) session = self.collection.find_one({ "id_beaker_session": self.id_beaker_session, 'username': self.user }) self.assertTrue(isinstance(session, dict)) self.assertEqual(got, session['last_ping']) self.assertEqual(self.path, session['last_visible_path']) got = self.manager.session_tracepath(self.id_beaker_session, self.user, self.path_bis) session = self.collection.find_one({ "id_beaker_session": self.id_beaker_session, 'username': self.user }) self.assertTrue(isinstance(session, dict)) self.assertEqual(got, session['last_ping']) self.assertEqual(self.path_bis, session['last_visible_path']) got = self.manager.session_tracepath(self.id_beaker_session, self.user, self.path_bis) session = self.collection.find_one({ "id_beaker_session": self.id_beaker_session, 'username': self.user }) self.assertTrue(isinstance(session, dict)) self.assertEqual(got, session['last_ping']) self.assertEqual(self.path_bis, session['last_visible_path']) def test_sessions_req(self): self.manager.session_start(self.id_beaker_session, self.user) sleep(1) session = self.collection.find_one({ "id_beaker_session": self.id_beaker_session, 'username': self.user }) session_req = self.manager.sessions_req(self.id_beaker_session, {"active": "true"}) self.assertEqual([session], session_req) session_req = self.manager.sessions_req(self.id_beaker_session, {"active": "false"}) self.assertEqual([], session_req) self.manager.session_start("azerty", "userTest") session2 = self.collection.find_one({ "id_beaker_session": "azerty", 'username': "******" }) session_req = self.manager.sessions_req( self.id_beaker_session, {"usernames[]": [self.user, "userTest"]}) self.assertEqual([session, session2], session_req)
class ViewAdapter(object): """ Adapter for the view collection. """ def __init__(self, logger): self.logger = logger self.view_collection = MongoCollection( MongoStore.get_default().get_collection(VIEWS_COLLECTION)) self.group_collection = MongoCollection( MongoStore.get_default().get_collection(GROUPS_COLLECTION)) def get_by_id(self, view_id): """ Get a view given its id. :param str view_id: the id of the view. """ return self.view_collection.find_one({ViewField.id: view_id}) def create(self, view): """ Create a new view and return its id. :param Dict[str, Any] view: :rtype: str """ view_id = str(uuid4()) view[ViewField.id] = view_id self.validate(view_id, view) self.view_collection.insert(view) return view_id def update(self, view_id, view): """ Update a view given its id. :param str view_id: the id of the view. :param Dict[str, Any] view: """ self.validate(view_id, view) self.view_collection.update({ViewField.id: view_id}, view, upsert=False) def remove_with_id(self, view_id): """ Remove a view given its id. :param str view_id: the id of the view. """ self.view_collection.remove({ViewField.id: view_id}) def list(self, name=None, title=None): """ Return a list of views, optionally filtered by name or title. :param str name: :param str title: :rtype: List[Dict[str, Any]] :raises: InvalidFilterError """ view_filter = {} if name is not None and title is not None: raise InvalidFilterError( 'Cannot filter both by name and by title.') if name is not None: view_filter[ViewField.name] = name if title is not None: view_filter[ViewField.title] = title views = self.view_collection.find(view_filter) groups = self.group_collection.find({}) response_groups = {} for group in groups: group_id = group[GroupField.id] del group[GroupField.id] group[GroupField.views] = [] response_groups[group_id] = group for view in views: try: group_id = view[ViewField.group_id] except KeyError: # This should never happen as long as the collections are only # modified using the API self.logger.exception( 'The view {0} is missing the group_id field.'.format( view[ViewField.id])) continue try: response_groups[group_id][GroupField.views].append(view) except KeyError: # This should never happen as long as the collections are only # modified using the API self.logger.exception('No group with id: {0}'.format(group_id)) return {ViewResponseField.groups: response_groups} def validate(self, view_id, view): """ Check that the view is valid, return InvalidViewError if it is not. :param Dict[str, Any] view: """ # Validate id field if view.get(ViewField.id, view_id) != view_id: raise InvalidViewError( 'The {0} field should not be modified'.format(ViewField.id)) # Validate group_id field if ViewField.group_id not in view: raise InvalidViewError('The view should have a group_id field.') group_id = view[ViewField.group_id] group = self.group_collection.find_one({GroupField.id: group_id}) if not group: raise InvalidViewError('No group with id: {0}'.format(group_id)) # Validate name field if ViewField.name not in view: raise InvalidViewError('The view should have a name field.') view_name = view[ViewField.name] same_name_view = self.view_collection.find_one({ ViewField.id: { '$ne': view_id }, ViewField.name: view_name }) if same_name_view: raise InvalidViewError( 'There is already a view with the name: {0}'.format(view_name)) # Validate title field if ViewField.title not in view: raise InvalidViewError('The view should have a title field.')
class GroupAdapter(object): """ Adapter for the group collection. """ def __init__(self, logger): self.logger = logger self.group_collection = MongoCollection( MongoStore.get_default().get_collection(GROUPS_COLLECTION)) self.view_collection = MongoCollection( MongoStore.get_default().get_collection(VIEWS_COLLECTION)) def get_by_id(self, group_id, name=None, title=None): """ Get a group given its id. :param str group_id: the id of the group. :param str name: :param str title: :rtype: Dict[str, Any] """ group = self.group_collection.find_one({GroupField.id: group_id}) if group: group[GroupField.views] = self.get_views(group_id, name, title) return group def get_views(self, group_id, name=None, title=None): """ Returns the list of views of a group, optionally filtered by name or title. :param str group_id: :param str name: :param str title: :rtype: List[Dict[str, Any]] :raises: InvalidFilterError """ view_filter = {ViewField.group_id: group_id} if name is not None and title is not None: raise InvalidFilterError( 'Cannot filter both by name and by title.') if name is not None: view_filter[ViewField.name] = name if title is not None: view_filter[ViewField.title] = title return list(self.view_collection.find(view_filter)) def is_empty(self, group_id): """ Return True if a group is empty. :param str group_id: :rtype: bool """ return self.view_collection.find({ ViewField.group_id: group_id }).limit(1).count() == 0 def create(self, group): """ Create a new group. :param Dict group: :rtype: str :raises: InvalidGroupError """ group_id = str(uuid4()) group[GroupField.id] = group_id self.validate(group_id, group) self.group_collection.insert(group) return group_id def update(self, group_id, group): """ Update a group given its id. :param str group_id: :param Dict group: :raises: InvalidGroupError """ self.validate(group_id, group) self.group_collection.update({GroupField.id: group_id}, group, upsert=False) def remove_with_id(self, group_id): """ Remove a group given its id. :param str group_id: :raises: NonEmptyGroupError """ if not self.is_empty(group_id): raise NonEmptyGroupError() self.group_collection.remove({GroupField.id: group_id}) def list(self, name=None): """ Return a list of groups, optionally filtered by name. :param str name: :rtype: List[Dict[str, Any]] """ group_filter = {} if name is not None: group_filter[GroupField.name] = name return list(self.group_collection.find(group_filter)) def validate(self, group_id, group): """ Check that the gorup is valid, return InvalidGroupError if it is not. :param Dict[str, Any] view: :raises: InvalidGroupError """ if group.get(GroupField.id, group_id) != group_id: raise InvalidGroupError( 'The {0} field should not be modified'.format(GroupField.id)) if GroupField.name not in group: raise InvalidGroupError('The group should have a name field.') group_name = group[GroupField.name] same_name_group = self.group_collection.find_one({ GroupField.id: { '$ne': group_id }, GroupField.name: group_name }) if same_name_group: raise InvalidGroupError( 'There is already a group with the name: {0}'.format( group_name)) def exists(self, group_id): """ Return True if a group exists. :param str group_id: :rtype: bool """ group = self.group_collection.find_one({GroupField.id: group_id}) return group is not None