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
