You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

273 lines
7.4 KiB

#!/usr/bin/env python3
# external
import sqlalchemy
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import scoped_session
from sqlalchemy.exc import IntegrityError
import core.state
import utils.fs
from utils.fs import remove_file, abs_path, copy_file, create_dir
import utils.io
from utils.io import write_json, get_valid_filename
import models.state
from models.state import StateVault
import models.database
from models.database import *
# BEGIN vault backend
class Database:
def __init__(self, path):
self._engine = create_engine('sqlite:///%s' % path)
self._SessionFactory = sessionmaker(bind=self._engine)
def _get_session(self):
DatabaseBase.metadata.create_all(self._engine)
return self._SessionFactory
def get_session(func):
def wrapper(self, *args, **kwargs):
s = self._get_session()
res = func(s, *args, **kwargs)
s.close()
return res
return wrapper
def _get(self, table):
s = self.get_session()
res = s.query(table).all()
s.close()
return res
def _clear(self, table):
s = self.get_session()
s.query(table).delete()
s.commit()
s.close()
def get_tests(self):
return self._get(VaultTest)
def clear_jobs(self):
self._clear(VaultJob)
def insert_evidence(self, name, path):
evidence = VaultEvidence(name, path)
self._insert(evidence)
return evidence.id
def insert_comparison(self, name, path):
comparison = VaultComparison(name, path)
self._insert(comparison)
return comparison.id
def insert_test(self, evidence_id, comparison_id):
test = VaultTest(evidence_id, comparison_id)
self._insert(test)
return test.id
def insert_event(self, name):
event = VaultEvent(name)
self._insert(event)
return event.id
def insert_job(self, test_id, batch_id):
job = VaultJob(test_id, batch_id)
self._insert(job)
return job.id
def insert_batch(self, name):
batch = VaultBatch(name)
self._insert(batch)
return batch.id
def insert_result(self, name, path):
result = VaultResult(name, path)
self._insert(result)
return result
# new api
def _insert(self, Table):
s = self.get_session()
try:
res = []
for t in Table:
s.add(t)
s.flush()
s.refresh(t)
res.append(t.id)
except:
s.add(Table)
s.flush()
s.refresh(Table)
res = Table.id
s.commit()
s.close()
return res
def insert(self, Table):
return self._insert(Table)
def _query(self, Table):
s = self.get_session()
q = s.query(Table)
s.close()
return q
def sizeof(self, Table, where=None):
if where:
return self._query(Table).filter(where).count()
else:
return self._query(Table).count()
def next(self, Table):
s = self.get_session()
q = s._query(Table).filter(event_id=1)
def select_all(self, Table):
return self._query(Table).all()
def select_first(self, Table):
return self._query(Table).first()
def select_last(self, Table):
return self._query(Table).order_by(Table.id.desc()).first()
# public functions
def new(database_name, database_root):
try:
path = f'{database_root}/database.db'
s = StateVault(database_name, database_root, path)
core.state.add_state(s)
except:
raise
def list_databases():
try:
databases = core.state.list_state(StateVault)
return databases
except:
raise
def delete(name):
vault = _state_get_by_name(name)
try:
remove_file(vault.root)
except :
pass
_state_remove_by_name(name)
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 upload(name, evidence, comparisons):
vault_state = _state_get_by_name(name)
vault_root = vault_state.root
vault = Vault(vault_state.path)
evidence_name = evidence['name']
evidence_file = '%s/input/%s.json' % (vault_root, evidence_name)
evidence_file = str(abs_path(evidence_file))
write_json(evidence_file, evidence)
evidence_id = vault.insert(VaultEvidence(evidence_name, evidence_file))
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 clear(name):
vault_state = _state_get_by_name(name)
vault = Vault(vault_state.path)
vault.clear_jobs()
def vault_path(name):
v = _state_get_by_name(name)
return v.path
def vault_name(path):
v = _state_get_by_path(path)
return v.name
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 copy(name, destination):
vault = _state_get_by_name(name)
copy_file(vault.path, destination)
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