Browse Source

nearly complete cjs2 database api

automate2
Stephen Lorenz 4 years ago
parent
commit
62137640fb
  1. 242
      cjs/core/database.py
  2. 68
      cjs/database.py
  3. 364
      cjs/debug.py
  4. 2
      cjs/utils/cli.py

242
cjs/core/database.py

@ -104,6 +104,35 @@ class Database:
try:
q = sess.query(Table).filter(expr).all()
return q
except:
# TODO: implement fine-grained handling
raise
@_pass_session
def first_record(self, sess, Table, expr=True):
'''Get first record in the given Table.'''
try:
q = sess.query(Table).filter(expr).first()
return q
except:
# TODO: implement fine-grained handling
raise
@_pass_session
def last_record(self, sess, Table, expr=True):
'''Get last record in the given Table.'''
try:
q = sess.query(Table).filter(expr).order_by(Table.id_.desc()).first()
return q
except:
# TODO: implement fine-grained handling
raise
@_pass_session
def count_table(self, sess, Table, expr=True):
try:
q = sess.query(Table).filter(expr).count()
return q
except:
# TODO: implement fine-grained handling
raise
@ -134,6 +163,54 @@ def _get_database(name):
# TODO: implement fine-grained handling
raise
# END: Database object functions
# END: Internal functions
# BEGIN: External functions
# BEGIN: Database State functions
def add_database(name, location, init=True):
try:
ds = DatabaseState(name, location)
add_state(ds)
# TODO: replace with a better method of doing this
if init:
db = _get_database(name)
db.initialize()
except:
# TODO: implement fine-grained handling
raise
def remove_database(name):
try:
remove_state(DatabaseState,
DatabaseState.name == name)
except:
# TODO: implement fine-grained handling
raise
def update_database(name):
try:
pass
except:
# TODO: implement fine-grained handling
raise
def list_database():
try:
# select_state's expr parameter is true by default
# and will return everything in the state table if
# not passed another value
ds = select_state(DatabaseState)
return ds
except:
# TODO: implement fine-grained handling # TODO: implement fine-grained handling
raise
# END: Database State functions
# BEGIN: Database object functions
# TODO: name is easily confused with add_database, consider refactoring
def _database_add(name, Tables):
try:
@ -152,7 +229,6 @@ def _database_compound_add(name, Table, Tables):
# TODO: implement fine-grained handling
raise
# TODO: name is easily confused with remove_database, consider refactoring
def _database_remove(name, Table, expr=True):
try:
@ -184,164 +260,72 @@ def _database_select(name, Table, expr=True):
# TODO: implement fine-grained handling
raise
# END: Database object functions
# END: Internal functions
# BEGIN: External functions
# BEGIN: Database State functions
def add_database(name, location, init=True):
def database_location(name):
try:
ds = DatabaseState(name, location)
add_state(ds)
# TODO: replace with a better method of doing this
if init:
db = _get_database(name)
db.initialize()
ds = _get_database(name)
return ds.location
except:
# TODO: implement fine-grained handling
raise
def remove_database(name):
def database_count(name, Table, expr=True):
try:
remove_state(DatabaseState,
DatabaseState.name == name)
ds = _get_database(name)
return ds.count_table(Table, expr)
except:
# TODO: implement fine-grained handling
raise
def update_database(name):
def database_last(name, Table, expr=True):
try:
pass
ds = _get_database(name)
return ds.last_record(Table, expr)
except:
# TODO: implement fine-grained handling
raise
def list_database():
def database_first(name, Table, expr=True):
try:
# select_state's expr parameter is true by default
# and will return everything in the state table if
# not passed another value
ds = select_state(DatabaseState)
return ds
ds = _get_database(name)
return ds.first_record(Table, expr)
except:
# TODO: implement fine-grained handling
raise
# END: Database State functions
# BEGIN: Database object functions
# END: Database object functions
# BEGIN: Automation API
def database_location(name):
def pending_jobs(name):
try:
ds = _get_database(name)
return ds.location
pass
except:
# TODO: implement fine-grained handling
raise
# END: Database object functions
# END: External functions
def tester():
name = 'debug'
#location = '.data/my_database'
#add_database(name, location)
#print(list_database())
#database_add(name, [EvidenceTable('my_file', '/path/to/file')])
print(_database_select(name, EventTable, EventTable.name == 'Pending')[0].name)
#tester()
'''
def upload(name, evidence, comparisons):
def next_job(name):
try:
db = _get_database(name)
# format
ds = _get_database(name)
pass
except:
# TODO: implement fine-grained handling
raise
evidence_name = evidence['name']
evidence_file = '%s/input/%s.json' % (db.location, evidence_name)
evidence_file = str(abs_path(evidence_file))
write_json(evidence_file, evidence)
evidence_id = vault.insert(VaultEvidence(evidence_name, evidence_file))
'''
'''
def stage(name, batch):
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
batch_id = vault.insert(VaultBatch(batch))
tests = vault.get_tests()
l = [VaultJob(t.id, batch_id) for t in tests]
vault.insert(l)
def current_batch(name):
try:
ds = _get_database(name)
pass
except:
# TODO: implement fine-grained handling
raise
def insert_result(name, filename):
try:
ds = _get_database(name)
pass
except:
# TODO: implement fine-grained handling
raise
for c in comparisons:
comparison_name = c['name']
comparison_file = '%s/input/%s.json' % (vault_root, comparison_name)
comparison_file = str(abs_path(comparison_file))
write_json(comparison_file, c)
comparison_id = vault.insert(VaultComparison(comparison_name, comparison_file))
vault.insert(VaultTest(evidence_id, comparison_id))
def next_job(name):
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
s = vault.get_session()
job_row = s.query(VaultJob).filter_by(event_id=1).first()
test_row = s.query(VaultTest).filter_by(id=job_row.test_id).first()
evidence = s.query(VaultEvidence).filter_by(id=test_row.evidence_id).first()
comparison = s.query(VaultComparison).filter_by(id=test_row.comparison_id).first()
s.close()
return [evidence, comparison]
def all_jobs(name):
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
s = vault.get_session()
jobs = []
for job_row in s.query(VaultJob).filter_by(event_id=1).all():
test_row = s.query(VaultTest).filter_by(id=job_row.test_id).first()
evidence = s.query(VaultEvidence).filter_by(id=test_row.evidence_id).first()
comparison = s.query(VaultComparison).filter_by(id=test_row.comparison_id).first()
jobs.append([evidence, comparison])
s.close()
return jobs
def insert_result(vault, name, path):
# get the current batch's id
batch_id = vault.select_last(VaultBatch).id
r = VaultResult(name, path, bath_id)
vault.insert(r)
def status(name):
status = {}
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
s = vault.get_session()
status['num_jobs'] = vault.sizeof(VaultJob)
if status['num_jobs'] == 0:
return status
status['num_pending'] = s.query(VaultJob).filter_by(event_id=1).count()
status['num_processing'] = s.query(VaultJob).filter_by(event_id=2).count()
status['num_complete'] = s.query(VaultJob).filter_by(event_id=3).count()
job_row = s.query(VaultJob).filter_by(event_id=1).first()
test_row = s.query(VaultTest).filter_by(id=job_row.test_id).first()
evidence_name = s.query(VaultEvidence).filter_by(id=test_row.evidence_id).first().name
comparison_name = s.query(VaultComparison).filter_by(id=test_row.comparison_id).first().name
status['next_job'] = [evidence_name, comparison_name]
s.close()
return status
'''
# END: Automation API
# END: External functions

