def datastore(datastore_connection, filestore):
global test_file, file_res_list
ds = datastore_connection
try:
create_users(ds)
# noinspection PyUnusedLocal
for _f in range(NUM_FILES):
f = random_model_obj(File)
if test_file is None:
test_file = f
ds.file.save(f.sha256, f)
filestore.put(f.sha256, f.sha256)
# noinspection PyUnusedLocal
for _r in range(random.randint(1, 3)):
r = random_model_obj(Result)
r.sha256 = f.sha256
file_res_list.append(r.build_key())
ds.result.save(r.build_key(), r)
ds.file.commit()
ds.result.commit()
yield ds
finally:
wipe_users(ds)
ds.file.wipe()
ds.result.wipe()
for key in file_res_list:
filestore.delete(key[:64])
def datastore(datastore_connection):
global user_list
ds = datastore_connection
try:
create_users(ds)
data = {
'alert': [],
'error': [],
'search': [],
'signature': [],
'submission': [],
}
for x in range(NUM_FAVS):
f = random_model_obj(Favorite)
f.name = f"test_{x+1}"
for key in data:
data[key].append(f)
ds.user_favorites.save('admin', data)
ds.user_favorites.save('user', data)
for x in range(NUM_USERS):
u = random_model_obj(User)
u.uname = f"test_{x+1}"
ds.user.save(u.uname, u)
ds.user_favorites.save(u.uname, data)
ds.user_avatar.save(u.uname, AVATAR)
user_list.append(u.uname)
yield ds
finally:
wipe_users(ds)
def datastore(datastore_connection):
ds = datastore_connection
try:
create_users(ds)
for _ in range(NUM_ITEMS):
f = random_model_obj(File)
f_hash_list.append(f.sha256)
ds.file.save(f.sha256, f)
for x in range(NUM_ITEMS):
a = random_model_obj(Alert)
a.file.sha256 = f_hash_list[x]
ds.alert.save(a.alert_id, a)
for x in range(NUM_ITEMS):
r = random_model_obj(Result)
r.sha256 = f_hash_list[x]
ds.result.save(r.build_key(), r)
ds.alert.commit()
ds.file.commit()
ds.submission.commit()
yield ds
finally:
ds.alert.wipe()
ds.file.wipe()
ds.submission.wipe()
wipe_users(ds)
def datastore(datastore_connection, filestore):
ds = datastore_connection
try:
random_data.create_heuristics(ds, heuristics_count=10)
random_data.create_services(ds)
random_data.create_signatures(ds)
random_data.create_users(ds)
random_data.create_workflows(ds)
submissions = []
for _ in range(2):
submissions.append(random_data.create_submission(
ds, filestore))
random_data.create_alerts(ds,
alert_count=10,
submission_list=submissions)
yield ds
finally:
# Cleanup test data
random_data.wipe_alerts(ds)
random_data.wipe_heuristics(ds)
random_data.wipe_services(ds)
random_data.wipe_signatures(ds)
random_data.wipe_submissions(ds, filestore)
random_data.wipe_users(ds)
random_data.wipe_workflows(ds)
def datastore(datastore_connection, filestore):
ds = datastore_connection
try:
create_users(ds)
create_services(ds)
for _ in range(NUM_FILES):
f = random_model_obj(File)
ds.file.save(f.sha256, f)
file_hashes.append(f.sha256)
filestore.put(f.sha256, f.sha256)
ds.file.commit()
yield ds
finally:
# Cleanup Elastic
ds.file.wipe()
wipe_services(ds)
wipe_users(ds)
# Cleanup Minio
for f in file_hashes:
filestore.delete(f)
# Cleanup Redis
nq.delete()
iq.delete()
def datastore(datastore_connection):
ds = datastore_connection
try:
create_users(ds)
for x in range(TEST_RESULTS):
f = random_model_obj(File)
ds.file.save(f.sha256, f)
file_list.append(f.sha256)
ds.file.commit()
for x in range(TEST_RESULTS):
e = random_model_obj(Error)
e.sha256 = file_list[x]
ds.error.save(e.build_key(), e)
error_key_list.append(e.build_key())
ds.error.commit()
for x in range(TEST_RESULTS):
r = random_model_obj(Result)
r.sha256 = file_list[x]
ds.result.save(r.build_key(), r)
result_key_list.append(r.build_key())
ds.result.commit()
yield ds
finally:
ds.error.wipe()
ds.file.wipe()
ds.result.wipe()
wipe_users(ds)
def datastore(datastore_connection):
try:
create_users(datastore_connection)
create_heuristics(datastore_connection)
yield datastore_connection
finally:
wipe_users(datastore_connection)
wipe_heuristics(datastore_connection)
def datastore(datastore_connection, filestore):
try:
create_users(datastore_connection)
create_submission(datastore_connection, filestore)
yield datastore_connection
finally:
wipe_users(datastore_connection)
wipe_submissions(datastore_connection, filestore)
def datastore(datastore_connection):
try:
create_users(datastore_connection)
create_services(datastore_connection)
yield datastore_connection
finally:
wipe_users(datastore_connection)
wipe_services(datastore_connection)
def purge_data(ds, fs):
wipe_alerts(ds)
wipe_heuristics(ds)
wipe_services(ds)
wipe_signatures(ds)
wipe_submissions(ds, fs)
wipe_users(ds)
wipe_workflows(ds)
def datastore(datastore_connection, filestore):
global submission
try:
create_users(datastore_connection)
submission = create_submission(datastore_connection, filestore)
yield datastore_connection
finally:
wipe_users(datastore_connection)
wipe_submissions(datastore_connection, filestore)
sq.delete()
def datastore(datastore_connection, filestore):
try:
create_users(datastore_connection)
for _ in range(NUM_SUBMISSIONS):
create_submission(datastore_connection, filestore)
yield datastore_connection
finally:
wipe_users(datastore_connection)
wipe_submissions(datastore_connection, filestore)
def datastore(datastore_connection):
ds = datastore_connection
try:
create_users(ds)
signatures.extend(create_signatures(ds))
ds.signature.commit()
for _ in range(TEST_SIZE):
f = random_model_obj(File)
ds.file.save(f.sha256, f)
file_list.append(f.sha256)
ds.file.commit()
for x in range(TEST_SIZE):
a = random_model_obj(Alert)
a.file.sha256 = file_list[x]
ds.alert.save(a.alert_id, a)
ds.alert.commit()
for x in range(TEST_SIZE):
r = random_model_obj(Result)
r.sha256 = file_list[x]
ds.result.save(r.build_key(), r)
ds.result.commit()
for x in range(TEST_SIZE):
s = random_model_obj(Submission)
for f in s.files:
f.sha256 = file_list[x]
ds.submission.save(s.sid, s)
ds.submission.commit()
for x in range(TEST_SIZE):
h = random_model_obj(Heuristic)
h.heur_id = f"AL_TEST_{x}"
ds.heuristic.save(h.heur_id, h)
ds.heuristic.commit()
for _ in range(TEST_SIZE):
w_id = get_random_id()
w = random_model_obj(Workflow)
ds.workflow.save(w_id, w)
ds.workflow.commit()
yield ds
finally:
ds.alert.wipe()
ds.file.wipe()
ds.result.wipe()
ds.signature.wipe()
ds.submission.wipe()
ds.heuristic.wipe()
ds.workflow.wipe()
wipe_users(ds)
def datastore(datastore_connection):
global test_submission
try:
create_users(datastore_connection)
test_submission = random_model_obj(Submission)
datastore_connection.submission.save(test_submission.sid,
test_submission)
datastore_connection.submission.commit()
yield datastore_connection
finally:
datastore_connection.submission.wipe()
wipe_users(datastore_connection)
def datastore(datastore_connection, filestore):
try:
create_users(datastore_connection)
submission = create_submission(datastore_connection, filestore)
alert = random_model_obj(Alert)
alert.alert_id = ALERT_ID
alert.sid = submission.sid
datastore_connection.alert.save(ALERT_ID, alert)
yield datastore_connection
finally:
wipe_users(datastore_connection)
wipe_submissions(datastore_connection, filestore)
datastore_connection.alert.delete(ALERT_ID)
def datastore(datastore_connection):
try:
global test_error
create_users(datastore_connection)
for _ in range(NUM_ERRORS):
e = random_model_obj(Error)
if test_error is None:
test_error = e
datastore_connection.error.save(e.build_key(), e)
datastore_connection.error.commit()
yield datastore_connection
finally:
wipe_users(datastore_connection)
datastore_connection.error.wipe()
def datastore(datastore_connection):
try:
create_users(datastore_connection)
create_services(datastore_connection)
for _ in range(NUM_WORKFLOWS):
workflow = random_model_obj(Workflow)
workflow_list.append(workflow.workflow_id)
datastore_connection.workflow.save(workflow.workflow_id, workflow)
datastore_connection.workflow.commit()
yield datastore_connection
finally:
wipe_users(datastore_connection)
wipe_services(datastore_connection)
datastore_connection.workflow.wipe()
def cleanup(ds, fs):
wipe_users(ds)
wipe_submissions(ds, fs)
wipe_alerts(ds)
def purge_alert(ds, fs):
wipe_users(ds)
wipe_submissions(ds, fs)
ds.alert.wipe()
