def _LaunchFlow(self, client, name, args):
"""Creates the specified flow, setting KeepAlive if requested.
Args:
client (object): GRR Client object on which to launch the flow.
name (str): name of the GRR flow.
args (object): arguments specific for type of flow, as defined in GRR
flow proto (FlowArgs).
Returns:
str: GRR identifier for launched flow, or an empty string if flow could
not be launched.
"""
# Start the flow and get the flow ID
flow = self._WrapGRRRequestWithApproval(client,
client.CreateFlow,
name=name,
args=args)
if not flow:
return ''
flow_id = flow.flow_id
self.logger.info('{0:s}: Scheduled'.format(flow_id))
if self.keepalive:
keepalive_flow = client.CreateFlow(name='KeepAlive',
args=flows_pb2.KeepAliveArgs())
self.logger.info('KeepAlive Flow:{0:s} scheduled'.format(
keepalive_flow.flow_id))
return flow_id
def _launch_flow(self, client, name, args):
"""Create specified flow, setting KeepAlive if requested.
Args:
client: GRR Client object on which to launch the flow.
name: string containing flow name.
args: proto (*FlowArgs) for type of flow, as defined in GRR flow proto.
Returns:
string containing ID of launched flow
"""
# Start the flow and get the flow ID
flow = self._check_approval_wrapper(client,
client.CreateFlow,
name=name,
args=args)
flow_id = flow.flow_id
print('{0:s}: Scheduled'.format(flow_id))
if self.keepalive:
keepalive_flow = client.CreateFlow(name='KeepAlive',
args=flows_pb2.KeepAliveArgs())
print('KeepAlive Flow:{0:s} scheduled'.format(
keepalive_flow.flow_id))
return flow_id
def testLaunchFlowKeepalive(self, mock_CreateFlow):
"""Tests that keepalive flows are correctly created."""
mock_CreateFlow.return_value = mock_grr_hosts.MOCK_FLOW
self.grr_flow_module.keepalive = True
flow_id = self.grr_flow_module._LaunchFlow(
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 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()
}))