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 test_pkg_load_bad(self):
ioc = GreaseContainer()
# clean up
for root, dirnames, filenames in os.walk(pkg_resources.resource_filename('tgt_grease.enterprise.Model', 'config/')):
for filename in fnmatch.filter(filenames, '*.config.json'):
self.assertIsNone(os.remove(os.path.join(root, filename)))
configList = [
{
"name": "test1",
"job": "fakeJob",
"exe_env": "windows",
"source": "swapi",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "badtest1",
"exe_env": "windows",
"source": "stackOverflow",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "test3",
"job": "fakeJob",
"exe_env": "windows",
"source": "Google",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
],
"exists": [
{
"field": "var"
}
]
}
}
]
GoodConfigList = [
{
"name": "test1",
"job": "fakeJob",
"exe_env": "windows",
"source": "swapi",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "test3",
"job": "fakeJob",
"exe_env": "windows",
"source": "Google",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
],
"exists": [
{
"field": "var"
}
]
}
}
]
i = 0
for conf in configList:
with open(pkg_resources.resource_filename('tgt_grease.enterprise.Model', 'config/') + 'conf{0}.config.json'.format(i), 'w') as fil:
fil.write(json.dumps(conf, indent=4))
i += 1
conf = PrototypeConfig(ioc)
conf.load(reloadConf=True)
self.assertEqual(len(conf.getConfiguration().get('configuration').get('pkg')), len(GoodConfigList))
self.assertEqual(len(conf.getConfiguration().get('raw')), len(GoodConfigList))
self.assertEqual(len(conf.getConfiguration().get('source').get('swapi')), 1)
self.assertEqual(len(conf.getConfiguration().get('source').get('Google')), 1)
self.assertEqual(2, len(conf.get_sources()))
self.assertEqual(2, len(conf.get_names()))
self.assertEqual(len(conf.get_source('Google')), 1)
self.assertTrue(isinstance(conf.get_config('test1'), dict))
self.assertTrue(conf.get_config('test1'))
# clean up
for root, dirnames, filenames in os.walk(pkg_resources.resource_filename('tgt_grease.enterprise.Model', 'config/')):
for filename in fnmatch.filter(filenames, '*.config.json'):
self.assertIsNone(os.remove(os.path.join(root, filename)))
# clear the config
conf.load(reloadConf=True)
def test_mongo_load_bad(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)))
configList = [
{
"name": "test1",
"job": "fakeJob",
"exe_env": "windows",
"source": "swapi",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "badtest1",
"exe_env": "windows",
"source": "stackOverflow",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "test3",
"job": "fakeJob",
"exe_env": "windows",
"source": "Google",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
],
"exists": [
{
"field": "var"
}
]
}
}
]
GoodConfigList = [
{
"name": "test1",
"job": "fakeJob",
"exe_env": "windows",
"source": "swapi",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "test3",
"job": "fakeJob",
"exe_env": "windows",
"source": "Google",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
],
"exists": [
{
"field": "var"
}
]
}
}
]
for conf in configList:
ioc.getCollection('Configuration').insert_one(conf)
ioc.getCollection('Configuration').update_many({}, {'$set': {'active': True, 'type': 'prototype_config'}})
# sleep because travis is slow sometimes
time.sleep(1.5)
conf = PrototypeConfig(ioc)
conf.load(reloadConf=True)
self.assertEqual(len(conf.getConfiguration().get('configuration').get('mongo')), len(GoodConfigList))
self.assertEqual(len(conf.getConfiguration().get('raw')), len(GoodConfigList))
self.assertEqual(len(conf.getConfiguration().get('source').get('swapi')), 1)
self.assertEqual(len(conf.getConfiguration().get('source').get('Google')), 1)
self.assertEqual(2, len(conf.get_sources()))
self.assertEqual(2, len(conf.get_names()))
self.assertEqual(len(conf.get_source('Google')), 1)
self.assertTrue(isinstance(conf.get_config('test1'), dict))
self.assertTrue(conf.get_config('test1'))
# clean up
ioc.getCollection('Configuration').drop()
# clear the config
conf.load(reloadConf=True)
def test_all_load_bad(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)))
# clean up
for root, dirnames, filenames in os.walk(pkg_resources.resource_filename('tgt_grease.enterprise.Model', 'config/')):
for filename in fnmatch.filter(filenames, '*.config.json'):
self.assertIsNone(os.remove(os.path.join(root, filename)))
configList = [
{
"name": "test1",
"job": "fakeJob",
"exe_env": "windows",
"source": "swapi",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "badtest1",
"exe_env": "windows",
"source": "stackOverflow",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "test3",
"job": "fakeJob",
"exe_env": "windows",
"source": "Google",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
],
"exists": [
{
"field": "var"
}
]
}
}
]
GoodConfigList = [
{
"name": "test1",
"job": "fakeJob",
"exe_env": "windows",
"source": "swapi",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
]
}
},
{
"name": "test3",
"job": "fakeJob",
"exe_env": "windows",
"source": "Google",
"logic": {
"regex": [
{
"field": "character",
"pattern": ".*skywalker.*"
}
],
"exists": [
{
"field": "var"
}
]
}
}
]
i = 0
length = len(configList) - 1
while i <= length:
if i == 0:
with open(ioc.getConfig().get('Configuration', 'dir') + 'conf{0}.config.json'.format(i), 'w') as fil:
fil.write(json.dumps(configList[i], indent=4))
if i == 1:
with open(pkg_resources.resource_filename('tgt_grease.enterprise.Model',
'config/') + 'conf{0}.config.json'.format(i), 'w') as fil:
fil.write(json.dumps(configList[i], indent=4))
if i == 2:
ioc.getCollection('Configuration').insert_one(configList[i])
i += 1
ioc.getCollection('Configuration').update_many({}, {'$set': {'active': True, 'type': 'prototype_config'}})
# sleep because travis is slow
time.sleep(1.5)
conf = PrototypeConfig(ioc)
conf.load(reloadConf=True)
self.assertEqual(len(conf.getConfiguration().get('configuration').get('mongo')), 1)
self.assertEqual(len(conf.getConfiguration().get('configuration').get('pkg')), 0)
self.assertEqual(len(conf.getConfiguration().get('configuration').get('fs')), 1)
self.assertEqual(len(conf.getConfiguration().get('raw')), len(GoodConfigList))
self.assertEqual(len(conf.getConfiguration().get('source').get('swapi')), 1)
self.assertEqual(len(conf.getConfiguration().get('source').get('Google')), 1)
self.assertEqual(2, len(conf.get_names()))
self.assertEqual(len(conf.get_source('Google')), 1)
self.assertTrue(isinstance(conf.get_config('test1'), dict))
self.assertTrue(conf.get_config('test1'))
# clean up
ioc.getCollection('Configuration').drop()
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)))
# clean up
for root, dirnames, filenames in os.walk(pkg_resources.resource_filename('tgt_grease.enterprise.Model', 'config/')):
for filename in fnmatch.filter(filenames, '*.config.json'):
self.assertIsNone(os.remove(os.path.join(root, filename)))
# clear the config
conf.load(reloadConf=True)