def class_from_id(type_, _id): """ Return an instantiated class object. :param type_: The CRIPTs top-level object type. :type type_: str :param _id: The ObjectId to search for. :type _id: str :returns: class which inherits from :class:`cripts.core.cripts_mongoengine.CriptsBaseAttributes` """ #Quick fail if not _id or not type_: return None # doing this to avoid circular imports from cripts.comments.comment import Comment from cripts.core.cripts_mongoengine import Action from cripts.core.source_access import SourceAccess from cripts.core.user_role import UserRole from cripts.events.event import Event from cripts.usernames.username import UserName from cripts.targets.target import Target from cripts.hashes.hash import Hash from cripts.datasets.dataset import Dataset from cripts.email_addresses.email_address import EmailAddress # make sure it's a string _id = str(_id) # Use bson.ObjectId to make sure this is a valid ObjectId, otherwise # the queries below will raise a ValidationError exception. if not ObjectId.is_valid(_id.decode('utf8')): return None if type_ == 'Comment': return Comment.objects(id=_id).first() elif type_ == 'Event': return Event.objects(id=_id).first() elif type_ == 'Action': return Action.objects(id=_id).first() elif type_ == 'SourceAccess': return SourceAccess.objects(id=_id).first() elif type_ == 'UserRole': return UserRole.objects(id=_id).first() elif type_ == 'UserName': return UserName.objects(id=_id).first() elif type_ == 'Target': return Target.objects(id=_id).first() elif type_ == 'Hash': return Hash.objects(id=_id).first() elif type_ == 'Dataset': return Dataset.objects(id=_id).first() elif type_ == 'EmailAddress': return EmailAddress.objects(id=_id).first() else: return None
def clean_db(): """ Clean up the DB after testing. """ src = SourceAccess.objects(name=TSRC).first() if src: src.delete() user = CRIPTsUser.objects(username=TUSER_NAME).first() if user: user.delete() TestObject.drop_collection() TestSourceObject.drop_collection() CRIPTsConfig.drop_collection()
def DelSource(self): self.assertTrue(SourceAccess.objects(name=TSRC).first()) SourceAccess.objects(name=TSRC).first().delete() self.assertFalse(SourceAccess.objects(name=TSRC).first())
def FindSource(self): self.assertEqual(SourceAccess.objects(name=TSRC).first().name, TSRC)
def setUp(self): src = SourceAccess.objects(name=TSRC).first() if src: src.delete()