def __init__(self, ioc=None): if isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.centralScheduler = Scheduling(self.ioc) self.scheduler = Scheduler(self.ioc)
def __init__(self, ioc=None): if isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.imp = ImportTool(self.ioc.getLogger()) self.monitor = NodeMonitoring(self.ioc)
def __init__(self, ioc=None): if ioc and isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.conf = PrototypeConfig(self.ioc) self.configs = []
def __init__(self, ioc=None): global GREASE_PROTOTYPE_CONFIGURATION if ioc and isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() if not GREASE_PROTOTYPE_CONFIGURATION: GREASE_PROTOTYPE_CONFIGURATION = self.load()
def __init__(self, *args, **kwargs): TestCase.__init__(self, *args, **kwargs) self.configuration = None self.enabled = False self.mock_data = {} self.expected_data = {} self.ioc = GreaseContainer() self.detect = Detect(self.ioc)
def test_registration(self): ioc = GreaseContainer() cmd = DaemonProcess(ioc) self.assertTrue(cmd.register()) collection = ioc.getMongo().Client().get_database('grease').get_collection('JobServer') self.assertTrue(collection.find({}).count()) del collection del cmd del ioc
def __init__(self, ioc=None): if ioc and isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.conf = PrototypeConfig(self.ioc) self.impTool = ImportTool(self.ioc.getLogger()) self.dedup = Deduplication(self.ioc) self.scheduler = Scheduling(self.ioc)
def __init__(self, ioc): if isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.current_real_second = datetime.utcnow().second if self.ioc.getConfig().NodeIdentity == "Unknown" and not self.register(): self.registered = False self.impTool = ImportTool(self.ioc.getLogger()) self.conf = PrototypeConfig(self.ioc)
def __init__(self, Logger=None): if Logging and isinstance(Logger, Logging): self.ioc = GreaseContainer(Logger) else: self.ioc = GreaseContainer() self.variable_storage = self.ioc.getMongo()\ .Client()\ .get_database(self.ioc.getConfig().get('Connectivity', 'MongoDB').get('db', 'grease'))\ .get_collection(self.__class__.__name__) self.start_time = datetime.utcnow() self.exec_data = {'execVal': False, 'retVal': False, 'data': {}}
def test_empty_conf(self): ioc = GreaseContainer() # clean up for root, dirnames, filenames in os.walk(ioc.getConfig().get('Configuration', 'dir')): for filename in fnmatch.filter(filenames, '*.config.json'): self.assertIsNone(os.remove(os.path.join(root, filename))) conf = PrototypeConfig() conf.load(reloadConf=True) self.assertTrue(conf.getConfiguration()) self.assertListEqual(conf.getConfiguration().get('configuration').get('pkg'), []) self.assertListEqual(conf.getConfiguration().get('configuration').get('fs'), []) self.assertListEqual(conf.getConfiguration().get('configuration').get('mongo'), []) conf.load(reloadConf=True)
def test_empty_source_schedule(self): ioc = GreaseContainer() sch = Scheduling(ioc) jServer = ioc.getCollection('JobServer') jID = jServer.insert_one({ 'jobs': 0, 'os': platform.system().lower(), 'roles': ["general"], 'prototypes': ["detect"], 'active': True, 'activationTime': datetime.datetime.utcnow() }).inserted_id time.sleep(1.5) self.assertFalse(sch.scheduleDetection('test', 'test_conf', [])) jServer.delete_one({'_id': ObjectId(jID)}) ioc.getCollection('SourceData').drop()
def setUp(self): self.ioc = GreaseContainer() self.good_config = { "source": "kafka", "max_backlog": 20, "min_backlog": 5, "servers": ["server"], "topics": ["topic"] } self.bad_config = {"source": "not kafka"} self.mock_thread = MagicMock() self.ks = KafkaSource()
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 parse_source(self, configuration): """This will Query the SQL Server to find data Args: configuration (dict): Configuration of Source. See Class Documentation above for more info Returns: bool: If True data will be scheduled for ingestion after deduplication. If False the engine will bail out """ ioc = GreaseContainer() if configuration.get('hour'): if datetime.datetime.utcnow().hour != int(configuration.get('hour')): # it is not the correct hour return True if configuration.get('minute'): if datetime.datetime.utcnow().minute != int(configuration.get('minute')): # it is not the correct hour return True if configuration.get('type') != 'postgresql': ioc.getLogger().error("Unsupported SQL Server Type; Currently Only supporting PostgreSQL", notify=False) return False else: # Attempt to get the DSN for the connection if os.environ.get(configuration.get('dsn')) and configuration.get('query'): # ensure the DSN is setup and the query is present try: DSN = os.environ.get(configuration.get('dsn')) with psycopg2.connect(DSN) as conn: with conn.cursor(cursor_factory=RealDictCursor) as cursor: cursor.execute(configuration.get('query')) data = cursor.fetchall() for row in data: self._data.append(row) del ioc return True except Exception as e: # Naked except to prevent issues around connections ioc.getLogger().error("Error processing configuration; Error [{0}]".format(e.message), notify=False) del ioc return False else: # could not get the DSN ioc.getLogger().error("Failed to locate the DSN variable", notify=False) del ioc return False
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)
class Command(object): """Abstract class for commands in GREASE Attributes: __metaclass__ (ABCMeta): Metadata class object purpose (str): The purpose of the command help (str): Help string for the command line __author__ (str): Authorship string __version__ (str): Command Version os_needed (str): If a specific OS is needed then set this ioc (GreaseContainer): IOC container for access to system resources variable_storage (pymongo.collection): collection object for command """ ### # Command Metadata information ### purpose = "Default" help = """ No Help Information Provided """ __author__ = "Jimmy The Programmer" __version__ = "1.0.0" os_needed = None __metaclass__ = ABCMeta def __init__(self, Logger=None): if Logging and isinstance(Logger, Logging): self.ioc = GreaseContainer(Logger) else: self.ioc = GreaseContainer() self.variable_storage = self.ioc.getMongo()\ .Client()\ .get_database(self.ioc.getConfig().get('Connectivity', 'MongoDB').get('db', 'grease'))\ .get_collection(self.__class__.__name__) self.start_time = datetime.utcnow() self.exec_data = {'execVal': False, 'retVal': False, 'data': {}} def getExecVal(self): """Get the execution attempt success Returns: bool: If the command executed without exception """ return self.exec_data.get('execVal', False) def getRetVal(self): """Get the execution boolean return state Returns: bool: the boolean return value of execute """ return self.exec_data.get('retVal', False) def getData(self): """Get any data the execute method wanted to put into telemetry Returns: dict: The Key/Value pairs from the execute method execution """ return self.exec_data.get('data', {}) def setData(self, Key, Data): """Put Data into the data object to be inserted into telemetry Args: Key (str): Key for the data to be stored Data (object): JSON-able object to store Returns: None: Void Method to put data """ self.exec_data['data'][Key] = Data def __del__(self): # close mongo connection self.ioc.getMongo().Close() def safe_execute(self, context=None): """Attempt execution and prevent MOST exceptions Args: context (dict): context for the command to use Returns: None: Void method to attempt exceptions """ if not context: context = {} try: try: self.exec_data['execVal'] = True self.exec_data['retVal'] = bool(self.execute(context)) except BaseException: self.exec_data['execVal'] = False exc_type, exc_obj, exc_tb = sys.exc_info() # Find initial traceback frame current_tb = exc_tb while current_tb.tb_next: current_tb = current_tb.tb_next self.ioc.getLogger().error( "Failed to execute [{0}] execute got exception!".format( self.__class__.__name__), additional={ 'file': os.path.split( current_tb.tb_frame.f_code.co_filename)[1], 'type': exc_type, 'line': current_tb.tb_lineno }) except: self.ioc.getLogger().error( "Failed to execute [{0}] execute got exception!".format( self.__class__.__name__), ) except: self.ioc.getLogger().error( "Failed to execute [{0}] execute major exception".format( self.__class__.__name__), ) @abstractmethod def execute(self, context): """Base Execute Method This method should *always* be overridden in child classes. This is the code that will run when your command is called. If this method is not implemented then the class will fail loading. Args: context (dict): context for the command to use Returns: bool: Command Success """ pass
class Scheduling(object): """Central scheduling class for GREASE This class routes data to nodes within GREASE Attributes: ioc (GreaseContainer): IoC access for DeDuplication """ def __init__(self, ioc=None): if isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.ioc.ensureRegistration() def scheduleDetection(self, source, configName, data): """Schedule a Source Parse to detection This method will take a list of single dimension dictionaries and schedule them for detection Args: source (str): Name of the source configName (str): Configuration Data was sourced from data (list[dict]): Data to be scheduled for detection Returns: bool: Scheduling success """ if len(data) is 0 or not isinstance(data, list): self.ioc.getLogger().trace( "Data provided empty or is not type list type: [{0}] len: [{1}]".format(str(type(data)), len(data)), trace=True ) return False self.ioc.getLogger().trace("Preparing to schedule [{0}] source objects".format(len(data)), trace=True) sourceCollect = self.ioc.getCollection('SourceData') jServerCollect = self.ioc.getCollection('JobServer') # begin scheduling loop of each block for elem in data: if not isinstance(elem, dict): self.ioc.getLogger().warning( "Element from data not of type dict! Got [{0}] DROPPED".format(str(type(elem))), notify=False ) continue server, jobCount = self.determineDetectionServer() if server: sourceCollect.insert_one({ 'grease_data': { 'sourcing': { 'server': ObjectId(self.ioc.getConfig().NodeIdentity) }, 'detection': { 'server': ObjectId(server), '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': str(source), 'configuration': str(configName), 'data': elem, 'createTime': datetime.datetime.utcnow(), 'expiry': Deduplication.generate_max_expiry_time(1) }) jServerCollect.update_one({ '_id': ObjectId(server)}, {'$set': {'jobs': int(jobCount) + 1}} ) else: self.ioc.getLogger().warning( "Failed to find detection server for data object from source [{0}]; DROPPED".format(source), notify=False ) self.ioc.getLogger().warning( "Detection scheduling failed. Could not find detection server", notify=False ) return False return True def scheduleScheduling(self, objectId): """Schedule a source for job scheduling This method schedules a source for job scheduling Args: objectId (str): MongoDB ObjectId to schedule Returns: bool: If scheduling was successful """ server, jobCount = self.determineSchedulingServer() if not server: self.ioc.getLogger().error("Failed to find scheduling server", notify=False) return False self.ioc.getCollection('SourceData').update_one( {'_id': ObjectId(objectId)}, { '$set': { 'grease_data.scheduling.server': ObjectId(server), 'grease_data.scheduling.start': None, 'grease_data.scheduling.end': None } } ) self.ioc.getCollection('SourceData').update_one({ '_id': ObjectId(server)}, {'$set': {'jobs': int(jobCount) + 1}} ) return True def determineDetectionServer(self): """Determines detection server to use Finds the detection server available for a new detection job Returns: tuple: MongoDB Object ID of server & current job count """ result = self.ioc.getCollection('JobServer').find({ 'active': True, 'prototypes': 'detect' }).sort('jobs', pymongo.ASCENDING).limit(1) if result.count(): return str(result[0]['_id']), int(result[0]['jobs']) else: return "", 0 def determineSchedulingServer(self): """Determines scheduling server to use Finds the scheduling server available for a new scheduling job Returns: tuple: MongoDB Object ID of server & current job count """ result = self.ioc.getCollection('JobServer').find({ 'active': True, 'prototypes': 'schedule' }).sort('jobs', pymongo.DESCENDING).limit(1) if result.count(): return str(result[0]['_id']), int(result[0]['jobs']) else: return "", 0 def determineExecutionServer(self, role): """Determines execution server to use Finds the execution server available for a new execution job Returns: str: MongoDB Object ID of server; if one cannot be found then string will be empty """ result = self.ioc.getCollection('JobServer').find({ 'active': True, 'roles': str(role) }).sort('jobs', pymongo.DESCENDING).limit(1) if result.count(): return str(result[0]['_id']), int(result[0]['jobs']) else: return "", 0
class DaemonProcess(object): """Actual daemon processing for GREASE Daemon Attributes: ioc (GreaseContainer): The Grease IOC current_real_second (int): Current second in time registered (bool): If the node is registered with MongoDB impTool (ImportTool): Instance of Import Tool conf (PrototypeConfig): Prototype Configuration Instance """ ioc = None current_real_second = None registered = True contextManager = {'jobs': {}, 'prototypes': {}} impTool = None def __init__(self, ioc): if isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.current_real_second = datetime.utcnow().second if self.ioc.getConfig( ).NodeIdentity == "Unknown" and not self.register(): self.registered = False self.impTool = ImportTool(self.ioc.getLogger()) self.conf = PrototypeConfig(self.ioc) def server(self): """Server process for ensuring prototypes & jobs are running By Running this method this will clear the DB of any jobs a node may have Returns: bool: Server Success """ # Ensure we aren't swamping the system cpu = cpu_percent(interval=.1) mem = virtual_memory().percent if \ cpu >= int(self.ioc.getConfig().get('NodeInformation', 'ResourceMax')) \ or mem >= int(self.ioc.getConfig().get('NodeInformation', 'ResourceMax')): self.ioc.getLogger().trace( "Thread Maximum Reached CPU: [{0}] Memory: [{1}]".format( cpu, mem), trace=True) # remove variables del cpu del mem return True if not self.registered: self.ioc.getLogger().trace("Server is not registered", trace=True) return False self.ioc.getLogger().trace("Server execution starting", trace=True) # establish job collection JobsCollection = self.ioc.getCollection("SourceData") self.ioc.getLogger().trace("Searching for Jobs", trace=True) jobs = JobsCollection.find({ 'grease_data.execution.server': ObjectId(self.ioc.getConfig().NodeIdentity), 'grease_data.execution.commandSuccess': False, 'grease_data.execution.executionSuccess': False, 'grease_data.execution.failures': { '$lt': 6 } }) # Get Node Information Node = self.ioc.getCollection('JobServer').find_one( {'_id': ObjectId(self.ioc.getConfig().NodeIdentity)}) if not Node: # If for some reason we couldn't find it self.ioc.getLogger().error("Failed To Load Node Information") return False # Get Prototypes prototypes = list(Node.get('prototypes')) # Del node instance del Node if prototypes: # We have prototypes to spin up for prototype in prototypes: self.ioc.getLogger().trace( "Passing ProtoType [{0}] to Runner".format(prototype), trace=True) self._run_prototype(prototype) if jobs.count(): self.ioc.getLogger().trace("Total Jobs to Execute: [{0}]".format( jobs.count())) for job in jobs: self.ioc.getLogger().trace( "Passing Job [{0}] to Runner".format(job.get("_id")), trace=True) self._run_job(job, JobsCollection) else: # Nothing to Run for Jobs self.ioc.getLogger().trace("No Jobs Scheduled to Server", trace=True) self.ioc.getLogger().trace("Server execution complete", trace=True) return True def _run_job(self, job, JobCollection): """Run a On-Demand Job Args: job (dict): Job Data to execute JobCollection (pymongo.collection.Collection): JobCollection to update for telemetry Returns: None: Void Method to kickoff execution """ if not self.contextManager['jobs'].get(job.get('_id')): # New Job to run if isinstance(job.get('configuration'), bytes): conf = job.get('configuration').decode() else: conf = job.get('configuration') inst = self.impTool.load(self.conf.get_config(conf).get('job', '')) if inst and isinstance(inst, Command): inst.ioc.getLogger().foreground = self.ioc.getLogger( ).foreground thread = threading.Thread( target=inst.safe_execute, args=(job.get('grease_data', {}).get('detection', {}).get('detection', {}), ), name="GREASE DAEMON COMMAND EXECUTION [{0}]".format( job.get('_id'))) thread.daemon = True thread.start() self.contextManager['jobs'][job.get("_id")] = { 'thread': thread, 'command': inst } else: # Invalid Job del inst self.ioc.getLogger().warning("Invalid Job", additional=job) JobCollection.update_one({'_id': ObjectId(job['_id'])}, { '$set': { 'grease_data.execution.failures': job.get('failures', 0) + 1 } }) return else: # Job already executing if self.contextManager['jobs'].get( job.get('_id')).get('thread').isAlive(): # thread still executing return else: # Execution has ended self.ioc.getLogger().trace("Job [{0}] finished running".format( job.get('_id')), trace=True) finishedJob = self.contextManager['jobs'].get( job.get('_id')).get('command') # type: Command if finishedJob.getRetVal(): # job completed successfully JobCollection.update_one( {'_id': ObjectId(job.get('_id'))}, { '$set': { 'grease_data.execution.commandSuccess': finishedJob.getRetVal(), 'grease_data.execution.executionSuccess': finishedJob.getExecVal(), 'grease_data.execution.completeTime': datetime.utcnow(), 'grease_data.execution.returnData': finishedJob.getData() } }) else: # Job Failure self.ioc.getLogger().warning( "Job Failed [{0}]".format(job.get('_id')), additional=finishedJob.getData()) # TODO: Job Execution cooldown timing JobCollection.update_one({'_id': ObjectId(job['_id'])}, { '$set': { 'grease_data.execution.failures': job.get('grease_data', {}).get( 'execution', {}).get('failures', 0) + 1 } }) # close out job finishedJob.__del__() del finishedJob # remove from contextManager del self.contextManager['jobs'][job.get('_id')] return def _run_prototype(self, prototype): """Startup a ProtoType Args: prototype (str): ProtoType to start Returns: None: Void method to start prototype """ if not self.contextManager['prototypes'].get(prototype): # ProtoType has not started inst = self.impTool.load(prototype) if not isinstance(inst, Command): # invalid ProtoType self.log_once_per_second( "Invalid ProtoType [{0}]".format(prototype), level=ERROR) return inst.ioc.getLogger().foreground = self.ioc.getLogger().foreground thread = threading.Thread( target=inst.safe_execute, args=({}), name="GREASE DAEMON PROTOTYPE [{0}]".format(prototype)) thread.daemon = True thread.start() self.contextManager['prototypes'][prototype] = thread return else: # ensure thread is alive if self.contextManager['prototypes'].get(prototype).isAlive(): self.ioc.getLogger().trace( "ProtoType [{0}] is alive".format(prototype)) return else: # Thread died for some reason self.log_once_per_second( "ProtoType [{0}] Stopped".format(prototype), level=INFO) inst = self.impTool.load(prototype) if not isinstance(inst, Command): self.log_once_per_second( "Invalid ProtoType [{0}]".format(prototype), level=ERROR) return inst.ioc.getLogger().foreground = self.ioc.getLogger( ).foreground thread = threading.Thread( target=inst.execute, name="GREASE DAEMON PROTOTYPE [{0}]".format(prototype)) thread.daemon = True thread.start() self.contextManager['prototypes'][prototype] = thread return def drain_jobs(self, JobCollection): """Will drain jobs from the current context This method is used to prevent abnormal ending of executions Args: JobCollection (pymongo.collection.Collection): Job Collection Object Returns: bool: When job queue is emptied """ Threads = True while Threads: if self.contextManager['jobs']: jobs = {} for key, val in self.contextManager['jobs'].items(): if val['thread'].isAlive(): jobs[key] = val continue else: # Execution has ended self.ioc.getLogger().trace( "Job [{0}] finished running".format(key), trace=True) finishedJob = self.contextManager['jobs'].get(key).get( 'command') # type: Command if finishedJob.getRetVal(): # job completed successfully JobCollection.update_one({'_id': ObjectId(key)}, { '$set': { 'grease_data.execution.commandSuccess': finishedJob.getRetVal(), 'grease_data.execution.executionSuccess': finishedJob.getExecVal(), 'grease_data.execution.completeTime': datetime.utcnow(), 'grease_data.execution.returnData': finishedJob.getData() } }) else: # Job Failure self.ioc.getLogger().warning( "Job Failed [{0}]".format(key), additional=finishedJob.getData()) JobCollection.update_one({'_id': ObjectId(key)}, { '$set': { 'grease_data.execution.failures': val['command'].get('failures', 0) + 1 } }) # close out job finishedJob.__del__() del finishedJob self.contextManager['jobs'] = jobs else: Threads = False return True def register(self): """Attempt to register with MongoDB Returns: bool: Registration Success """ return self.ioc.ensureRegistration() def log_once_per_second(self, message, level=DEBUG, additional=None): """Log Message once per second Args: message (str): Message to log level (int): Log Level additional (object): Additional information that is able to be str'd Returns: None: Void Method to fire log message """ if self._has_time_progressed(): self.ioc.getLogger().TriageMessage(message=message, level=level, additional=additional) def _has_time_progressed(self): """Determines if the current second and the real second are not the same Returns: bool: if true then time has passed in a meaningful way """ if self.current_real_second != datetime.utcnow().second: self.current_real_second = datetime.utcnow().second return True else: return False
def __init__(self, ioc=None): if isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.ioc.ensureRegistration()
def test_registration(self): ioc = GreaseContainer() self.assertTrue(ioc.ensureRegistration())
def test_logger(self): ioc = GreaseContainer() self.assertTrue(isinstance(ioc.getLogger(), Logging))
def test_mongo(self): ioc = GreaseContainer() self.assertTrue(isinstance(ioc.getMongo(), Mongo))
def test_get_collection(self): ioc = GreaseContainer() self.assertTrue(isinstance(ioc.getMongo(), Mongo)) coll = ioc.getCollection('TestCollection') self.assertTrue(isinstance(coll, Collection)) self.assertEqual(coll.name, "TestCollection")
def test_config(self): ioc = GreaseContainer() self.assertTrue(isinstance(ioc.getConfig(), Configuration))
def test_notifications(self): ioc = GreaseContainer() self.assertTrue(isinstance(ioc.getNotification(), Notifications))
def test_detectionScheduling(self): ioc = GreaseContainer() ioc.ensureRegistration() sch = Scheduling(ioc) jServer = ioc.getCollection('JobServer') jID1 = jServer.insert_one({ 'jobs': 0, 'os': platform.system().lower(), 'roles': ["general"], 'prototypes': ["detect"], 'active': True, 'activationTime': datetime.datetime.utcnow() }).inserted_id time.sleep(1) jID2 = jServer.insert_one({ 'jobs': 0, 'os': platform.system().lower(), 'roles': ["general"], 'prototypes': ["detect"], 'active': True, 'activationTime': datetime.datetime.utcnow() }).inserted_id time.sleep(1) self.assertTrue( sch.scheduleDetection('test', 'test_conf', [ { 'test0': 'var0', 'test1': 'var1', 'test2': 'var2', 'test3': 'var3', 'test4': 'var4', 'test5': 'var5', 'test6': 'var6', 'test7': 'var7', 'test8': 'var8', 'test9': 'var9', 'test10': 'var10', }, { 'test0': 'var0', 'test1': 'var1', 'test2': 'var2', 'test3': 'var3', 'test4': 'var4', 'test5': 'var5', 'test6': 'var6', 'test7': 'var7', 'test8': 'var8', 'test9': 'var9', 'test10': 'var10', }, { 'test0': 'var0', 'test1': 'var1', 'test2': 'var2', 'test3': 'var3', 'test4': 'var4', 'test5': 'var5', 'test6': 'var6', 'test7': 'var7', 'test8': 'var8', 'test9': 'var9', 'test10': 'var10', }, { 'test0': 'var0', 'test1': 'var1', 'test2': 'var2', 'test3': 'var3', 'test4': 'var4', 'test5': 'var5', 'test6': 'var6', 'test7': 'var7', 'test8': 'var8', 'test9': 'var9', 'test10': 'var10', }, { 'test0': 'var0', 'test1': 'var1', 'test2': 'var2', 'test3': 'var3', 'test4': 'var4', 'test5': 'var5', 'test6': 'var6', 'test7': 'var7', 'test8': 'var8', 'test9': 'var9', 'test10': 'var10', }, { 'test0': 'var0', 'test1': 'var1', 'test2': 'var2', 'test3': 'var3', 'test4': 'var4', 'test5': 'var5', 'test6': 'var6', 'test7': 'var7', 'test8': 'var8', 'test9': 'var9', 'test10': 'var10', }, ])) time.sleep(1) self.assertEqual( ioc.getCollection('SourceData').find({ 'grease_data.detection.server': ObjectId(jID1) }).count(), 3) self.assertEqual( ioc.getCollection('SourceData').find({ 'grease_data.detection.server': ObjectId(jID2) }).count(), 3) self.assertEqual( ioc.getCollection('JobServer').find_one({'_id': ObjectId(jID1)})['jobs'], 3) self.assertEqual( ioc.getCollection('JobServer').find_one({'_id': ObjectId(jID2)})['jobs'], 3) jServer.delete_one({'_id': ObjectId(jID1)}) jServer.delete_one({'_id': ObjectId(jID2)}) ioc.getCollection('SourceData').drop()
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)
class PrototypeConfig(object): """Responsible for Scanning/Detection/Scheduling configuration Structure of Configuration:: { 'configuration': { 'pkg': [], # <-- Loaded from pkg_resources.resource_filename('tgt_grease.enterprise.Model', 'config/') 'fs': [], # <-- Loaded from `<GREASE_DIR>/etc/*.config.json` 'mongo': [] # <-- Loaded from the Configuration Mongo Collection }, 'raw': [], # <-- All loaded configurations 'sources': [], # <-- list of sources found in configurations 'source': {} # <-- keys will be source values list of configs for that source 'names': [], # <-- all configs via their name so to allow dialing 'name': {} # <-- all configs via their name so to allow being dialing } Structure of a configuration file:: { "name": String, "job": String, "exe_env": String, # <-- If not provided will be default as 'general' "source": String, "logic": { # I need to be the logical blocks for Detection } } Attributes: ioc (GreaseContainer): IOC access """ def __init__(self, ioc=None): global GREASE_PROTOTYPE_CONFIGURATION if ioc and isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() if not GREASE_PROTOTYPE_CONFIGURATION: GREASE_PROTOTYPE_CONFIGURATION = self.load() def getConfiguration(self): """Returns the Configuration Object loaded into memory Returns: dict: Configuration object """ global GREASE_PROTOTYPE_CONFIGURATION if not GREASE_PROTOTYPE_CONFIGURATION: self.load(reloadConf=True) return GREASE_PROTOTYPE_CONFIGURATION def load(self, reloadConf=False, ConfigurationList=None): """[Re]loads configuration data about the current execution node Configuration data loads from 3 places in GREASE. The first is internal to the package, if one were to manually add their own files into the package in the current directory following the file pattern. The next is following the same pattern but loaded from `<GREASE_DIR>/etc/`. The final place GREASE looks for configuration data is from the `configuration` collection in MongoDB Args: reloadConf (bool): If True this will reload the global object. False will return the object ConfigurationList (list of dict): If provided will load the list of dict for config after validation Note: Providing a configuration *automatically* reloads the memory structure of prototype configuration Returns: dict: Current Configuration information """ global GREASE_PROTOTYPE_CONFIGURATION if ConfigurationList: conf = dict() conf['configuration'] = dict() conf['configuration'][ 'ConfigurationList'] = self.validate_config_list( ConfigurationList) conf['raw'] = conf['configuration']['ConfigurationList'] # split by configuration sets # the list of configured sources conf['sources'] = list() # the actual configurations for each source conf['source'] = dict() # configurations to get via name conf['names'] = list() # the actual configurations for each config name conf['name'] = dict() for config in conf.get('raw'): # type: dict if config.get('source') in conf['sources']: conf['source'][config.get('source')].append(config) else: conf['sources'].append(config.get('source')) conf['source'][config.get('source')] = list() conf['source'][config.get('source')].append(config) if config.get('name') in conf['names']: self.ioc.getLogger().error( "Prototype Configuration [{0}] already found! Overwriting" .format(config.get('name'))) conf['name'][config.get('name')] = config else: conf['names'].append(config.get('name')) conf['name'][config.get('name')] = config GREASE_PROTOTYPE_CONFIGURATION = conf return conf # fill out raw results conf = dict() conf['configuration'] = dict() conf['raw'] = [] pkg = self.validate_config_list( self.load_from_fs( pkg_resources.resource_filename('tgt_grease.enterprise.Model', 'config/'))) for newConfig in pkg: conf['raw'].append(newConfig) conf['configuration']['pkg'] = pkg del pkg fs = self.validate_config_list( self.load_from_fs(self.ioc.getConfig().get('Configuration', 'dir'))) for newConfig in fs: conf['raw'].append(newConfig) conf['configuration']['fs'] = fs del fs mongo = self.validate_config_list(self.load_from_mongo()) for newConfig in mongo: conf['raw'].append(newConfig) conf['configuration']['mongo'] = mongo del mongo # split by configuration sets # the list of configured sources conf['sources'] = list() # the actual configurations for each source conf['source'] = dict() # configurations to get via name conf['names'] = list() # the actual configurations for each config name conf['name'] = dict() for config in conf.get('raw'): # type: dict if config.get('source') in conf['sources']: conf['source'][config.get('source')].append(config) else: conf['sources'].append(config.get('source')) conf['source'][config.get('source')] = list() conf['source'][config.get('source')].append(config) if config.get('name') in conf['names']: self.ioc.getLogger().error( "Prototype Configuration [{0}] already found! Overwriting". format(config.get('name'))) conf['name'][config.get('name')] = config else: conf['names'].append(config.get('name')) conf['name'][config.get('name')] = config # return block if not reloadConf: return conf else: GREASE_PROTOTYPE_CONFIGURATION = conf return conf def get_sources(self): """Returns the list of sources to be scanned Returns: list: List of sources """ global GREASE_PROTOTYPE_CONFIGURATION # type: dict if GREASE_PROTOTYPE_CONFIGURATION: return GREASE_PROTOTYPE_CONFIGURATION.get('sources', []) else: self.ioc.getLogger().error( "GREASE Prototype configuration is not loaded", trace=True, notify=False) return [] def get_source(self, name): """Get all configuration by source by name Args: name (str): Source name to get Returns: list[dict]: Configuration if found else empty dict """ global GREASE_PROTOTYPE_CONFIGURATION if GREASE_PROTOTYPE_CONFIGURATION: return GREASE_PROTOTYPE_CONFIGURATION.get('source').get(name, []) else: self.ioc.getLogger().error( "GREASE Prototype configuration not loaded", notify=False, trace=True) return [] def get_names(self): """Returns the list of names of configs Returns: list: List of config names """ global GREASE_PROTOTYPE_CONFIGURATION # type: dict if GREASE_PROTOTYPE_CONFIGURATION: return GREASE_PROTOTYPE_CONFIGURATION.get('names', []) else: self.ioc.getLogger().error( "GREASE Prototype configuration is not loaded", trace=True, notify=False) return [] def get_config(self, name): """Get Configuration by name Args: name (str): Configuration name to get Returns: dict: Configuration if found else empty dict """ global GREASE_PROTOTYPE_CONFIGURATION if GREASE_PROTOTYPE_CONFIGURATION: if GREASE_PROTOTYPE_CONFIGURATION.get('name'): return GREASE_PROTOTYPE_CONFIGURATION.get('name').get(name, {}) else: self.ioc.getLogger().error("GREASE Configuration Not Found", notify=False, trace=True) return {} else: self.ioc.getLogger().error( "GREASE Prototype configuration not loaded", notify=False, trace=True) return {} def load_from_fs(self, directory): """Loads configurations from provided directory Note: Pattern is `*.config.json` Args: directory (str): Directory to load from Returns: list of dict: configurations """ self.ioc.getLogger().trace( "Loading Configurations from directory [{0}]".format(directory), trace=True) intermediate = list() matches = [] for root, dirnames, filenames in os.walk(directory): for filename in fnmatch.filter(filenames, '*.config.json'): matches.append(os.path.join(root, filename)) for doc in matches: self.ioc.getLogger().trace("Attempting to load [{0}]".format(doc), trace=True) with open(doc, 'rb') as current_file: content = current_file.read() if isinstance(content, bytes): content = content.decode() try: intermediate.append(json.loads(content)) self.ioc.getLogger().trace( "Successfully loaded [{0}]".format(doc), trace=True) except ValueError: self.ioc.getLogger().error("Failed to load [{0}]".format(doc), trace=True, notify=False) continue self.ioc.getLogger().trace( "total documents returned from fs [{0}]".format(len(intermediate)), trace=True) return intermediate def load_from_mongo(self): """Returns all active configurations from the mongo collection Configuration Structure of Configuration expected in Mongo:: { "name": String, "job": String, "exe_env": String, # <-- If not provided will be default as 'general' "active": Boolean, # <-- set to true to load configuration "type": "prototype_config", # <-- MUST BE THIS VALUE; For it is the config type :) "source": String, "logic": { # I need to be the logical blocks for Detection } } Returns: list of dict: Configurations """ self.ioc.getLogger().trace("Loading Configurations from mongo", trace=True) mConf = [] for conf in self.ioc.getCollection('Configuration').find({ 'active': True, 'type': 'prototype_config' }): mConf.append(dict(conf)) self.ioc.getLogger().trace( "total documents returned from mongo [{0}]".format(len(mConf)), trace=True) return mConf def validate_config_list(self, configs): """Validates a configuration List Args: configs (list[dict]): Configuration List Returns: list: The Valid configurations """ final = [] self.ioc.getLogger().trace( "Total configurations to validate [{0}]".format(len(configs))) for conf in configs: if self.validate_config(conf): final.append(conf) return final def validate_config(self, config): """Validates a configuration The default JSON Schema is this:: { "name": String, "job": String, "exe_env": String, # <-- If not provided will be default as 'general' "source": String, "logic": { # I need to be the logical blocks for Detection } } Args: config (dict): Configuration to validate Returns: bool: If it is a valid configuration """ self.ioc.getLogger().trace( "Configuration to be validated: [{0}]".format(config), trace=True) if not isinstance(config, dict): self.ioc.getLogger().error( "Configuration Validation Failed! Not of Type Dict::Got [{0}]". format(str(type(config))), trace=True, notify=False) if config.get('name'): if not isinstance(config.get('name'), str): config['name'] = str(config.get('name')) else: self.ioc.getLogger().error( "Configuration does not have valid name field", trace=True, notify=False) return False if config.get('job'): if not isinstance(config.get('job'), str): config['job'] = str(config.get('job')) else: self.ioc.getLogger().error( "Configuration does not have valid job field", trace=True, notify=False) return False if config.get('source'): if not isinstance(config.get('source'), str): config['source'] = str(config.get('source')) else: self.ioc.getLogger().error( "Configuration does not have valid source field", trace=True, notify=False) return False if not isinstance(config.get('logic'), dict): self.ioc.getLogger().error( "Configuration does not have valid logic field", trace=True, notify=False) return False if not config.get('logic'): # empty dictionary check AKA no logical blocks return False for key, params in config.get('logic').items(): if not isinstance(params, list): self.ioc.getLogger().error( "Configuration logic field was not list!", trace=True, notify=False) return False for block in params: if not isinstance(block, dict): self.ioc.getLogger().error( "Configuration logical block was not dict", trace=True, notify=False) return False return True
class BridgeCommand(object): """Methods for Cluster Administration Attributes: imp (ImportTool): Import Tool Instance monitor (NodeMonitoring): Node Monitoring Model Instance """ def __init__(self, ioc=None): if isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.imp = ImportTool(self.ioc.getLogger()) self.monitor = NodeMonitoring(self.ioc) def action_register(self): """Ensures Registration of server Returns: bool: Registration status """ self.ioc.getLogger().debug("Registration Requested") if self.ioc.ensureRegistration(): print("Registration Complete!") self.ioc.getLogger().info("Registration Completed Successfully") return True print("Registration Failed!") self.ioc.getLogger().info("Registration Failed") return False def action_info(self, node=None, jobs=None, prototypeJobs=None): """Gets Node Information Args: node (str): MongoDB Object ID to get information about jobs (bool): If true then will retrieve jobs executed by this node prototypeJobs (bool): If true then prototype jobs will be printed as well Note: provide a node argument via the CLI --node=4390qwr2fvdew458239 Note: provide a jobs argument via teh CLI --jobs Note: provide a prototype jobs argument via teh CLI --pJobs Returns: bool: If Info was found """ if not self.ioc.ensureRegistration(): self.ioc.getLogger().error("Server not registered with MongoDB") print("Unregistered servers cannot talk to the cluster") return False valid, serverId = self.valid_server(node) if not valid: print("Invalid ObjectID") return False server = self.ioc.getCollection('JobServer').find_one( {'_id': ObjectId(str(serverId))}) if server: server = dict(server) print(""" <<<<<<<<<<<<<< SERVER: {0} >>>>>>>>>>>>>> Activation State: {1} Date: {2} Jobs: {3} Operating System: {4} Prototypes: {5} Execution Roles: {6} """.format(server.get('_id'), server.get('active'), server.get('activationTime'), server.get('jobs'), server.get('os'), server.get('prototypes'), server.get('roles'))) if jobs and prototypeJobs: print( "======================= SOURCING =======================") for job in self.ioc.getCollection('SourceData').find( {'grease_data.sourcing.server': ObjectId(serverId)}): print( """ ------------------------------- Job: {0} ------------------------------- """, job['_id']) if jobs and prototypeJobs: print( "======================= DETECTION =======================" ) for job in self.ioc.getCollection('SourceData').find( {'grease_data.detection.server': ObjectId(serverId)}): print(""" ------------------------------- Job: {0} Start Time: {1} End Time: {2} Context: {3} ------------------------------- """.format(job['_id'], job['grease_data']['detection']['start'], job['grease_data']['detection']['end'], job['grease_data']['detection']['detection'])) if jobs and prototypeJobs: print( "======================= SCHEDULING =======================" ) for job in self.ioc.getCollection('SourceData').find( {'grease_data.scheduling.server': ObjectId(serverId)}): print(""" ------------------------------- Job: {0} Start Time: {1} End Time: {2} ------------------------------- """.format(job['_id'], job['grease_data']['scheduling']['start'], job['grease_data']['scheduling']['end'])) if jobs: print( "======================= EXECUTION =======================" ) for job in self.ioc.getCollection('SourceData').find( {'grease_data.execution.server': ObjectId(serverId)}): print(""" ------------------------------- Job: {0} Assignment Time: {1} Completed Time: {2} Execution Success: {3} Command Success: {4} Failures: {5} Return Data: {6} ------------------------------- """.format( job['_id'], job['grease_data']['execution']['assignmentTime'], job['grease_data']['execution']['completeTime'], job['grease_data']['execution']['executionSuccess'], job['grease_data']['execution']['commandSuccess'], job['grease_data']['execution']['failures'], job['grease_data']['execution']['returnData'])) return True print("Unable to locate server") self.ioc.getLogger().error( "Unable to load [{0}] server for information".format(serverId)) return False def action_assign(self, prototype=None, role=None, node=None): """Assign prototypes/roles to a node either local or remote Args: prototype (str): Prototype Job to assign role (str): Role to assign node (str): MongoDB ObjectId of node to assign to, if not provided will default to the local node Returns: bool: If successful true else false """ assigned = False if prototype: job = self.imp.load(str(prototype)) if not job or not isinstance(job, Command): print( "Cannot find prototype [{0}] to assign check search path!". format(prototype)) self.ioc.getLogger().error( "Cannot find prototype [{0}] to assign check search path!". format(prototype)) return False # Cleanup job job.__del__() del job valid, serverId = self.valid_server(node) if not valid: print("Invalid ObjectID") return False updated = self.ioc.getCollection('JobServer').update_one( { '_id': ObjectId(serverId) }, { '$addToSet': { 'prototypes': prototype } }).acknowledged if updated: print("Prototype Assigned") self.ioc.getLogger().info( "Prototype [{0}] assigned to server [{1}]".format( prototype, serverId)) assigned = True else: print("Prototype Assignment Failed!") self.ioc.getLogger().info( "Prototype [{0}] assignment failed to server [{1}]".format( prototype, serverId)) return False if role: valid, serverId = self.valid_server(node) if not valid: print("Invalid ObjectID") return False updated = self.ioc.getCollection('JobServer').update_one( { '_id': ObjectId(serverId) }, { '$push': { 'roles': role } }).acknowledged if updated: print("Role Assigned") self.ioc.getLogger().info( "Role [{0}] assigned to server [{1}]".format( prototype, serverId)) assigned = True else: print("Role Assignment Failed!") self.ioc.getLogger().info( "Role [{0}] assignment failed to server [{1}]".format( prototype, serverId)) return False if not assigned: print("Assignment failed, please check logs for details") return assigned def action_unassign(self, prototype=None, role=None, node=None): """Unassign prototypes to a node either local or remote Args: prototype (str): Prototype Job to unassign role (str): Role to unassign node (str): MongoDB ObjectId of node to unassign to, if not provided will default to the local node Returns: bool: If successful true else false """ unassigned = False if prototype: job = self.imp.load(str(prototype)) if not job or not isinstance(job, Command): print( "Cannot find prototype [{0}] to unassign check search path!" .format(prototype)) self.ioc.getLogger().error( "Cannot find prototype [{0}] to unassign check search path!" .format(prototype)) return False # Cleanup job job.__del__() del job valid, serverId = self.valid_server(node) if not valid: print("Invalid ObjectID") return False updated = self.ioc.getCollection('JobServer').update_one( { '_id': ObjectId(serverId) }, { '$pull': { 'prototypes': prototype } }).acknowledged if updated: print("Prototype Assignment Removed") self.ioc.getLogger().info( "Prototype [{0}] unassigned from server [{1}]".format( prototype, serverId)) unassigned = True else: print("Prototype Unassignment Failed!") self.ioc.getLogger().info( "Prototype [{0}] unassignment failed from server [{1}]". format(prototype, serverId)) return False if role: valid, serverId = self.valid_server(node) if not valid: print("Invalid ObjectID") return False updated = self.ioc.getCollection('JobServer').update_one( { '_id': ObjectId(serverId) }, { '$pull': { 'roles': role } }).acknowledged if updated: print("Role Removed") self.ioc.getLogger().info( "Role [{0}] removed to server [{1}]".format( prototype, serverId)) unassigned = True else: print("Role Removal Failed!") self.ioc.getLogger().info( "Role [{0}] removal failed to server [{1}]".format( prototype, serverId)) return False if not unassigned: print("Unassignment failed, please check logs for details") return unassigned def action_cull(self, node=None): """Culls a server from the active cluster Args: node (str): MongoDB ObjectId to cull; defaults to local node """ if not self.ioc.ensureRegistration(): self.ioc.getLogger().error("Server not registered with MongoDB") print("Unregistered servers cannot talk to the cluster") return False valid, serverId = self.valid_server(node) if not valid: print("Invalid ObjectID") return False if not self.monitor.deactivateServer(serverId): self.ioc.getLogger().error( "Failed deactivating server [{0}]".format(serverId)) print("Failed deactivating server [{0}]".format(serverId)) return False self.ioc.getLogger().warning( "Server [{0}] preparing to reallocate detect jobs".format( serverId)) if not self.monitor.rescheduleDetectJobs(serverId): self.ioc.getLogger().error( "Failed rescheduling detect jobs [{0}]".format(serverId)) print("Failed rescheduling detect jobs [{0}]".format(serverId)) return False self.ioc.getLogger().warning( "Server [{0}] preparing to reallocate schedule jobs".format( serverId)) if not self.monitor.rescheduleScheduleJobs(serverId): self.ioc.getLogger().error( "Failed rescheduling detect jobs [{0}]".format(serverId)) print("Failed rescheduling detect jobs [{0}]".format(serverId)) return False self.ioc.getLogger().warning( "Server [{0}] preparing to reallocate jobs".format(serverId)) if not self.monitor.rescheduleJobs(serverId): self.ioc.getLogger().error( "Failed rescheduling detect jobs [{0}]".format(serverId)) print("Failed rescheduling detect jobs [{0}]".format(serverId)) return False print("Server Deactivated") return True def action_activate(self, node=None): """activates server in cluster Args: node (str): MongoDB ObjectId to activate; defaults to local node Returns: bool: If activation is successful """ if not self.ioc.ensureRegistration(): self.ioc.getLogger().error("Server not registered with MongoDB") print("Unregistered servers cannot talk to the cluster") return False valid, serverId = self.valid_server(node) if not valid: print("Invalid ObjectID") return False if self.ioc.getCollection('JobServer').update_one( { '_id': ObjectId(serverId) }, { '$set': { 'active': True, 'activationTime': datetime.datetime.utcnow() } }).modified_count < 1: self.ioc.getLogger().warning( "Server [{0}] failed to be activated".format(serverId)) return False self.ioc.getLogger().warning("Server [{0}] activated".format(serverId)) return True def valid_server(self, node=None): """Validates node is in the MongoDB instance connected to Args: node (str): MongoDB Object ID to validate; defaults to local node Returns: tuple: first element is boolean if valid second is objectId as string """ if node: try: server = self.ioc.getCollection('JobServer').find_one( {'_id': ObjectId(str(node))}) except InvalidId: self.ioc.getLogger().error( "Invalid ObjectID passed to bridge info [{0}]".format( node)) return False, "" if server: return True, dict(server).get('_id') self.ioc.getLogger().error( "Failed to find server [{0}] in the database".format(node)) return False, "" return True, self.ioc.getConfig().NodeIdentity
class Scan(object): """Scanning class for GREASE Scanner This is the model to actually utilize the scanners to parse the configured environments Attributes: ioc (GreaseContainer): IOC for scanning conf (PrototypeConfig): Prototype configuration instance impTool (ImportTool): Import Utility Instance dedup (Deduplication): Deduplication instance to be used """ def __init__(self, ioc=None): if ioc and isinstance(ioc, GreaseContainer): self.ioc = ioc else: self.ioc = GreaseContainer() self.conf = PrototypeConfig(self.ioc) self.impTool = ImportTool(self.ioc.getLogger()) self.dedup = Deduplication(self.ioc) self.scheduler = Scheduling(self.ioc) def Parse(self, source=None, config=None): """This will read all configurations and attempt to scan the environment This is the primary business logic for scanning in GREASE. This method will use configurations to parse the environment and attempt to schedule Note: If a Source is specified then *only* that source is parsed. If a configuration is set then *only* that configuration is parsed. If both are provided then the configuration will *only* be parsed if it is of the source provided Note: **If mocking is enabled**: Deduplication *will not occur* Args: source (str): If set will only parse for the source listed config (str): If set will only parse the specified config Returns: bool: True unless error """ self.ioc.getLogger().trace("Starting Parse of Environment", trace=True) Configuration = self.generate_config_set(source=source, config=config) for conf in Configuration: inst = self.impTool.load(conf.get('source', str(uuid4()))) if not isinstance(inst, BaseSourceClass): self.ioc.getLogger().error("Invalid Source [{0}]".format( conf.get('source')), notify=False) del inst continue else: # If mock mode enabled if self.ioc.getConfig().get('Sourcing', 'mock'): data = inst.mock_data(conf) # else actually do sourcing else: if inst.parse_source(conf): # deduplicate data data = self.dedup.Deduplicate( data=inst.get_data(), source=conf.get('source'), threshold=inst.deduplication_strength, expiry_hours=inst.deduplication_expiry, expiry_max=inst.deduplication_expiry_max, collection='Dedup_Sourcing', field_set=inst.field_set) else: self.ioc.getLogger().warning( "Source [{0}] parsing failed".format( conf.get('source')), notify=False) data = [] if len(data) > 0: if self.scheduler.scheduleDetection( conf.get('source'), conf.get('name'), data): self.ioc.getLogger().info( "Data scheduled for detection from source [{0}]". format(conf.get('source')), trace=True) del inst continue else: self.ioc.getLogger().error( "Scheduling failed for source document!", notify=False) del inst continue else: self.ioc.getLogger().trace( "Length of data was empty; was not scheduled", trace=True) del inst continue return True def generate_config_set(self, source=None, config=None): """Examines configuration and returns list of configs to parse Note: If a Source is specified then *only* that source is parsed. If a configuration is set then *only* that configuration is parsed. If both are provided then the configuration will *only* be parsed if it is of the source provided Args: source (str): If set will only parse for the source listed config (str): If set will only parse the specified config Returns: list[dict]: Returns Configurations to Parse for data """ ConfigList = [] if source and config: if self.conf.get_config(config).get('source') == source: ConfigList.append(self.conf.get_config(config)) return ConfigList else: self.ioc.getLogger().warning( "Configuration [{0}] Not Found With Correct Source [{1}]". format(config, source), trace=True, notify=False) elif source and not config: if source in self.conf.get_sources(): for configuration in self.conf.get_source(source): ConfigList.append(configuration) return ConfigList else: self.ioc.getLogger().warning( "Source not found in Configuration [{0}]".format(source), trace=True, notify=False) elif not source and config: if self.conf.get_config(config): ConfigList.append(self.conf.get_config(config)) return ConfigList else: self.ioc.getLogger().warning( "Config not found in Configuration [{0}]".format(config), trace=True, notify=False) else: ConfigList = self.conf.getConfiguration().get('raw') return ConfigList