def test_expire_missing_service():
redis = mock.MagicMock(spec=Redis)
redis.keys.return_value = [b'service-queue-not-service-a']
redis_persist = mock.MagicMock(spec=Redis)
datastore = mock.MagicMock()
service_a = random_model_obj(Service)
service_a.name = 'a'
service_a.enabled = True
datastore.list_all_services.return_value = [service_a]
with mock.patch('assemblyline_core.plumber.run_plumber.time'):
plumber = Plumber(redis=redis,
redis_persist=redis_persist,
datastore=datastore)
plumber.get_service_stage = mock.MagicMock(
return_value=ServiceStage.Running)
plumber.dispatch_client = mock.MagicMock()
task = random_model_obj(Task)
plumber.dispatch_client.request_work.side_effect = [task, None, None]
plumber.running = TrueCountTimes(count=1)
plumber.try_run()
assert plumber.dispatch_client.service_failed.call_count == 1
args = plumber.dispatch_client.service_failed.call_args
assert args[0][0] == task.sid
def test_add(datastore, client):
# Select a random service
service = random_id_from_collection(datastore, 'service_delta')
# Make sure the service can generate signatures
service_data = datastore.get_service_with_delta(service, as_obj=False)
service_data['update_config'] = {
'generates_signatures': True,
'sources': []
}
res = client.service.set(service, service_data)
assert res['success']
# Add a new source
added_source_name = "NEW_TEST_SOURCE"
new_source = random_model_obj(UpdateSource, as_json=True)
new_source['name'] = added_source_name
res = client.signature.sources.add(service, new_source)
assert res['success']
# Test if new source is there
service_data = datastore.get_service_with_delta(service, as_obj=False)
assert service_data['update_config']['sources'][0][
'name'] == added_source_name
def test_add(datastore, client):
new_workflow = random_model_obj(Workflow, as_json=True)
res = client.workflow.add(new_workflow)
assert res['success']
saved_workflow = datastore.workflow.get(res['workflow_id'], as_obj=False)
assert saved_workflow is not None
assert saved_workflow['query'] == new_workflow['query']
def test_add(datastore, client):
user_id = 'to_be_added'
user_data = random_model_obj(User).as_primitives()
user_data['email'] = user_data['email']
user_data['uname'] = user_id
assert datastore.user.get(user_id, as_obj=False) is None
res = client.user.add(user_id, user_data)
assert res['success']
assert datastore.user.get(user_id, as_obj=False) == user_data
def test_delete(datastore, client):
user_id = 'to_be_deleted'
user_data = random_model_obj(User)
user_data['uname'] = user_id
datastore.user.save(user_id, user_data)
assert datastore.user.get(user_id, as_obj=False) is not None
res = client.user.delete(user_id)
assert res['success']
datastore.user.commit()
assert datastore.user.get(user_id, as_obj=False) is None
def test_flush_paused_queues():
redis = mock.MagicMock(spec=Redis)
redis.keys.return_value = [b'service-queue-a']
redis.zcard.return_value = 0
redis_persist = mock.MagicMock(spec=Redis)
datastore = mock.MagicMock()
service_a = random_model_obj(Service)
service_a.name = 'a'
service_a.enabled = True
datastore.list_all_services.return_value = [service_a]
plumber = Plumber(redis=redis,
redis_persist=redis_persist,
datastore=datastore,
delay=1)
plumber.get_service_stage = mock.MagicMock(
return_value=ServiceStage.Running)
plumber.dispatch_client = mock.MagicMock()
task = random_model_obj(Task)
plumber.dispatch_client.request_work.side_effect = [task, None, None]
plumber.running = TrueCountTimes(count=1)
plumber.try_run()
assert plumber.dispatch_client.service_failed.call_count == 0
plumber.get_service_stage = mock.MagicMock(
return_value=ServiceStage.Paused)
plumber.running = TrueCountTimes(count=1)
plumber.try_run()
assert plumber.dispatch_client.service_failed.call_count == 1
args = plumber.dispatch_client.service_failed.call_args
assert args[0][0] == task.sid
def test_favorites(datastore, client):
possible_fav_types = {
'alert', 'error', 'search', 'signature', 'submission'
}
user_id = random_id_from_collection(datastore, 'user')
# Test getting favorites
favorites = client.user.favorites(user_id)
assert possible_fav_types.issubset(set(favorites.keys()))
# Test adding favorites
fav_type = random.choice(list(possible_fav_types))
new_fav = random_model_obj(Favorite, as_json=True)
res = client.user.favorites.add(user_id, fav_type, new_fav)
assert res['success']
# Test if favorite was added
found = False
for fav in datastore.user_favorites.get(user_id, as_obj=False)[fav_type]:
if fav['name'] == new_fav['name'] and fav['query'] == new_fav['query']:
found = True
break
assert found
# Test delete favorite
client.user.favorites.delete(user_id, fav_type, new_fav['name'])
assert res['success']
# Test if favorite was deleted
found = False
for fav in datastore.user_favorites.get(user_id, as_obj=False)[fav_type]:
if fav['name'] == new_fav['name'] and fav['query'] == new_fav['query']:
found = True
break
assert not found
# Test update favorites
new_favs = {
'alert': [new_fav],
'error': [new_fav],
'search': [new_fav],
'signature': [new_fav],
'submission': [new_fav]
}
res = client.user.favorites.update(user_id, new_favs)
assert res['success']
# Test is favorites were really updated
for favs in datastore.user_favorites.get(user_id, as_obj=False).values():
assert favs[0]['name'] == new_fav['name']
def test_add_update_many(datastore, client):
# Insert a dummy signature
source = "source"
s_type = "type"
sig_list = []
for x in range(10):
data = random_model_obj(Signature).as_primitives(hidden_fields=True)
data['signature_id'] = f"test_sig_{x}"
data['name'] = f"sig_name_{x}"
data['status'] = "DEPLOYED"
data['source'] = source
data['type'] = s_type
sig_list.append(data)
res = client.signature.add_update_many(source, s_type, sig_list)
assert res == {'errors': False, 'success': 10, 'skipped': []}
# Test the signature data
datastore.signature.commit()
data = Signature(random.choice(sig_list)).as_primitives()
key = "%s_%s_%s" % (data['type'], data['source'], data['signature_id'])
added_sig = datastore.signature.get(key, as_obj=False)
assert data == added_sig
# This signature should fail
random_sig_fail = deepcopy(data)
random_sig_fail['signature_id'] = "FAIL"
res = client.signature.add_update_many(source, s_type, [random_sig_fail])
assert res['success'] == 0
random_key = "%s_%s_%s" % (random_sig_fail['type'],
random_sig_fail['source'],
random_sig_fail['signature_id'])
assert random_key in res['skipped']
# Does not fail if we don't dedup names
res = client.signature.add_update_many(source,
s_type, [random_sig_fail],
dedup_name=False)
assert res['success'] == 1
def test_add_update(datastore, client):
random_sig = random_model_obj(Signature, as_json=True)
random_sig['signature_id'] = "random"
random_sig['name'] = "sig_name_random"
random_sig['status'] = "DEPLOYED"
random_sig['source'] = "source"
random_sig['type'] = "type"
res = client.signature.add_update(random_sig)
assert res['success']
datastore.signature.commit()
# This signature should fail
random_sig_fail = deepcopy(random_sig)
random_sig_fail['signature_id'] = "FAIL"
with pytest.raises(ClientError):
client.signature.add_update(random_sig_fail)
datastore.signature.commit()
res = client.signature.add_update(random_sig_fail, dedup_name=False)
assert res['success']