def _InitializeFlow(self, **kwargs) -> str:
return flow_test_lib.TestFlowHelper(
osquery_flow.OsqueryFlow.__name__,
action_mocks.OsqueryClientMock(),
client_id=self.client_id,
creator=self.test_username,
query="<<<FAKE OSQUERY FLOW QUERY PLACEHOLDER>>>",
**kwargs)
def _InitializeFlow(
self,
query: str,
file_collection_columns: List[str],
**kwargs,
) -> None:
session_id = flow_test_lib.TestFlowHelper(
osquery_flow.OsqueryFlow.__name__,
action_mocks.OsqueryClientMock(),
client_id=self.client_id,
token=self.token,
query=query,
file_collection_columns=file_collection_columns,
**kwargs)
return session_id
def _RunOsqueryExportResults(self,
stdout: str) -> utils.BinaryChunkIterator:
client_id = self.SetupClient(0)
with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=""):
flow_id = flow_test_lib.TestFlowHelper(
osquery_flow.OsqueryFlow.__name__,
action_mocks.OsqueryClientMock(),
client_id=client_id,
token=self.token,
query="doesn't matter")
result_flow = self.api.Client(client_id=client_id).Flow(flow_id)
result_flow.WaitUntilDone()
format_csv = api_osquery_pb2.ApiGetOsqueryResultsArgs.Format.CSV
return result_flow.GetOsqueryResults(format_csv)
def testFailure(self):
stderr = "Error: query syntax error"
with osquery_test_lib.FakeOsqueryiOutput(stdout="", stderr=stderr):
flow_id = flow_test_lib.StartFlow(
flow_cls=osquery_flow.OsqueryFlow,
client_id=self.client_id,
creator=self.test_username,
query="<<<FAKE OSQUERY FLOW QUERY PLACEHOLDER>>>")
with self.assertRaises(RuntimeError):
flow_test_lib.RunFlow(
client_id=self.client_id,
flow_id=flow_id,
client_mock=action_mocks.OsqueryClientMock())
flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(flow.flow_state,
rdf_flow_objects.Flow.FlowState.ERROR)
self.assertIn(stderr, flow.error_message)