68
cjs/database.py

@ -31,6 +31,10 @@ from core.database import _database_compound_add
from core.database import _database_remove
from core.database import _database_add
from core.database import _database_select
from core.database import database_count
from core.database import database_first
from core.database import database_last
import models.reqbt
from models.reqbt import EvidenceSchema
@ -279,8 +283,8 @@ def stage(ctx, database_name, batch_name):
click.echo('No changes were made.')
@database.command()
@click.argument('database')
def status(db):
@click.argument('database_name')
def status(database_name):
'''
Display information about a database.\f
@ -288,23 +292,49 @@ def status(db):
ctx: Click command-line interface context.
src: Source directory containing input file(s)
'''
from core.database import database_status
s = core.database.status(db)
if s['num_jobs'] > 0:
plabel('Job Tracker', '', 'yellow')
plabel(' Total', s['num_jobs'], 'yellow')
plabel(' Pending', (s['num_pending']/s['num_jobs'])*100, 'yellow')
plabel(' Processing', (s['num_processing']/s['num_jobs'])*100, 'yellow')
plabel(' Complete', (s['num_complete']/s['num_jobs'])*100, 'yellow')
click.echo('')
evidence, comparison = s['next_job']
plabel('Next Job', '%s and %s' % (evidence, comparison), 'yellow')
else:
click.echo('No job currently staged.')
get_count = lambda Table, expr=True: database_count(database_name, Table, expr)
evidence_count = get_count(EvidenceTable)
comparison_count = get_count(ComparisonTable)
test_count = get_count(TestTable)
job_count = get_count(JobTable)
pending_count = get_count(JobTable, JobTable.event_id == 1)
processing_count = get_count(JobTable, JobTable.event_id == 2)
complete_count = get_count(JobTable, JobTable.event_id == 3)
failed_count = get_count(JobTable, JobTable.event_id == 4)
next_job = database_first(database_name,
JobTable,
JobTable.event_id == 1)
next_test = database_first(database_name,
TestTable,
TestTable.id_ == next_job.test_id)
next_evidence = _database_select(database_name,
EvidenceTable,
EvidenceTable.id_ == next_test.evidence_id)[0]
next_comparison = _database_select(database_name,
ComparisonTable,
ComparisonTable.id_ == next_test.comparison_id)[0]
next_job = '%s and %s' % (next_evidence.name, next_comparison.name)
current_batch = database_last(database_name, BatchTable).name
echo_title('Database Status', fg='yellow', bold=True)
echo_field('Evidence', evidence_count, fg='yellow')
echo_field('Comparison', comparison_count, fg='yellow')
echo_field('Tests', test_count, fg='yellow')
click.echo('')
echo_field('Jobs', job_count, fg='yellow')
echo_field(' Pending', pending_count, fg='yellow')
echo_field(' Processing', processing_count, fg='yellow')
echo_field(' Complete', complete_count, fg='yellow')
echo_field(' Failed', failed_count, fg='yellow')
click.echo('')
echo_field('Current Batch', current_batch, fg='yellow')
echo_field('Next Job', next_job, fg='yellow')
echo_field('Average Time', '23.3 sec', fg='yellow')
# END command-line interface

