def test_task():
task = Task(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13)
need_set = set([
"path",
"category",
"package",
"timeout",
"priority",
"options",
"machine",
"platform",
"tags",
"custom",
"memory",
"clock",
"enforce_timeout",
"main_task_id",
"retrieved",
"route",
])
have_set = set(dir(task))
assert need_set & have_set == need_set
def submit_tasks(self, node_id, pend_tasks_num):
db = session()
node = db.query(Node).filter_by(id=node_id).first()
if node.name != "master":
# don't do nothing if nothing in pending
# Get tasks from main_db submitted through web interface
main_db_tasks = main_db.list_tasks(status=TASK_PENDING,
order_by=desc("priority"))
if main_db_tasks:
log.debug("going to upload {} tasks to node id {}".format(
pend_tasks_num, node_id))
limit = 0
for t in main_db_tasks:
# convert the options string to a dict, e.g. {'opt1': 'val1', 'opt2': 'val2', 'opt3': 'val3'}
if "node=" in t.options:
try:
options = dict((value.strip()
for value in option.split("=", 1))
for option in t.options.split(","))
if options.get("node", node.name) != node.name:
# skip task
continue
except Exception as e:
log.error(e)
tasks = db.query(Task).filter_by(main_task_id=t.id).all()
if not tasks:
# Check if file exist, if no wipe from db and continue, rare cases
if not os.path.exists(t.target):
main_db.delete_task(t.id)
log.info(
"Task id: {} - File doesn't exist: {}".format(
t.id, t.target))
continue
# Convert array of tags into comma separated list
tags = ','.join([tag.name for tag in t.tags])
# Append a comma, to make LIKE searches more precise
if tags: tags += ','
if "msoffice-crypt-tmp" in t.target and "password="******"password="******"pwd=")
args = dict(package=t.package,
category=t.category,
timeout=t.timeout,
priority=t.priority,
options=t.options +
",main_task_id={}".format(t.id),
machine=t.machine,
platform=t.platform,
tags=tags,
custom=t.custom,
memory=t.memory,
clock=t.clock,
enforce_timeout=t.enforce_timeout,
main_task_id=t.id)
task = Task(path=t.target, **args)
db.add(task)
try:
db.commit()
except Exception as e:
#main_db.set_status(t.id, TASK_FAILED_REPORTING)
db.rollback()
log.info(e)
continue
if reporting_conf.distributed.enable_tags:
# Only get tasks that have not been pushed yet.
q = db.query(Task).filter(
or_(Task.node_id == None,
Task.task_id == None),
Task.finished == False)
# Order by task priority and task id.
q = q.order_by(-Task.priority, Task.main_task_id)
# Get available node tags
machines = db.query(Machine).filter_by(
node_id=node.id).all()
# Get available tag combinations
ta = set()
for m in machines:
for i in xrange(1, len(m.tags) + 1):
for tag in combinations(m.tags, i):
ta.add(','.join(tag))
ta = list(ta)
# Create filter query from tasks in ta
tags = [getattr(Task, "tags") == ""]
for tg in ta:
if len(tg.split(',')) == 1:
tags.append(
getattr(Task, "tags") == (tg + ','))
else:
tg = tg.split(',')
# ie. LIKE '%,%,%,'
t_combined = [
getattr(Task, "tags").like(
"%s" % ('%,' * len(tg)))
]
for tag in tg:
t_combined.append(
getattr(Task, "tags").like(
"%%%s%%" % (tag + ',')))
tags.append(and_(*t_combined))
# Filter by available tags
q = q.filter(or_(*tags))
# Submit appropriate tasks to node
if pend_tasks_num > 0:
for task in q.limit(pend_tasks_num).all():
submitted = node_submit_task(
task.id, node.id)
if submitted:
main_db.set_status(t.id, TASK_RUNNING)
db.close()
return
else:
# Submit appropriate tasks to node
submitted = node_submit_task(task.id, node.id)
if submitted:
main_db.set_status(t.id, TASK_RUNNING)
limit += 1
if limit == pend_tasks_num:
db.close()
return
else:
for task in tasks:
log.info(
"Deleting incorrectly uploaded file from dist db, main_task_id: {}"
.format(t.id))
db.delete(task)
db.commit()
db.close()
def submit_tasks(self, node_id, pend_tasks_num, options_like=False, force_push_push=False, db=None):
# HACK do not create a new session if the current one (passed as parameter) is still valid.
try:
node = db.query(Node).filter_by(name=node_id).first()
except (OperationalError, SQLAlchemyError) as e:
log.warning(
"Got an operational Exception when trying to submit tasks: {}".format(e))
return False
limit = 0
if node.name != "master":
# don't do nothing if nothing in pending
# Get tasks from main_db submitted through web interface
main_db_tasks = main_db.list_tasks(status=TASK_PENDING, order_by=desc("priority"), options_like=options_like, limit=pend_tasks_num)
if not main_db_tasks:
return True
if main_db_tasks:
for t in main_db_tasks:
force_push = False
try:
# convert the options string to a dict, e.g. {'opt1': 'val1', 'opt2': 'val2', 'opt3': 'val3'}
options = get_options(t.options)
# check if node exist and its correct
if "node=" in t.options:
requested_node = options.get("node")
if requested_node not in STATUSES:
# if the requested node is not available
force_push = True
elif requested_node != node.name:
# otherwise keep looping
continue
if "timeout=" in t.options:
t.timeout = options["timeout"]
except Exception as e:
log.error(e)
# wtf are you doing in pendings?
tasks = db.query(Task).filter_by(main_task_id=t.id).all()
if tasks:
for task in tasks:
print(task.id, task.task_id,task.main_task_id, task.node_id)
#log.info("Deleting incorrectly uploaded file from dist db, main_task_id: {}".format(t.id))
#db.delete(task)
#db.commit()
if node.name == "master":
main_db.set_status(t.id, TASK_RUNNING)
else:
main_db.set_status(t.id, TASK_DISTRIBUTED)
continue
# Check if file exist, if no wipe from db and continue, rare cases
if t.category in ("file", "pcap", "static") and not os.path.exists(t.target):
log.info("Task id: {} - File doesn't exist: {}".format(t.id, t.target))
main_db.delete_task(t.id)
continue
# Convert array of tags into comma separated list
tags = ','.join([tag.name for tag in t.tags])
# Append a comma, to make LIKE searches more precise
if tags:
tags += ','
if "msoffice-crypt-tmp" in t.target and "password="******"password="******"pwd=")
args = dict(package=t.package, category=t.category, timeout=t.timeout, priority=t.priority,
options=t.options+",main_task_id={}".format(t.id), machine=t.machine, platform=t.platform,
tags=tags, custom=t.custom, memory=t.memory, clock=t.clock,
enforce_timeout=t.enforce_timeout, main_task_id=t.id)
task = Task(path=t.target, **args)
db.add(task)
try:
db.commit()
except Exception as e:
log.exception(e)
log.info("TASK_FAILED_REPORTING")
db.rollback()
log.info(e)
continue
if force_push or force_push_push:
# Submit appropriate tasks to node
submitted = node_submit_task(task.id, node.id)
if submitted:
if node.name == "master":
main_db.set_status(t.id, TASK_RUNNING)
else:
main_db.set_status(t.id, TASK_DISTRIBUTED)
limit += 1
if limit == pend_tasks_num or limit == len(main_db_tasks):
return True
# Only get tasks that have not been pushed yet.
q = db.query(Task).filter(
or_(Task.node_id == None, Task.task_id == None), Task.finished == False)
if q is None:
return True
# Order by task priority and task id.
q = q.order_by(-Task.priority, Task.main_task_id)
# if we have node set in options push
if reporting_conf.distributed.enable_tags:
# Create filter query from tasks in ta
tags = [getattr(Task, "tags") == ""]
for tg in SERVER_TAGS[node.name]:
if len(tg.split(',')) == 1:
tags.append(getattr(Task, "tags") == (tg + ','))
else:
tg = tg.split(',')
# ie. LIKE '%,%,%,'
t_combined = [getattr(Task, "tags").like(
"%s" % ('%,' * len(tg)))]
for tag in tg:
t_combined.append(
getattr(Task, "tags").like("%%%s%%" % (tag + ',')))
tags.append(and_(*t_combined))
# Filter by available tags
q = q.filter(or_(*tags))
to_upload = q.limit(pend_tasks_num).all()
if not to_upload:
return True
# Submit appropriate tasks to node
log.debug("going to upload {} tasks to node {}".format(
pend_tasks_num, node.name))
for task in to_upload:
submitted = node_submit_task(task.id, node.id)
if submitted:
if node.name == "master":
main_db.set_status(task.main_task_id, TASK_RUNNING)
else:
main_db.set_status(
task.main_task_id, TASK_DISTRIBUTED)
else:
print(
"something is wrong with submission of task: {}".format(task.id))
db.delete(task)
db.commit()
limit += 1
if limit == pend_tasks_num:
return True
return True