def test_get(self):
"""
Attempt to fetch a GroupAction document
"""
group_action1 = Database.create_group_action()
group_action2 = Database.get_group_action(group_action1.group_action_id)
self.assertEqual(group_action1.group_action_id, group_action2.group_action_id)
def test_find_no_name_fail(self):
"""
This target tests to ensure mongoengine throws an error if the target
is not found by name.
"""
with self.assertRaises(DoesNotExist):
Database.get_target('HI. I DONT EXIST.')
def test_target_rename_association(self):
"""
Tests the RenameTarget API function, check to make sure Sessions, Targets, and Groups.
"""
target = Database.create_target()
target_name = target.name
session_id = Database.create_session(target_name).session_id
orig_group = Database.create_group('some_group')
orig_group.whitelist_member(target.name)
action_id = Database.create_action(target_name).action_id
data = APIClient.rename_target(self.client, target_name, 'TEST')
self.assertEqual(data['error'], False)
target = Database.get_target('TEST')
self.assertIsNotNone(target)
with self.assertRaises(DoesNotExist):
Database.get_target(target_name)
self.assertEqual(target.name, 'TEST')
session = Database.get_session(session_id)
self.assertEqual(session.target_name, 'TEST')
action = Database.get_action(action_id)
self.assertEqual(action.target_name, 'TEST')
group = Database.get_group(orig_group.name)
self.assertIn(target.name, group.member_names)
def test_target_set_facts(self):
"""
This test will pass if the facts are correctly set.
"""
initial_facts = {
'some fact': 54,
'some other fact': 'Pi',
'A list fact': ['sdasd', 'asdasd']
}
fact_update = {
'new fact': 'Wow. I am new!',
'A list fact': ['asdasd', 'sdasd'],
'some fact': 55
}
final_facts = {
'new fact': 'Wow. I am new!',
'some other fact': 'Pi',
'A list fact': ['asdasd', 'sdasd'],
'some fact': 55
}
target = Database.create_target('FACT TEST', ['AA:BB:CC:DD:EE:FF'],
initial_facts)
data = APIClient.set_target_facts(self.client, 'FACT TEST',
fact_update)
self.assertEqual(data['error'], False)
target = Database.get_target('FACT TEST')
self.assertIsNotNone(target)
self.assertDictEqual(final_facts, target.facts)
def test_target_status_pass(self):
"""
This tests to ensure a target's status is being set properly
based on it's session statuses.
"""
session1 = Database.create_session(None, 0)
target = Database.get_target(session1.target_name)
self.assertEqual(target.status, SESSION_STATUSES.get('inactive'))
missing_timer = 2401 + SESSION_CHECK_THRESHOLD
session2 = Database.create_session(target.name,
time.time() - missing_timer, 1800,
600)
self.assertEqual(target.status, SESSION_STATUSES.get('missing'))
session3 = Database.create_session(target.name)
self.assertEqual(target.status, SESSION_STATUSES.get('active'))
session3.timestamp = 0
session3.save()
self.assertEqual(target.status, SESSION_STATUSES.get('missing'))
session2.timestamp = 0
session2.save()
self.assertEqual(target.status, SESSION_STATUSES.get('inactive'))
def test_find_pass(self):
"""
This test will attempt to create a target model object,
save it to the database, and then find it.
"""
target1 = Database.create_target(self.TEST_NAME)
target2 = Database.get_target(self.TEST_NAME)
self.assertEqual(target1, target2)
def test_whitelist_member(self):
"""
Test the whitelist_member function.
"""
group = Database.create_group('SOME GROUP')
target = Database.create_target()
group.whitelist_member(target.name)
group = Database.get_group('SOME GROUP')
self.assertListEqual([target.name], group.whitelist_members)
def test_delete(self):
"""
Test the DeleteGroup function.
"""
group = Database.create_group()
data = APIClient.delete_group(self.client, group.name)
self.assertEqual(data['error'], False)
with self.assertRaises(DoesNotExist):
Database.get_group(group.name)
def test_create_dup_name_fail(self):
"""
This test will attempt to create two targets with identitical names,
and it will fail as Mongo should throw a not unique exception.
"""
with self.assertRaises(NotUniqueError):
target1 = Database.create_target(self.TEST_NAME)
target2 = Database.create_target(self.TEST_NAME)
self.assertEqual(target1.name, target2.name)
def test_find_pass(self):
"""
This test will attempt to create an action model object,
save it to the database, and then find it.
"""
action1 = Database.create_action()
action2 = Database.get_action(action1.action_id)
self.assertIsNotNone(action1)
self.assertIsNotNone(action2)
self.assertEqual(action1, action2)
self.assertEqual(action1.action_id, action2.action_id)
def test_create_pass(self):
"""
This test will attempt to create an action model object.
"""
target = Database.create_target()
action = Database.create_action(target.name, self.TEST_ACTION_STRING)
self.assertEqual(action.target_name, target.name)
self.assertEqual(action.action_string, self.TEST_ACTION_STRING)
self.assertEqual(action.action_type, ACTION_TYPES.get('exec', 1))
self.assertEqual(action.bound_session_id, '')
self.assertIsNone(action.session_id)
def test_get_by_name(self):
"""
Test the get_by_name function.
"""
Database.create_group('some group')
Database.create_group('other group')
group = Group.get_by_name('some group')
self.assertEqual(group.name, 'some group')
with self.assertRaises(DoesNotExist):
Group.get_by_name('not group')
def test_remove_member(self):
"""
Test the remove_member function.
"""
target1 = Database.create_target()
target2 = Database.create_target()
group = Database.create_group('test_group')
group.whitelist_member(target1.name)
group.whitelist_member(target2.name)
group.remove_member(target2.name)
self.assertNotIn(target2.name, group.members)
def test_exec_parsing(self):
"""
Perform more extensive tests on exec.
"""
action_tests = [
# pipe
(
Database.parse_action_string('exec echo hi | tee output.txt'),
{
'action_type': ACTION_TYPES.get('exec', 1),
'command': 'echo',
'args': ['hi', '|', 'tee', 'output.txt']
}
),
# conflicting args
(
Database.parse_action_string('exec date --time time'),
{
'action_type': ACTION_TYPES.get('exec', 1),
'command': 'date',
'args': ['--time', 'time']
}
),
# subshell
(
Database.parse_action_string('exec find $(which bash)'),
{
'action_type': ACTION_TYPES.get('exec', 1),
'command': 'find',
'args': ['$(which', 'bash)']
}
),
# backtick subshell
(
Database.parse_action_string('exec rm -rf `which bash`'),
{
'action_type': ACTION_TYPES.get('exec', 1),
'command': 'rm',
'args': ['-rf', '`which', 'bash`']
}
),
# special chars (must be quoted)
(
Database.parse_action_string('exec echo -e "Hello \n World"'),
{
'action_type': ACTION_TYPES.get('exec', 1),
'command': 'echo',
'args': ['-e', 'Hello \n World']
}
),
]
for test in action_tests:
self.assertDictEqual(test[0], test[1])
def test_add_member(self):
"""
Test the GetGroup API function.
"""
Database.create_group('TEST GROUP')
target = Database.create_target()
data = APIClient.add_group_member(self.client, 'TEST GROUP',
target.name)
self.assertEqual(data['error'], False)
group = Database.get_group('TEST GROUP')
self.assertListEqual(group.whitelist_members, [target.name])
def test_get_unassigned_actions(self):
"""
This test ensures that the proper unassigned actions are returned for a target.
"""
target = Database.create_target()
action1 = Database.create_action(target.name, 'exec echo hello')
action2 = Database.create_action(target.name, 'exec echo world')
actions = Action.get_target_unassigned_actions(target.name)
for action in actions:
self.assertIn(action.action_id, [action1.action_id, action2.action_id])
self.assertEqual(len(actions), 2)
def test_get(self):
"""
This test will pass if it finds the correct action.
"""
target = Database.create_target()
action_id = Database.create_action(target.name, 'exec ls')['action_id']
data = APIClient.get_action(self.client, action_id)
action = Database.get_action(action_id)
self.assertEqual(data['error'], False)
self.assertIsNotNone(data['action'])
self.assertEqual(action_id, action.action_id)
self.assertEqual(target.name, action.target_name)
self.assertEqual('ls', action.command)
def test_find_pass(self):
"""
This test creates a session and attempts to find it,
it's target, and it's history document.
"""
session1 = Database.create_session()
session2 = Database.get_session(session1.session_id)
self.assertEqual(session1, session2)
self.assertIsNotNone(Database.get_target(session2.target_name))
self.assertIsNotNone(session1.history)
self.assertIsNotNone(session2.history)
def test_cancel(self):
"""
This test will pass if an action is successfully cancelled.
"""
target = Database.create_target()
action_id = Database.create_action(
target.name, 'exec echo hello world')['action_id']
action = Database.get_action(action_id)
self.assertEqual(action.cancelled, False)
data = APIClient.cancel_action(self.client, action_id)
self.assertEqual(data['error'], False)
action = Database.get_action(action_id)
self.assertEqual(action.status, ACTION_STATUSES.get('cancelled'))
self.assertEqual(action.cancelled, True)
def test_status(self):
"""
Ensure that the proper status is set for a group action
based on it's included actions.
"""
# After creation, status should be queued
group_action = Database.create_group_action()
self.assertEqual(group_action.get_status(),
GROUP_ACTION_STATUSES.get('queued', 'queued'))
# Send to a session, status should be in progress
session = Database.create_session()
group_action.actions[0].assign_to(session.session_id)
self.assertEqual(
group_action.get_status(),
GROUP_ACTION_STATUSES.get('in progress', 'in progress'))
# Submit a response to all actions, status should be success
for action in group_action.actions:
action.assign_to(session.session_id)
response = Database.create_response()
action.submit_response(response)
self.assertEqual(group_action.get_status(),
GROUP_ACTION_STATUSES.get('success', 'success'))
# Create a new group action, make all actions stale. Status should be stale
group_action = Database.create_group_action()
for action in group_action.actions:
action.queue_time = 0
action.save()
self.assertEqual(group_action.get_status(),
GROUP_ACTION_STATUSES.get('stale', 'stale'))
# Have a session check in, status should update to in progress
group_action.actions[0].assign_to(session.session_id)
self.assertEqual(
group_action.get_status(),
GROUP_ACTION_STATUSES.get('in progress', 'in progress'))
# Submit a response, status should be mixed success
response = Database.create_response()
group_action.actions[0].submit_response(response)
self.assertEqual(
group_action.get_status(),
GROUP_ACTION_STATUSES.get('mixed success', 'mixed success'))
# Set all actions to failed. Status should be failed.
group_action = Database.create_group_action()
session = Database.create_session(None, 0)
for action in group_action.actions:
action.assign_to(session.session_id)
self.assertEqual(group_action.get_status(),
GROUP_ACTION_STATUSES.get('failed', 'failed'))
# Create a new group action, and cancel it. Status should be cancelled
group_action = Database.create_group_action()
group_action.cancel()
self.assertEqual(group_action.get_status(),
GROUP_ACTION_STATUSES.get('cancelled', 'cancelled'))
def test_create(self):
"""
This test will pass if the action is created and content matches.
"""
target = Database.create_target()
data = APIClient.create_action(self.client, target.name,
'exec ls -al /dir')
action_id = data['action_id']
self.assertEqual(False, data['error'])
action = Database.get_action(action_id)
self.assertIsNotNone(action)
self.assertEqual(action.action_type, ACTION_TYPES.get('exec', 1))
self.assertEqual(action.command, 'ls')
self.assertListEqual(action.args, ['-al', '/dir'])
self.assertEqual(action.cancelled, False)
def test_list(self):
"""
Populates the database with sample groups, and calls the list API
function to ensure that all are returned.
"""
groups = [
Database.create_group(),
Database.create_group(),
Database.create_group(),
Database.create_group(),
Database.create_group(),
]
data = APIClient.list_groups(self.client)
self.assertEqual(data['error'], False)
self.assertEqual(sorted(list(data['groups'].keys())),
sorted([group.name for group in groups]))
def test_create_pass(self):
"""
This test will attempt to create a target model object.
"""
target = Database.create_target(self.TEST_NAME, None, self.TEST_FACTS)
self.assertEqual(target.name, self.TEST_NAME)
self.assertEqual(target.facts, self.TEST_FACTS)
def test_create(self):
"""
Test the CreateGroup API function.
"""
data = APIClient.create_group(self.client, 'SOME GROUP')
self.assertEqual(data['error'], False)
self.assertIsNotNone(Database.get_group('SOME GROUP'))
def test_members(self):
"""
Test the members function.
"""
target1 = Database.create_target()
target2 = Database.create_target()
target3 = Database.create_target()
group = Database.create_group('test_group')
group.whitelist_member(target1.name)
group.whitelist_member(target2.name)
member_names = group.members
self.assertIn(target1.name, member_names)
self.assertIn(target2.name, member_names)
self.assertNotIn(target3.name, member_names)
def test_get_by_id(self):
"""
Attempt to get a GroupAction by ID
"""
group_action1 = Database.create_group_action().group_action_id
group_action2 = GroupAction.get_by_id(group_action1)
self.assertIsNotNone(group_action2)
self.assertEqual(group_action1, group_action2.group_action_id)
def test_create_dup_uuid_fail(self):
"""
This test will attempt to create targets with the same uuids,
and it will fail as mongo should throw a not unique exception.
"""
# basic test
with self.assertRaises(NotUniqueError):
target1 = Database.create_target(None, 'AA:BB:CC:DD:EE:FF')
target2 = Database.create_target(None, 'AA:BB:CC:DD:EE:FF')
self.assertEqual(target1.uuid, target2.uuid)
# different encoding
with self.assertRaises(NotUniqueError):
target1 = Database.create_target(None, 'AA:BB:CC:DD:EE:FF'.encode('utf-8'))
target2 = Database.create_target(None, 'AA:BB:CC:DD:EE:FF'.encode('ascii'))
self.assertEqual(target1.uuid, target2.uuid)
def test_cancel(self):
"""
This test will pass if an action is successfully cancelled.
"""
group_action = Database.create_group_action()
self.assertEqual(group_action.cancelled, False)
data = APIClient.cancel_group_action(self.client,
group_action.group_action_id)
self.assertEqual(data['error'], False)
group_action = Database.get_group_action(group_action.group_action_id)
self.assertEqual(group_action.cancelled, True)
self.assertEqual(group_action.get_status(),
GROUP_ACTION_STATUSES.get('cancelled', 'cancelled'))
for action in group_action.actions:
self.assertEqual(action.cancelled, True)
self.assertEqual(action.status,
ACTION_STATUSES.get('cancelled', 'cancelled'))
def test_target_list(self):
"""
Populates the database with sample targets, and calls the list API
function to ensure that all are returned.
"""
targets = [
Database.create_target(),
Database.create_target(),
Database.create_target(),
Database.create_target(),
]
data = APIClient.list_targets(self.client)
self.assertEqual(data['error'], False)
self.assertListEqual(sorted(list(data['targets'].keys())),
sorted([target.name for target in targets]))
def test_target_groups(self):
"""
Test the get_target_groups function.
"""
target = Database.create_target()
group1 = Database.create_group('group1')
group2 = Database.create_group('group2')
Database.create_group('group3')
group1.whitelist_member(target.name)
group2.whitelist_member(target.name)
group_names = [
group.name for group in Group.get_target_groups(target.name)
]
self.assertIn('group1', group_names)
self.assertIn('group2', group_names)
self.assertNotIn('group3', group_names)