def setUp(self):
self.test_state = state.DFTimewolfState(config.Config)
self.grr_flow_module = grr_hosts.GRRFlow(self.test_state)
self.grr_flow_module.setup(
reason='random reason',
grr_server_url='http://fake/endpoint',
grr_username='******',
grr_password='******',
approvers='[email protected],[email protected]',
verify=True)
def setUp(self, mock_InitHttp):
self.mock_grr_api = mock.Mock()
mock_InitHttp.return_value = self.mock_grr_api
self.test_state = state.DFTimewolfState(config.Config)
self.grr_flow_module = grr_hosts.GRRFlow(self.test_state)
self.grr_flow_module.SetUp(
reason='random reason',
grr_server_url='http://fake/endpoint',
grr_username='******',
grr_password='******',
approvers='[email protected],[email protected]',
verify=True)
def testGetClientByHostname(self, mock_SearchClients):
"""Tests that GetClientByHostname fetches the most recent GRR client."""
mock_SearchClients.return_value = mock_grr_hosts.MOCK_CLIENT_LIST
test_state = state.DFTimewolfState()
base_grr_flow_collector = grr_hosts.GRRFlow(test_state)
base_grr_flow_collector.setup(
'random reason', 'http://fake/endpoint', ('admin', 'admin'),
'[email protected],[email protected]')
# pylint: disable=protected-access
client = base_grr_flow_collector._get_client_by_hostname('tomchop')
mock_SearchClients.assert_called_with('tomchop')
self.assertEqual(client.data.client_id,
mock_grr_hosts.MOCK_CLIENT_RECENT.data.client_id)
def testLaunchFlow(self, mock_CreateFlow):
"""Tests that CreateFlow is correctly called."""
mock_CreateFlow.return_value = mock_grr_hosts.MOCK_FLOW
test_state = state.DFTimewolfState()
base_grr_flow_collector = grr_hosts.GRRFlow(test_state)
base_grr_flow_collector.setup(
'random reason', 'http://fake/endpoint', ('admin', 'admin'),
'[email protected],[email protected]')
# pylint: disable=protected-access
flow_id = base_grr_flow_collector._launch_flow(
mock_grr_hosts.MOCK_CLIENT, "FlowName", "FlowArgs")
self.assertEqual(flow_id, 'F:12345')
mock_CreateFlow.assert_called_once_with(name="FlowName",
args="FlowArgs")
def testLaunchFlowKeepalive(self, mock_CreateFlow):
"""Tests that keepalive flows are correctly created."""
mock_CreateFlow.return_value = mock_grr_hosts.MOCK_FLOW
test_state = state.DFTimewolfState()
base_grr_flow_collector = grr_hosts.GRRFlow(test_state)
base_grr_flow_collector.setup(
'random reason', 'http://fake/endpoint', ('admin', 'admin'),
'[email protected],[email protected]')
base_grr_flow_collector.keepalive = True
# pylint: disable=protected-access
flow_id = base_grr_flow_collector._launch_flow(
mock_grr_hosts.MOCK_CLIENT, "FlowName", "FlowArgs")
self.assertEqual(flow_id, 'F:12345')
self.assertEqual(mock_CreateFlow.call_count, 2)
self.assertEqual(mock_CreateFlow.call_args,
((), {
'name': 'KeepAlive',
'args': flows_pb2.KeepAliveArgs()
}))
def testInitialization(self):
"""Tests that the collector can be initialized."""
test_state = state.DFTimewolfState()
base_grr_flow_collector = grr_hosts.GRRFlow(test_state)
self.assertIsNotNone(base_grr_flow_collector)