def test_import_bundle(datastore, login_session, filestore):
_, session, host = login_session
ds = datastore
# Create a temporary bundle
submission = random.choice(
ds.submission.search('id:*', rows=100, as_obj=False)['items'])
bundle_file = create_bundle(submission['sid'], working_dir='/tmp/bundle')
# Delete associated submission
ds.delete_submission_tree(submission['sid'], transport=filestore)
ds.error.commit()
ds.file.commit()
ds.result.commit()
ds.submission.commit()
with open(bundle_file, 'rb') as bfh:
resp = get_api_data(session,
f"{host}/api/v4/bundle/",
method="POST",
data=bfh.read())
assert resp['success']
ds.submission.commit()
assert submission == random.choice(
ds.submission.search('id:*', rows=100, as_obj=False)['items'])
def test_alert_bundle(datastore_connection, filestore, config):
# Cleanup previous runs
datastore_connection.alert.delete(ALERT_ID)
# Create a temporary submission
submission = create_submission(datastore_connection, filestore)
sid = submission['sid']
# Create a random alert
alert = random_model_obj(Alert)
alert.alert_id = ALERT_ID
alert.sid = sid
datastore_connection.alert.save(ALERT_ID, alert)
# Create the submission's bundle
path = create_bundle(ALERT_ID, use_alert=True)
# Test if the bundle
assert os.path.exists(path)
with open(path, 'rb') as fh:
assert is_cart(fh.read(256))
# Remove alert and submission from DB
datastore_connection.alert.delete(alert.alert_id)
datastore_connection.delete_submission_tree(sid, transport=filestore)
assert datastore_connection.alert.get_if_exists(alert.alert_id) is None
assert datastore_connection.submission.get_if_exists(sid) is None
# Restore bundle
new_submission = import_bundle(path)
# Validate restored submission
assert new_submission['sid'] == sid
assert new_submission['metadata']['bundle.source'] == config.ui.fqdn
# Validate restored alert
new_alert = datastore_connection.alert.get_if_exists(alert.alert_id, as_obj=False)
assert new_alert['alert_id'] == ALERT_ID
assert new_alert['sid'] == sid
assert new_alert['metadata']['bundle.source'] == config.ui.fqdn
# Cleanup
assert not os.path.exists(path)
datastore_connection.alert.delete(alert.alert_id)
datastore_connection.delete_submission_tree(sid, transport=filestore)
assert datastore_connection.alert.get_if_exists(alert.alert_id) is None
assert datastore_connection.submission.get_if_exists(sid) is None
def test_alert_no_submission_bundle(datastore_connection, config):
# Cleanup previous runs
datastore_connection.alert.delete(ALERT_ID)
datastore_connection.submission.delete(SUBMISSION_ID)
# Create a random alert
alert = random_model_obj(Alert)
alert.alert_id = ALERT_ID
alert.sid = SUBMISSION_ID
datastore_connection.alert.save(ALERT_ID, alert)
# Create the submission's bundle
path = create_bundle(ALERT_ID, use_alert=True)
# Test if the bundle
assert os.path.exists(path)
with open(path, 'rb') as fh:
assert is_cart(fh.read(256))
# Remove alert from Datastore
datastore_connection.alert.delete(alert.alert_id)
assert datastore_connection.alert.get_if_exists(alert.alert_id) is None
assert datastore_connection.submission.get_if_exists(alert.sid) is None
# Restore bundle
new_submission = import_bundle(path)
# Validate restored submission
assert new_submission is None
# Validate restored alert
new_alert = datastore_connection.alert.get_if_exists(alert.alert_id)
assert new_alert['alert_id'] == ALERT_ID
assert new_alert['sid'] == SUBMISSION_ID
assert new_alert['metadata']['bundle.source'] == config.ui.fqdn
# Cleanup
datastore_connection.alert.delete(ALERT_ID)
datastore_connection.submission.delete(SUBMISSION_ID)
def test_submission_bundle(datastore_connection, filestore, config):
# Create a temporary submission
submission = create_submission(datastore_connection, filestore)
sid = submission['sid']
# Create the submission's bundle
path = create_bundle(sid)
# Test if the bundle
assert os.path.exists(path)
with open(path, 'rb') as fh:
assert is_cart(fh.read(256))
# Remove submission from DB
datastore_connection.delete_submission_tree(sid, transport=filestore)
assert datastore_connection.submission.get_if_exists(sid) is None
# Restore bundle
new_submission = import_bundle(path, cleanup=False)
# Validate restored submission
assert new_submission['sid'] == sid
assert new_submission['metadata']['bundle.source'] == config.ui.fqdn
# Test inserting failure
with pytest.raises(SubmissionAlreadyExist):
import_bundle(path, cleanup=False)
# Test skip failure on exist
new_submission = import_bundle(path, exist_ok=True)
# Validate restored submission
assert new_submission['sid'] == sid
assert new_submission['metadata']['bundle.source'] == config.ui.fqdn
# Cleanup
assert not os.path.exists(path)
datastore_connection.delete_submission_tree(sid, transport=filestore)
assert datastore_connection.submission.get_if_exists(sid) is None
def test_alert_import_bundle(datastore, login_session, filestore):
_, session, host = login_session
ds = datastore
# Create a temporary bundle
alert = ds.alert.get_if_exists(ALERT_ID, as_obj=False)
submission = ds.submission.get_if_exists(alert['sid'], as_obj=False)
bundle_file = create_bundle(ALERT_ID,
working_dir='/tmp/bundle',
use_alert=True)
# Delete associated alert and submission
ds.alert.delete(ALERT_ID)
ds.delete_submission_tree(alert['sid'], transport=filestore)
ds.alert.commit()
ds.error.commit()
ds.file.commit()
ds.result.commit()
ds.submission.commit()
with open(bundle_file, 'rb') as bfh:
resp = get_api_data(session,
f"{host}/api/v4/bundle/",
method="POST",
data=bfh.read())
assert resp['success']
ds.submission.commit()
new_submission = ds.submission.get_if_exists(alert['sid'],
as_obj=False)
assert new_submission['sid'] == alert['sid']
assert 'bundle.source' in new_submission['metadata']
new_alert = ds.alert.get_if_exists(ALERT_ID, as_obj=False)
assert new_alert['alert_id'] == ALERT_ID
assert new_alert['sid'] == submission['sid']
def create_submission_bundle(self, sid, bundle_path):
temp_bundle_file = create_bundle(sid, working_dir=os.path.dirname(bundle_path))
os.rename(temp_bundle_file, bundle_path)
def create_alert_bundle(self, alert_id, bundle_path):
temp_bundle_file = create_bundle(alert_id, working_dir=os.path.dirname(bundle_path), use_alert=True)
os.rename(temp_bundle_file, bundle_path)
def test_failed_alert_bundle():
# Test creation failure
with pytest.raises(AlertNotFound):
create_bundle("ThisAlertIDDoesNotExists", use_alert=True)
def test_failed_submission_bundle():
# Test creation failure
with pytest.raises(SubmissionNotFound):
create_bundle("ThisSIDDoesNotExists")