def test_object_score_low_duplication(self): obj1 = { 'field1': 'value', 'field2': 'value1', 'field3': 'value2', 'field4': 'value3', 'field5': 'value4' } obj2 = { 'field1': str(uuid.uuid4()), 'field2': str(uuid.uuid4()), 'field3': str(uuid.uuid4()), 'field4': str(uuid.uuid4()), 'field5': str(uuid.uuid4()) } ioc = GreaseContainer() parent1 = ioc.getCollection('test_scoring').insert_one({ 'expiry': Deduplication.generate_expiry_time(1), 'max_expiry': Deduplication.generate_max_expiry_time(1), 'type': 1, 'score': 1, 'source': 'test_source', 'hash': Deduplication.generate_hash_from_obj(obj1) }).inserted_id score1 = Deduplication.object_field_score('test_scoring', ioc, 'test_source', 'test_configuration', obj1, parent1, 1, 1) parent2 = ioc.getCollection('test_scoring').insert_one({ 'expiry': Deduplication.generate_expiry_time(1), 'max_expiry': Deduplication.generate_max_expiry_time(1), 'type': 1, 'score': 1, 'source': 'test_source', 'hash': Deduplication.generate_hash_from_obj(obj2) }).inserted_id score2 = Deduplication.object_field_score('test_scoring', ioc, 'test_source', 'test_configuration', obj2, parent2, 1, 1) print("++++++++++++++++++++++++++++++++++") print("score1: {0}".format(score1)) print("score2: {0}".format(score2)) print("++++++++++++++++++++++++++++++++++") self.assertTrue(score1 == 0.0) self.assertTrue(score2 <= 20.0) ioc.getCollection('test_scoring').drop() time.sleep(1.5)
def test_generate_hash_other_type(self): obj = 7 self.assertEqual(Deduplication.generate_hash_from_obj(obj), hashlib.sha256(str(obj).encode('utf-8')).hexdigest()) obj = 'test' self.assertEqual(Deduplication.generate_hash_from_obj(obj), hashlib.sha256(str(obj).encode('utf-8')).hexdigest()) obj = u'test' self.assertEqual(Deduplication.generate_hash_from_obj(obj), hashlib.sha256(str(obj).encode('utf-8')).hexdigest()) obj = 7.8 self.assertEqual(Deduplication.generate_hash_from_obj(obj), hashlib.sha256(str(obj).encode('utf-8')).hexdigest()) obj = ['test', 'var', 8, 8.43] self.assertEqual(Deduplication.generate_hash_from_obj(obj), hashlib.sha256(str(obj).encode('utf-8')).hexdigest())
def scanComplete(self): """Enters a completed source so that this local server is alive next run This method is so that the server's 'heart' beats after each run. It will insert a completed SourceData document and increments the job counter in the JobServer Document Returns: None: Writes a MongoDB Document """ self.ioc.getCollection('SourceData').insert_one({ 'grease_data': { 'sourcing': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity) }, 'detection': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity), 'start': datetime.datetime.utcnow(), 'end': datetime.datetime.utcnow(), 'detection': {} }, 'scheduling': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity), 'start': datetime.datetime.utcnow(), 'end': datetime.datetime.utcnow() }, 'execution': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity), 'assignmentTime': datetime.datetime.utcnow(), 'completeTime': datetime.datetime.utcnow(), 'returnData': {}, 'executionSuccess': True, 'commandSuccess': True, 'failures': 0 } }, 'source': 'grease_internal_node_monitoring', 'configuration': None, 'data': {}, 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }) server = self.ioc.getCollection('JobServer').find_one( {'_id': ObjectId(self.ioc.getConfig().NodeIdentity)}) if not server: self.ioc.getLogger().critical( "Failed to find server [{0}] after monitoring occurred!". format(self.ioc.getConfig().NodeIdentity)) self.ioc.getCollection('JobServer').update_one( {'_id': ObjectId(self.ioc.getConfig().NodeIdentity)}, {'$set': { 'jobs': dict(server).get('jobs', 0) + 1 }})
def test_get_schedule_staged(self): d = Detect() d.ioc.getCollection('SourceData').insert_one({ 'grease_data': { 'sourcing': { 'server': ObjectId(d.ioc.getConfig().NodeIdentity) }, 'detection': { 'server': ObjectId(d.ioc.getConfig().NodeIdentity), 'detectionStart': None, 'detectionEnd': None, 'detection': {} }, 'scheduling': { 'schedulingServer': None, 'schedulingStart': None, 'schedulingEnd': None }, 'execution': { 'server': None, 'assignmentTime': None, 'executionStart': None, 'executionEnd': None, 'context': {}, 'executionSuccess': False, 'commandSuccess': False, 'failures': 0, 'retryTime': datetime.datetime.utcnow() } }, 'source': str('test').encode('utf-8'), 'configuration': str('testConfig').encode('utf-8'), 'data': { 'dev': 'test' }, 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }) self.assertTrue(d.getScheduledSource()) self.assertEqual( d.getScheduledSource().get('grease_data').get('detection').get( 'server'), ObjectId(d.ioc.getConfig().NodeIdentity)) d.ioc.getCollection('SourceData').drop()
def test_deduplicate_object(self): ioc = GreaseContainer() ioc.getConfig().set('verbose', True, 'Logging') obj = [{ 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': str(uuid.uuid4()), 'field1': 'var1', 'field2': str(uuid.uuid4()), 'field3': str(uuid.uuid4()), 'field4': 'var4', 'field5': str(uuid.uuid4()), }] finalObj = [] Deduplication.deduplicate_object(ioc, obj[0], 1, 1, 40.0, 'test_source', 'test_configuration', finalObj, 'test_source') self.assertEqual(len(finalObj), 1) Deduplication.deduplicate_object(ioc, obj[1], 1, 1, 40.0, 'test_source', 'test_configuration', finalObj, 'test_source') self.assertEqual(len(finalObj), 1) Deduplication.deduplicate_object(ioc, obj[2], 1, 1, 40.0, 'test_source', 'test_configuration', finalObj, 'test_source') self.assertGreaterEqual(len(finalObj), 1) ioc.getConfig().set('verbose', False, 'Logging') ioc.getCollection('test_source').drop() time.sleep(1.5)
def test_scheduleScheduling(self): d = Detect() p = PrototypeConfig(d.ioc) source = { 'key': 'var', 'ver': 'key', 'greg': 'old', 'created': '2017-11-24' } configuration = { 'name': 'demoConfig', 'job': 'otherThing', 'exe_env': 'general', 'source': 'Google', 'logic': { 'DateRange': [{ 'field': 'created', 'format': '%Y-%m-%d', 'min': '2017-11-23', 'max': '2017-11-25' }], 'Regex': [{ 'field': 'key', 'pattern': '.*', 'variable': True, 'variable_name': 'field' }, { 'field': 'ver', 'pattern': '.*' }], 'Exists': [{ 'field': 'greg', 'variable': True, 'variable_name': 'greg' }] } } p.load(True, [configuration]) sourceId = d.ioc.getCollection('SourceData').insert_one({ 'grease_data': { 'sourcing': { 'server': ObjectId(d.ioc.getConfig().NodeIdentity) }, 'detection': { 'server': ObjectId(d.ioc.getConfig().NodeIdentity), 'detectionStart': None, 'detectionEnd': None, 'detection': {} }, 'scheduling': { 'schedulingServer': None, 'schedulingStart': None, 'schedulingEnd': None }, 'execution': { 'server': None, 'assignmentTime': None, 'executionStart': None, 'executionEnd': None, 'context': {}, 'executionSuccess': False, 'commandSuccess': False, 'failures': 0, 'retryTime': datetime.datetime.utcnow() } }, 'source': str('test').encode('utf-8'), 'configuration': configuration.get('name'), 'data': source, 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }).inserted_id scheduleServer = d.ioc.getCollection('JobServer').insert_one({ 'jobs': 0, 'os': platform.system().lower(), 'roles': ["general"], 'prototypes': ["schedule"], 'active': True, 'activationTime': datetime.datetime.utcnow() }).inserted_id self.assertTrue(d.detectSource()) self.assertFalse(d.getScheduledSource()) self.assertTrue( len( d.ioc.getCollection('SourceData').find_one( { 'grease_data.scheduling.server': ObjectId(scheduleServer), 'grease_data.scheduling.start': None, 'grease_data.scheduling.end': None }, sort=[('createTime', pymongo.DESCENDING)]))) d.ioc.getCollection('JobServer').delete_one( {'_id': ObjectId(scheduleServer)}) d.ioc.getCollection('SourceData').delete_one( {'_id': ObjectId(sourceId)})
def execute(self, context): """Command execute method This will run continuously waiting for `timeToRun` to return true to call `run` Args: context (dict): context for the command to use **Not Used Here** """ while True: if self.timeToRun(): result = self.run() self.ioc.getLogger().debug( "Scheduled command result of [{0}] was [{1}]".format( self.__class__.__name__, result), verbose=True) self.ioc.getLogger().info( "Scheduled Command [{0}] has executed".format( self.__class__.__name__)) self.ioc.getCollection('SourceData').insert_one({ 'grease_data': { 'sourcing': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity) }, 'detection': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity), 'start': datetime.datetime.utcnow(), 'end': datetime.datetime.utcnow(), 'detection': {} }, 'scheduling': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity), 'start': datetime.datetime.utcnow(), 'end': datetime.datetime.utcnow() }, 'execution': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity), 'assignmentTime': datetime.datetime.utcnow(), 'completeTime': datetime.datetime.utcnow(), 'returnData': {}, 'executionSuccess': True, 'commandSuccess': True, 'failures': 0 } }, 'source': 'grease_internal_scheduled_job_{0}'.format( self.__class__.__name__), 'configuration': None, 'data': {}, 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }) else: self.ioc.getLogger().trace( "Scheduled Job not scheduled to run [{0}]".format( self.__class__.__name__), trace=True) continue
def test_deduplication(self): ioc = GreaseContainer() dedup = Deduplication(ioc) ioc.getConfig().set('verbose', True, 'Logging') obj = [{ 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': str(uuid.uuid4()), 'field1': str(uuid.uuid4()), 'field2': str(uuid.uuid4()), 'field3': str(uuid.uuid4()), 'field4': str(uuid.uuid4()), 'field5': str(uuid.uuid4()) }, { 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': 'var', 'field1': 'var1', 'field2': 'var2', 'field3': 'var3', 'field4': 'var4', 'field5': 'var5', }, { 'field': str(uuid.uuid4()), 'field1': str(uuid.uuid4()), 'field2': str(uuid.uuid4()), 'field3': str(uuid.uuid4()), 'field4': str(uuid.uuid4()), 'field5': str(uuid.uuid4()) }, { 'field': str(uuid.uuid4()), 'field1': str(uuid.uuid4()), 'field2': str(uuid.uuid4()), 'field3': str(uuid.uuid4()), 'field4': str(uuid.uuid4()), 'field5': str(uuid.uuid4()) }, { 'field': str(uuid.uuid4()), 'field1': str(uuid.uuid4()), 'field2': str(uuid.uuid4()), 'field3': str(uuid.uuid4()), 'field4': str(uuid.uuid4()), 'field5': str(uuid.uuid4()) }] finalObj = dedup.Deduplicate(obj, 'test_source', 'test_configuration', 85.0, 1, 1, 'test_source') self.assertGreaterEqual(len(finalObj), 4) ioc.getConfig().set('verbose', False, 'Logging') ioc.getCollection('test_source').drop() time.sleep(1.5)
def test_generate_hash_multi_str_type(self): obj = {'test': u'var', 'test1': 5, 'test2': 7.89, 'test3': 'ver'} self.assertEqual(Deduplication.generate_hash_from_obj(obj), hashlib.sha256(str(obj).encode('utf-8')).hexdigest())
def test_generate_max_expiry_time(self): print(str(Deduplication.generate_max_expiry_time(7))) print(str(datetime.datetime.utcnow() + datetime.timedelta(days=7))) self.assertTrue( Deduplication.generate_max_expiry_time(7).day == ( datetime.datetime.utcnow() + datetime.timedelta(days=7)).day)
def test_generate_expiry_time(self): self.assertTrue( Deduplication.generate_expiry_time(12).hour == ( datetime.datetime.utcnow() + datetime.timedelta(hours=12)).hour)
def test_comparison(self): self.assertTrue( Deduplication.string_match_percentage("Hello", "Hallo") == 0.8)
def test_scheduling_fail(self): d = Detect() p = PrototypeConfig(d.ioc) s = Scheduler(d.ioc) source = { 'key': 'var', 'ver': 'key', 'greg': 'old', 'created': '2017-11-24' } configuration = { 'name': 'demoConfig', 'job': 'otherThing', 'exe_env': 'minix', 'source': 'Google', 'logic': { 'DateRange': [ { 'field': 'created', 'format': '%Y-%m-%d', 'min': '2017-11-23', 'max': '2017-11-25' } ], 'Regex': [ { 'field': 'key', 'pattern': '.*', 'variable': True, 'variable_name': 'field' }, { 'field': 'ver', 'pattern': '.*' } ], 'Exists': [ { 'field': 'greg', 'variable': True, 'variable_name': 'greg' } ] } } p.load(True, [configuration]) sourceId = d.ioc.getCollection('SourceData').insert_one({ 'grease_data': { 'sourcing': { 'server': ObjectId(d.ioc.getConfig().NodeIdentity) }, 'detection': { 'server': ObjectId(d.ioc.getConfig().NodeIdentity), 'detectionStart': None, 'detectionEnd': None, 'detection': {} }, 'scheduling': { 'schedulingServer': None, 'schedulingStart': None, 'schedulingEnd': None }, 'execution': { 'server': None, 'assignmentTime': None, 'executionStart': None, 'executionEnd': None, 'context': {}, 'executionSuccess': False, 'commandSuccess': False, 'failures': 0, 'retryTime': datetime.datetime.utcnow() } }, 'source': str('test').encode('utf-8'), 'configuration': configuration.get('name'), 'data': source, 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }).inserted_id d.ioc.getCollection('JobServer').update_one( {'_id': ObjectId(d.ioc.getConfig().NodeIdentity)}, { '$set': { 'prototypes': ["scan", "detect", "schedule"] } } ) self.assertTrue(d.detectSource()) self.assertFalse(d.getScheduledSource()) self.assertTrue(s.getDetectedSource()) self.assertFalse(s.scheduleExecution()) self.assertTrue(d.ioc.getCollection('JobServer').update_one( {'_id': ObjectId(d.ioc.getConfig().NodeIdentity)}, { '$set': { 'prototypes': [] } } )) d.ioc.getCollection('SourceData').delete_one({'_id': ObjectId(sourceId)})
def test_rescheduleJobsFailed(self): n = NodeMonitoring() p = PrototypeConfig(n.ioc) server1 = n.ioc.getCollection('JobServer').insert_one({ 'jobs': 9, 'os': platform.system().lower(), 'roles': [], 'prototypes': [], 'active': True, 'activationTime': datetime.datetime.utcnow() }).inserted_id server2 = n.ioc.getCollection('JobServer').insert_one({ 'jobs': 9, 'os': platform.system().lower(), 'roles': ['test1'], 'prototypes': ['detect', 'schedule'], 'active': True, 'activationTime': datetime.datetime.utcnow() }).inserted_id config = n.ioc.getCollection('Configuration').insert_one({ 'active': True, 'type': 'prototype_config', "name": "test", "job": "help", "exe_env": "test", "source": "test", "logic": { "Regex": [{ "field": "url", "pattern": ".*", 'variable': True, 'variable_name': 'url' }], 'Range': [{ 'field': 'status_code', 'min': 199, 'max': 201 }] }, 'constants': { 'test': 'ver' } }).inserted_id p.load(reloadConf=True) source = n.ioc.getCollection('SourceData').insert_one({ 'grease_data': { 'sourcing': { 'server': server1 }, 'detection': { 'server': server1, 'start': datetime.datetime.utcnow(), 'end': datetime.datetime.utcnow(), 'detection': {} }, 'scheduling': { 'server': server1, 'start': datetime.datetime.utcnow(), 'end': datetime.datetime.utcnow(), }, 'execution': { 'server': server1, 'assignmentTime': datetime.datetime.utcnow(), 'completeTime': None, 'returnData': {}, 'executionSuccess': False, 'commandSuccess': False, 'failures': 0 } }, 'source': 'test', 'configuration': 'test', 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }).inserted_id self.assertTrue(n.deactivateServer(str(server1))) self.assertFalse(n.rescheduleJobs(str(server1))) self.assertFalse( n.ioc.getCollection('SourceData').find({ '_id': source, 'grease_data.execution.server': server2 }).count()) self.assertEqual( n.ioc.getCollection('JobServer').delete_one({ '_id': server1 }).deleted_count, 1) self.assertEqual( n.ioc.getCollection('JobServer').delete_one({ '_id': server2 }).deleted_count, 1) n.ioc.getCollection('SourceData').delete_one({'_id': source}) n.ioc.getCollection('Configuration').drop() n.ioc.getCollection('ServerHealth').drop() p.load(reloadConf=True)
def test_rescheduleDetectJobsFailed(self): n = NodeMonitoring() server1 = n.ioc.getCollection('JobServer').insert_one({ 'jobs': 9, 'os': platform.system().lower(), 'roles': [], 'prototypes': [], 'active': True, 'activationTime': datetime.datetime.utcnow() }).inserted_id n.ioc.getCollection('JobServer').update_many( {'active': True}, {'$set': { 'prototypes': [] }}) source = n.ioc.getCollection('SourceData').insert_one({ 'grease_data': { 'sourcing': { 'server': server1 }, 'detection': { 'server': server1, 'start': None, 'end': None, 'detection': {} }, 'scheduling': { 'server': None, 'start': None, 'end': None }, 'execution': { 'server': None, 'assignmentTime': None, 'completeTime': None, 'returnData': {}, 'executionSuccess': False, 'commandSuccess': False, 'failures': 0 } }, 'source': 'test', 'configuration': 'test', 'data': { 'test': 'ver' }, 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }).acknowledged self.assertTrue(n.deactivateServer(str(server1))) self.assertFalse(n.rescheduleDetectJobs(str(server1))) self.assertEqual( n.ioc.getCollection('JobServer').delete_one({ '_id': server1 }).deleted_count, 1) n.ioc.getCollection('SourceData').delete_one({'_id': source}) n.ioc.getCollection('ServerHealth').drop()
def test_job_execution(self): ioc = GreaseContainer() cmd = DaemonProcess(ioc) proto = PrototypeConfig(ioc) ioc.getCollection('Configuration').insert_one( { 'active': True, 'type': 'prototype_config', "name": "exe_test", "job": "help", "exe_env": "general", "source": "url_source", "logic": { "Regex": [ { "field": "url", "pattern": ".*", 'variable': True, 'variable_name': 'url' } ], 'Range': [ { 'field': 'status_code', 'min': 199, 'max': 201 } ] }, 'constants': { 'test': 'ver' } } ) proto.load(reloadConf=True) jobid = ioc.getCollection('SourceData').insert_one({ 'grease_data': { 'sourcing': { 'server': ObjectId(ioc.getConfig().NodeIdentity) }, 'detection': { 'server': ObjectId(ioc.getConfig().NodeIdentity), 'start': datetime.datetime.utcnow(), 'end': datetime.datetime.utcnow(), 'detection': {} }, 'scheduling': { 'server': ObjectId(ioc.getConfig().NodeIdentity), 'start': datetime.datetime.utcnow(), 'end': datetime.datetime.utcnow(), }, 'execution': { 'server': ObjectId(ioc.getConfig().NodeIdentity), 'assignmentTime': datetime.datetime.utcnow(), 'completeTime': None, 'returnData': {}, 'executionSuccess': False, 'commandSuccess': False, 'failures': 0 } }, 'source': 'dev', 'configuration': 'exe_test', 'data': {}, 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }).inserted_id # Run for a bit self.assertTrue(cmd.server()) self.assertTrue(cmd.drain_jobs(ioc.getCollection('SourceData'))) result = ioc.getCollection('SourceData').find_one({'_id': ObjectId(jobid)}) self.assertTrue(result) self.assertTrue(result.get('grease_data').get('execution').get('executionSuccess')) self.assertTrue(result.get('grease_data').get('execution').get('commandSuccess')) ioc.getCollection('SourceData').drop() ioc.getCollection('Configuration').drop()