243 lines
No EOL
11 KiB
Python
243 lines
No EOL
11 KiB
Python
|
|
import ops
|
|
import ops.db
|
|
import os.path
|
|
import random
|
|
import pickle
|
|
import dsz
|
|
import sqlite3
|
|
|
|
def get_job(session):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT rowid FROM scansweep_queue WHERE session=? AND inprogress=? AND complete=?', query_list=[session, 'False', 'False'], no_return=False, max_tries=10)
|
|
job_list = []
|
|
for row in curs:
|
|
job_list.append(row)
|
|
if (len(job_list) == 0):
|
|
return False
|
|
random.shuffle(job_list)
|
|
job = job_list.pop()
|
|
query_execute_wrapper(conn, query_string="UPDATE scansweep_queue SET inprogress='True' WHERE session=? AND rowid=? AND inprogress=? AND complete=?", query_list=[session, job['rowid'], 'False', 'False'], max_tries=10)
|
|
curs = query_execute_wrapper(conn, query_string='SELECT job,target FROM scansweep_queue WHERE rowid=?', query_list=[job['rowid']], no_return=False, max_tries=10)
|
|
for row in curs:
|
|
return (row['job'], row['target'])
|
|
|
|
def write_job(session, job, target):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT count(*) FROM scansweep_queue WHERE session=? AND job=? AND target=?', query_list=[session, job, target], no_return=False)
|
|
if (int(curs.fetchone()['count(*)']) > 0):
|
|
return False
|
|
query_execute_wrapper(conn, query_string='INSERT INTO scansweep_queue (session, job, target, inprogress, complete) VALUES (?,?,?,?,?)', query_list=[session, job, target, 'False', 'False'])
|
|
return True
|
|
|
|
def write_job_list(session, job_list):
|
|
conn = get_database_conn()
|
|
with conn:
|
|
curs = conn.executemany(('INSERT INTO scansweep_queue (session, job, target, inprogress, complete) VALUES ("%s",?,?,"False","False")' % session), job_list)
|
|
return True
|
|
|
|
def reset_jobs(session):
|
|
conn = get_database_conn()
|
|
query_execute_wrapper(conn, query_string="UPDATE scansweep_queue SET inprogress='False' WHERE session=? AND inprogress=? AND complete=?", query_list=[session, 'True', 'False'])
|
|
query_execute_wrapper(conn, query_string="UPDATE scansweep_metadata SET kill='False' WHERE session=?", query_list=[session])
|
|
return True
|
|
|
|
def list_jobs(session, inprogress='False'):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string="SELECT * FROM scansweep_queue WHERE session=? AND inprogress=? AND complete='False'", query_list=[session, inprogress], no_return=False)
|
|
job_list = []
|
|
for row in curs:
|
|
job_list.append({'job': row['job'], 'target': row['target']})
|
|
return job_list
|
|
|
|
def num_jobs(session, inprogress='False', type=None):
|
|
conn = get_database_conn()
|
|
if (type is None):
|
|
curs = query_execute_wrapper(conn, query_string="SELECT count(*) FROM scansweep_queue WHERE session=? AND inprogress=? AND complete='False'", query_list=[session, inprogress], no_return=False)
|
|
else:
|
|
curs = query_execute_wrapper(conn, query_string="SELECT count(*) FROM scansweep_queue WHERE session=? AND job LIKE ? AND inprogress=? AND complete='False'", query_list=[session, ('%s%%' % type), inprogress], no_return=False)
|
|
num = curs.fetchone()['count(*)']
|
|
return num
|
|
|
|
def all_num_jobs(session):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT job,inprogress,complete FROM scansweep_queue WHERE session=?', query_list=[session], no_return=False)
|
|
num = []
|
|
for row in curs:
|
|
num.append({'type': row['job'].split('|')[0], 'inprogress': row['inprogress'], 'complete': row['complete']})
|
|
return num
|
|
|
|
def set_jobtype(session, type):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='INSERT INTO scansweep_jobtype (session, type) VALUES (?,?)', query_list=[session, type])
|
|
return True
|
|
|
|
def get_jobtypes(session):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT type FROM scansweep_jobtype WHERE session=?', query_list=[session], no_return=False)
|
|
type_dict = {}
|
|
for row in curs:
|
|
type_dict[row['type']] = True
|
|
return type_dict
|
|
|
|
def get_metadata(scansweep_env, var=None):
|
|
conn = get_database_conn()
|
|
if (var is None):
|
|
curs = query_execute_wrapper(conn, query_string='SELECT * FROM scansweep_metadata WHERE scansweep_env=?', query_list=[scansweep_env], no_return=False)
|
|
metadata = curs.fetchone()
|
|
else:
|
|
curs = query_execute_wrapper(conn, query_string=('SELECT %s FROM scansweep_metadata WHERE scansweep_env=?' % var), query_list=[scansweep_env], no_return=False)
|
|
metadata = curs.fetchone()
|
|
if (metadata is not None):
|
|
metadata = metadata[var]
|
|
else:
|
|
return ''
|
|
return metadata
|
|
|
|
def write_metadata(scansweep_env, session, scansweep_logfile, scansweep_results, verbose):
|
|
conn = get_database_conn()
|
|
query_execute_wrapper(conn, query_string='INSERT INTO scansweep_metadata (scansweep_env, session, scansweep_logfile, scansweep_results, verbose, kill) VALUES (?,?,?,?,?,?)', query_list=[scansweep_env, session, scansweep_logfile, scansweep_results, verbose, 'False'])
|
|
return True
|
|
|
|
def kill_session(session):
|
|
conn = get_database_conn()
|
|
query_execute_wrapper(conn, query_string="UPDATE scansweep_metadata SET kill='True' WHERE session=?", query_list=[session])
|
|
return True
|
|
|
|
def check_kill(session):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT kill FROM scansweep_metadata WHERE session=?', query_list=[session], no_return=False)
|
|
kill_data = curs.fetchone()
|
|
if (kill_data['kill'] == 'True'):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def get_excludes(session):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT * FROM scansweep_excludes WHERE session=?', query_list=[session], no_return=False)
|
|
excludes_list = []
|
|
for row in curs:
|
|
excludes_list.append(row['target'])
|
|
if (len(excludes_list) == 0):
|
|
return []
|
|
return excludes_list
|
|
|
|
def write_excludes(session, target):
|
|
conn = get_database_conn()
|
|
query_execute_wrapper(conn, query_string='INSERT INTO scansweep_excludes (session, target) VALUES (?,?)', query_list=[session, target])
|
|
return True
|
|
|
|
def write_excludes_list(session, target_list):
|
|
conn = get_database_conn()
|
|
input_list = []
|
|
for target in target_list:
|
|
input_list.append([target])
|
|
with conn:
|
|
curs = conn.executemany(('INSERT INTO scansweep_excludes (session, target) VALUES ("%s",?)' % session), input_list)
|
|
return True
|
|
|
|
def list_sessions():
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT session FROM scansweep_metadata', no_return=False)
|
|
session_list = []
|
|
for row in curs:
|
|
session_list.append(row['session'])
|
|
return session_list
|
|
|
|
def write_escalate_rule(session, rule):
|
|
conn = get_database_conn()
|
|
pickled_data = pickle.dumps(rule)
|
|
query_execute_wrapper(conn, query_string='INSERT INTO scansweep_escalate (session, rule) VALUES (?,?)', query_list=[session, pickled_data])
|
|
return True
|
|
|
|
def delete_escalate_rule(session, rule):
|
|
conn = get_database_conn()
|
|
pickled_data = pickle.dumps(rule)
|
|
query_execute_wrapper(conn, query_string='DELETE FROM scansweep_escalate WHERE session=? AND rule=?', query_list=[session, pickled_data])
|
|
return True
|
|
|
|
def get_escalate_rules(session):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT * FROM scansweep_escalate WHERE session=?', query_list=[session], no_return=False)
|
|
rule_list = []
|
|
for row in curs:
|
|
unpickled_data = pickle.loads(row['rule'])
|
|
rule_list.append(unpickled_data)
|
|
return rule_list
|
|
|
|
def check_escalation(session):
|
|
rule_list = get_escalate_rules(session)
|
|
if (len(rule_list) > 0):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def get_results(session, type, success=True):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT * FROM scansweep_results WHERE session=? AND type=? AND success=?', query_list=[session, type, success], no_return=False)
|
|
results_list = []
|
|
for row in curs:
|
|
unpickled_data = pickle.loads(row['data'])
|
|
results_list.append({'target': row['target'], 'data': unpickled_data})
|
|
return results_list
|
|
|
|
def num_results(session, type, success):
|
|
conn = get_database_conn()
|
|
curs = query_execute_wrapper(conn, query_string='SELECT count(*) FROM scansweep_results WHERE session=? AND type=? AND success=?', query_list=[session, type, success], no_return=False)
|
|
num = curs.fetchone()['count(*)']
|
|
return num
|
|
|
|
def write_result(session, type, target, data, success, job):
|
|
conn = get_database_conn()
|
|
pickled_data = pickle.dumps(data)
|
|
query_execute_wrapper(conn, query_string='INSERT INTO scansweep_results (session, type, target, data, success) VALUES (?,?,?,?,?)', query_list=[session, type, target, pickled_data, success], max_tries=10)
|
|
query_execute_wrapper(conn, query_string="UPDATE scansweep_queue SET inprogress='False', complete='True' WHERE session=? AND job=? AND target=? AND inprogress=? AND complete=?", query_list=[session, job, target, 'True', 'False'], max_tries=10)
|
|
return True
|
|
|
|
def get_database_conn():
|
|
scansweepdb = get_database()
|
|
return scansweepdb.connection
|
|
|
|
def get_database():
|
|
scansweep_db_file = os.path.join(ops.PROJECTLOGDIR, 'scansweep.db')
|
|
scansweepdb = ops.db.get_this_db(scansweep_db_file)
|
|
return scansweepdb
|
|
|
|
def query_execute_wrapper(db_conn, query_string=None, query_list=None, max_tries=3, no_return=True):
|
|
for i in range(0, max_tries):
|
|
try:
|
|
with db_conn:
|
|
if (query_list is None):
|
|
curs = db_conn.execute(query_string)
|
|
else:
|
|
curs = db_conn.execute(query_string, query_list)
|
|
if no_return:
|
|
return None
|
|
else:
|
|
return curs
|
|
except sqlite3.OperationalError:
|
|
dsz.Sleep(((1 + i) * random.randint(250, 1250)))
|
|
continue
|
|
raise
|
|
|
|
def setup_db():
|
|
scansweepdb = get_database()
|
|
scansweepdb.ensureTable('scansweep_queue', 'CREATE TABLE scansweep_queue (session, job, target, inprogress, complete)')
|
|
scansweepdb.ensureTable('scansweep_results', 'CREATE TABLE scansweep_results (session, type, target, data, success)')
|
|
scansweepdb.ensureTable('scansweep_metadata', 'CREATE TABLE scansweep_metadata (scansweep_env, session, scansweep_logfile, scansweep_results, verbose, kill)')
|
|
scansweepdb.ensureTable('scansweep_escalate', 'CREATE TABLE scansweep_escalate (session, rule)')
|
|
scansweepdb.ensureTable('scansweep_jobtype', 'CREATE TABLE scansweep_jobtype (session, type)')
|
|
scansweepdb.ensureTable('scansweep_excludes', 'CREATE TABLE scansweep_excludes (session, target)')
|
|
return True
|
|
|
|
def delete_db():
|
|
conn = get_database_conn()
|
|
with conn:
|
|
conn.execute('DROP TABLE if exists scansweep_queue')
|
|
conn.execute('DROP TABLE if exists scansweep_results')
|
|
conn.execute('DROP TABLE if exists scansweep_metadata')
|
|
conn.execute('DROP TABLE if exists scansweep_escalate')
|
|
conn.execute('DROP TABLE if exists scansweep_jobtype')
|
|
conn.execute('DROP TABLE if exists scansweep_excludes')
|
|
return True |