364
cjs/debug.py

@ -1,364 +0,0 @@
#!/usr/bin/env python3
# standard library modules
import functools
from functools import wraps
# pip modules
import sqlalchemy
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import scoped_session
from sqlalchemy.exc import IntegrityError
# local modules
import core.state
from core.state import add_state
from core.state import remove_state
from core.state import select_state
import models.state
from models.state import DatabaseState
import models.database
#from models.database import DatabaseBase
from models.database import *
# BEGIN: Generic Database API
# modified from https://stackoverflow.com/a/36944992
# NOTE: ideally this would be inside the database class
def _pass_session(fn):
@wraps(fn) # get around the fact a decorator cannot have self
def wrapper(self, *fn_args, **fn_kwargs):
s = self._get_session()
res = fn(self, s, *fn_args, **fn_kwargs)
s.close()
return res
return wrapper
class Database:
def __init__(self, location):
self.location = location # working directory of the database
self._engine = create_engine('sqlite:///%s/database.db' % location)
self._SessionFactory = sessionmaker(bind=self._engine)
def _get_session(self):
DatabaseBase.metadata.create_all(self._engine)
return self._SessionFactory()
# TODO: replace this function with a better method of doing this
@_pass_session
def initialize(self, sess):
# NOTE: this method was created because I was unable to get sqlalchemy's
# initial values to work
names = ['Pending', 'Processing', 'Complete', 'Failed']
events = [EventTable(n) for n in names]
self.add_tables(events)
@_pass_session
def add_tables(self, sess, Tables):
try:
for T in Tables:
sess.add(T)
sess.commit()
except:
# TODO: implement fine-grained handling
raise
@_pass_session
def add_compound_tables(self, sess, Table, relations):
# NOTE: this may not be the best way to do it, but it is a much faster
# way of inserting the evidence + comparison tables into test
# it accepts this input: [PrimaryTable, [SubTable1, SubTable2]]
# and insert the result in Table(PrimaryTable.id, SubTable.id)
try:
for r in relations:
left_table, Ts = r
sess.add(left_table)
for right_table in Ts:
sess.add(right_table)
sess.flush()
sess.add(Table(left_table.id_, right_table.id_))
sess.commit()
except:
# TODO: implement fine-grained handling
raise
@_pass_session
def remove_table(sess, Table, expr=True):
try:
sess.query(Table).filter(expr).delete()
sess.commit()
except:
# TODO: implement fine-grained handling
raise
@_pass_session
def update_table(sess, Table, expr):
pass
@_pass_session
def select_tables(self, sess, Table, expr=True):
try:
q = sess.query(Table).filter(expr).all()
return q
except:
# TODO: implement fine-grained handling
raise
# END: Generic Database API
# BEGIN: Internal functions
# BEGIN: Database State functions
def _get_state(name):
try:
ds = select_state(DatabaseState,
DatabaseState.name == name)
return ds[0] # select_state returns a list
except:
# TODO: implement fine-grained handling
raise
# END: Database State functions
# BEGIN: Database object functions
def _get_database(name):
try:
ds = _get_state(name)
db = Database(ds.location)
return db
except:
# TODO: implement fine-grained handling
raise
# TODO: name is easily confused with add_database, consider refactoring
def _database_add(name, Tables):
try:
db = _get_database(name)
db.add_tables(Tables)
except:
# TODO: implement fine-grained handling
raise
# TODO: name is easily confused with add_database, consider refactoring
def _database_compound_add(name, Table, Tables):
try:
db = _get_database(name)
db.add_compound_tables(Table, Tables)
except:
# TODO: implement fine-grained handling
raise
# TODO: name is easily confused with remove_database, consider refactoring
def _database_remove(name, Table):
try:
db = _get_database(name)
# remove_table's expr parameter is true by default
# and will remove everything in the given table if
# not passed another value
db.remove_table(Table)
except:
# TODO: implement fine-grained handling
raise
# TODO: name is easily confused with add_database, consider refactoring
def _database_update(name, Table):
try:
db = _get_database(name)
pass
except:
# TODO: implement fine-grained handling
raise
# TODO: name is easily confused with add_database, consider refactoring
def _database_select(name, Table, expr):
try:
db = _get_database(name)
Ts = db.select_tables(Table, expr) # returns list of Tables
return Ts
except:
# TODO: implement fine-grained handling
raise
# END: Database object functions
# END: Internal functions
# BEGIN: External functions
# BEGIN: Database State functions
def add_database(name, location, init=True):
try:
ds = DatabaseState(name, location)
add_state(ds)
# TODO: replace with a better method of doing this
if init:
db = _get_database(name)
db.initialize()
except:
# TODO: implement fine-grained handling
raise
def remove_database(name):
try:
remove_state(DatabaseState,
DatabaseState.name == name)
except:
# TODO: implement fine-grained handling
raise
def update_database(name):
try:
pass
except:
# TODO: implement fine-grained handling
raise
def list_database():
try:
# select_state's expr parameter is true by default
# and will return everything in the state table if
# not passed another value
ds = select_state(DatabaseState)
return ds
except:
# TODO: implement fine-grained handling
raise
# END: Database State functions
# BEGIN: Database object functions
def database_location(name):
try:
ds = _get_database(name)
return ds.location
except:
raise
# END: Database object functions
# END: External functions
def tester():
name = 'debug'
#location = '.data/my_database'
#add_database(name, location)
#print(list_database())
#database_add(name, [EvidenceTable('my_file', '/path/to/file')])
print(_database_select(name, EventTable, EventTable.name == 'Pending')[0].name)
tables = [
[
EvidenceTable('debug-1', '/path/to/file1'),
[
ComparisonTable('debug-1-1', '/path/to/file2'),
ComparisonTable('debug-1-2', '/path/to/file3')
]
],
[
EvidenceTable('debug-2', '/path/to/file4'),
[
ComparisonTable('debug-2-1', '/path/to/file5'),
ComparisonTable('debug-2-2', '/path/to/file6')
]
]
]
_database_compound_add('debug', TestTable, tables)
#tester()
'''
def upload(name, evidence, comparisons):
try:
db = _get_database(name)
# format
except:
# TODO: implement fine-grained handling
raise
evidence_name = evidence['name']
evidence_file = '%s/input/%s.json' % (db.location, evidence_name)
evidence_file = str(abs_path(evidence_file))
write_json(evidence_file, evidence)
evidence_id = vault.insert(VaultEvidence(evidence_name, evidence_file))
'''
'''
def stage(name, batch):
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
batch_id = vault.insert(VaultBatch(batch))
tests = vault.get_tests()
l = [VaultJob(t.id, batch_id) for t in tests]
vault.insert(l)
for c in comparisons:
comparison_name = c['name']
comparison_file = '%s/input/%s.json' % (vault_root, comparison_name)
comparison_file = str(abs_path(comparison_file))
write_json(comparison_file, c)
comparison_id = vault.insert(VaultComparison(comparison_name, comparison_file))
vault.insert(VaultTest(evidence_id, comparison_id))
def next_job(name):
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
s = vault.get_session()
job_row = s.query(VaultJob).filter_by(event_id=1).first()
test_row = s.query(VaultTest).filter_by(id=job_row.test_id).first()
evidence = s.query(VaultEvidence).filter_by(id=test_row.evidence_id).first()
comparison = s.query(VaultComparison).filter_by(id=test_row.comparison_id).first()
s.close()
return [evidence, comparison]
def all_jobs(name):
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
s = vault.get_session()
jobs = []
for job_row in s.query(VaultJob).filter_by(event_id=1).all():
test_row = s.query(VaultTest).filter_by(id=job_row.test_id).first()
evidence = s.query(VaultEvidence).filter_by(id=test_row.evidence_id).first()
comparison = s.query(VaultComparison).filter_by(id=test_row.comparison_id).first()
jobs.append([evidence, comparison])
s.close()
return jobs
def insert_result(vault, name, path):
# get the current batch's id
batch_id = vault.select_last(VaultBatch).id
r = VaultResult(name, path, bath_id)
vault.insert(r)
def status(name):
status = {}
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
s = vault.get_session()
status['num_jobs'] = vault.sizeof(VaultJob)
if status['num_jobs'] == 0:
return status
status['num_pending'] = s.query(VaultJob).filter_by(event_id=1).count()
status['num_processing'] = s.query(VaultJob).filter_by(event_id=2).count()
status['num_complete'] = s.query(VaultJob).filter_by(event_id=3).count()
job_row = s.query(VaultJob).filter_by(event_id=1).first()
test_row = s.query(VaultTest).filter_by(id=job_row.test_id).first()
evidence_name = s.query(VaultEvidence).filter_by(id=test_row.evidence_id).first().name
comparison_name = s.query(VaultComparison).filter_by(id=test_row.comparison_id).first().name
status['next_job'] = [evidence_name, comparison_name]
s.close()
return status
'''

2
cjs/utils/cli.py

@ -9,7 +9,7 @@ def echo_title(title, *args, **kwargs):
def echo_field(label, value, *args, **kwargs):
# TODO: add dynamic indent
echo(style(f' {label}: ', *args, **kwargs) + value)
echo(f"{style(f' {label}: ', *args, **kwargs)}{value}")
def echo_enum(i, value, *args, **kwargs):
echo(f' {i}. ' + style(f'{value}', *args, **kwargs))
Loading…
Cancel
Save