def testHuntIsStoppedIfAveragePerClientResultsCountTooHigh(self):
client_ids = self.SetupClients(5)
hunt_args = rdf_hunt_objects.HuntArguments.Standard(
flow_name=compatibility.GetName(processes.ListProcesses))
hunt_id = self._CreateHunt(
client_rule_set=foreman_rules.ForemanClientRuleSet(),
client_rate=0,
avg_results_per_client_limit=1,
args=hunt_args)
single_process = [rdf_client.Process(pid=1, exe="a.exe")]
with utils.Stubber(hunt, "MIN_CLIENTS_FOR_AVERAGE_THRESHOLDS", 4):
def CheckState(hunt_state, num_results):
hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id)
self.assertEqual(hunt_obj.hunt_state, hunt_state)
hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_id)
self.assertEqual(hunt_counters.num_results, num_results)
self._RunHunt(
client_ids[:2],
client_mock=action_mocks.ListProcessesMock(single_process))
# Hunt should still be running: we got 1 response from 2 clients. We need
# at least 3 clients to start calculating the average.
CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 2)
self._RunHunt([client_ids[2]],
client_mock=action_mocks.ListProcessesMock(
single_process * 2))
# Hunt should still be running: we got 1 response for first 2 clients and
# 2 responses for the third. This is over the limit but we need at least 4
# clients to start applying thresholds.
CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 4)
self._RunHunt([client_ids[3]],
client_mock=action_mocks.ListProcessesMock([]))
# Hunt should still be running: we got 1 response for first 2 clients,
# 2 responses for the third and zero for the 4th. This makes it 1 result
# per client on average. This is within the limit of 1.
CheckState(rdf_hunt_objects.Hunt.HuntState.STARTED, 4)
self._RunHunt(client_ids[4:5],
client_mock=action_mocks.ListProcessesMock(
single_process * 2))
# Hunt should be terminated: 5 clients did run and we got 6 results.
# That's more than the allowed average of 1.
# Note that this check also implicitly checks that the 6th client didn't
# run at all (otherwise total number of results would be 8, not 6).
CheckState(rdf_hunt_objects.Hunt.HuntState.STOPPED, 6)
self._CheckHuntStoppedNotification(
"reached the average results per client")
def _StartFlow(self, client_id, flow_cls, **kw):
if data_store.RelationalDBEnabled():
flow_id = flow.StartFlow(flow_cls=flow_cls,
client_id=client_id,
**kw)
# Lease the client message.
data_store.REL_DB.LeaseClientActionRequests(
client_id, lease_time=rdfvalue.Duration("10000s"))
# Write some responses. In the relational db, the client queue will be
# cleaned up as soon as all responses are available. Therefore we cheat
# here and make it look like the request needs more responses so it's not
# considered complete.
# Write the status first. This will mark the request as waiting for 2
# responses.
status = rdf_flow_objects.FlowStatus(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=2)
data_store.REL_DB.WriteFlowResponses([status])
# Now we read the request, adjust the number, and write it back.
reqs = data_store.REL_DB.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
req = reqs[0][0]
req.nr_responses_expected = 99
data_store.REL_DB.WriteFlowRequests([req])
# This response now won't trigger any deletion of client messages.
response = rdf_flow_objects.FlowResponse(
client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1,
payload=rdf_client.Process(name="test_process"))
data_store.REL_DB.WriteFlowResponses([response])
# This is not strictly needed as we don't display this information in the
# UI.
req.nr_responses_expected = 2
data_store.REL_DB.WriteFlowRequests([req])
return flow_id
else:
flow_id = flow.StartAFF4Flow(
flow_name=compatibility.GetName(flow_cls),
client_id=client_id,
token=self.token,
**kw).Basename()
# Have the client write some responses.
test_process = rdf_client.Process(name="test_process")
mock = flow_test_lib.MockClient(client_id,
action_mocks.ListProcessesMock(
[test_process]),
token=self.token)
mock.Next()
return flow_id
def testProcessListingOnly(self):
"""Test that the ListProcesses flow works."""
client_id = test_lib.TEST_CLIENT_ID
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=long(1333718907.167083 * 1e6))
])
flow_urn = flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name=flow_processes.ListProcesses.__name__,
token=self.token)
for s in flow_test_lib.TestFlowHelper(flow_urn,
client_mock,
client_id=client_id,
token=self.token):
session_id = s
# Check the output collection
processes = flow.GRRFlow.ResultCollectionForFID(session_id)
self.assertEqual(len(processes), 1)
self.assertEqual(processes[0].ctime, 1333718907167083L)
self.assertEqual(processes[0].cmdline, ["cmd.exe"])
def testGetExportedResultsArchiveForListProcessesFlow(self):
process = rdf_client.Process(
pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083,
RSS_size=42)
client_id = self.SetupClient(0)
flow_id = flow_test_lib.TestFlowHelper(
compatibility.GetName(processes.ListProcesses),
client_id=client_id,
client_mock=action_mocks.ListProcessesMock([process]),
token=self.token)
result_flow = self.api.Client(client_id=client_id).Flow(flow_id)
exported_archive = result_flow.GetExportedResultsArchive(
csv_plugin.CSVInstantOutputPlugin.plugin_name)
out_fd = io.BytesIO()
exported_archive.WriteToStream(out_fd)
out_fd.seek(0)
zip_fd = zipfile.ZipFile(out_fd, "r")
self.assertNotEmpty([n for n in zip_fd.namelist() if "MANIFEST" not in n])
def testDoesNotFetchDuplicates(self):
process1 = rdf_client.Process(pid=2,
ppid=1,
cmdline=["test_img.dd"],
exe=os.path.join(self.base_path,
"test_img.dd"),
ctime=long(1333718907.167083 * 1e6))
process2 = rdf_client.Process(pid=3,
ppid=1,
cmdline=["test_img.dd", "--arg"],
exe=os.path.join(self.base_path,
"test_img.dd"),
ctime=long(1333718907.167083 * 1e6))
client_mock = action_mocks.ListProcessesMock([process1, process2])
for s in flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
client_id=test_lib.TEST_CLIENT_ID,
fetch_binaries=True,
token=self.token):
session_id = s
processes = flow.GRRFlow.ResultCollectionForFID(session_id)
self.assertEqual(len(processes), 1)
def testProcessListingOnlyFleetspeak(self):
"""Test that the ListProcesses flow works with Fleetspeak."""
client_id = self.SetupClient(0)
data_store.REL_DB.WriteClientMetadata(client_id,
fleetspeak_enabled=True)
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe=r"c:\windows\cmd.exe",
ctime=1333718907167083)
])
flow_id = flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
client_id=client_id,
token=self.token)
processes = flow_test_lib.GetFlowResults(client_id, flow_id)
self.assertLen(processes, 1)
process, = processes
self.assertEqual(process.ctime, 1333718907167083)
self.assertEqual(process.cmdline, ["cmd.exe"])
def testListResultsForListProcessesFlow(self):
process = rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083,
RSS_size=42)
client_urn = self.SetupClient(0)
flow_urn = flow_test_lib.TestFlowHelper(
compatibility.GetName(processes.ListProcesses),
client_id=client_urn,
client_mock=action_mocks.ListProcessesMock([process]),
token=self.token)
if isinstance(flow_urn, rdfvalue.RDFURN):
flow_id = flow_urn.Basename()
else:
flow_id = flow_urn
result_flow = self.api.Client(
client_id=client_urn.Basename()).Flow(flow_id)
results = list(result_flow.ListResults())
self.assertLen(results, 1)
self.assertEqual(process.AsPrimitiveProto(), results[0].payload)
def testWhenFetchingFiltersOutProcessesWithoutExeAndConnectionState(self):
client_id = test_lib.TEST_CLIENT_ID
p1 = rdf_client.Process(pid=2,
ppid=1,
cmdline=["test_img.dd"],
ctime=long(1333718907.167083 * 1e6))
p2 = rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=long(1333718907.167083 * 1e6),
connections=rdf_client.NetworkConnection(
family="INET", state="ESTABLISHED"))
client_mock = action_mocks.ListProcessesMock([p1, p2])
for s in flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
fetch_binaries=True,
client_id=client_id,
connection_states=["LISTEN"],
token=self.token):
session_id = s
# No output matched.
processes = flow.GRRFlow.ResultCollectionForFID(session_id)
self.assertEqual(len(processes), 0)
def testProcessListingOnly(self):
"""Test that the ListProcesses flow works."""
client_id = self.SetupClient(0)
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(
pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083)
])
flow_urn = flow.StartAFF4Flow(
client_id=client_id,
flow_name=flow_processes.ListProcesses.__name__,
token=self.token)
session_id = flow_test_lib.TestFlowHelper(
flow_urn, client_mock, client_id=client_id, token=self.token)
# Check the output collection
processes = flow.GRRFlow.ResultCollectionForFID(session_id)
self.assertLen(processes, 1)
self.assertEqual(processes[0].ctime, 1333718907167083)
self.assertEqual(processes[0].cmdline, ["cmd.exe"])
def testWhenFetchingIgnoresMissingFiles(self):
client_id = self.SetupClient(0)
process1 = rdf_client.Process(
pid=2,
ppid=1,
cmdline=["test_img.dd"],
exe=os.path.join(self.base_path, "test_img.dd"),
ctime=1333718907167083)
process2 = rdf_client.Process(
pid=2,
ppid=1,
cmdline=["file_that_does_not_exist"],
exe=os.path.join(self.base_path, "file_that_does_not_exist"),
ctime=1333718907167083)
client_mock = action_mocks.ListProcessesMock([process1, process2])
session_id = flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
client_id=client_id,
fetch_binaries=True,
creator=self.test_username,
check_flow_errors=False)
binaries = flow_test_lib.GetFlowResults(client_id, session_id)
self.assertLen(binaries, 1)
self.assertEqual(binaries[0].pathspec.path, process1.exe)
def testDoesNotFetchDuplicates(self):
client_id = self.SetupClient(0)
process1 = rdf_client.Process(
pid=2,
ppid=1,
cmdline=["test_img.dd"],
exe=os.path.join(self.base_path, "test_img.dd"),
ctime=1333718907167083)
process2 = rdf_client.Process(
pid=3,
ppid=1,
cmdline=["test_img.dd", "--arg"],
exe=os.path.join(self.base_path, "test_img.dd"),
ctime=1333718907167083)
client_mock = action_mocks.ListProcessesMock([process1, process2])
session_id = flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
client_id=client_id,
fetch_binaries=True,
creator=self.test_username)
processes = flow_test_lib.GetFlowResults(client_id, session_id)
self.assertLen(processes, 1)
def testPidFiltering(self):
client_id = self.SetupClient(0)
proc_foo = rdf_client.Process()
proc_foo.pid = 42
proc_foo.exe = "/usr/bin/foo"
proc_bar = rdf_client.Process()
proc_bar.pid = 108
proc_bar.exe = "/usr/bin/bar"
proc_baz = rdf_client.Process()
proc_baz.pid = 1337
proc_baz.exe = "/usr/bin/baz"
args = flow_processes.ListProcessesArgs()
args.pids = [42, 1337]
client_mock = action_mocks.ListProcessesMock([proc_foo, proc_bar, proc_baz])
flow_id = flow_test_lib.StartAndRunFlow(
flow_processes.ListProcesses,
client_mock=client_mock,
client_id=client_id,
flow_args=args,
)
results = flow_test_lib.GetFlowResults(client_id=client_id, flow_id=flow_id)
self.assertLen(results, 2)
result_exes = {result.exe for result in results}
self.assertIn("/usr/bin/foo", result_exes)
self.assertIn("/usr/bin/baz", result_exes)
self.assertNotIn("/usr/bin/bar", result_exes)
def testListResultsForListProcessesFlow(self):
process = rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083,
RSS_size=42)
client_urn = self.SetupClient(0)
client_mock = action_mocks.ListProcessesMock([process])
flow_urn = flow.StartAFF4Flow(
client_id=client_urn,
flow_name=processes.ListProcesses.__name__,
token=self.token)
flow_test_lib.TestFlowHelper(flow_urn,
client_mock,
client_id=client_urn,
token=self.token)
result_flow = self.api.Client(client_id=client_urn.Basename()).Flow(
flow_urn.Basename())
results = list(result_flow.ListResults())
self.assertEqual(len(results), 1)
self.assertEqual(process.AsPrimitiveProto(), results[0].payload)
def ProcessFlow():
time.sleep(1)
client_mock = action_mocks.ListProcessesMock([])
flow_test_lib.TestFlowHelper(flow_urn,
client_mock,
client_id=client_urn,
token=self.token)
def testWhenFetchingFiltersOutProcessesWithoutExeAndConnectionState(self):
client_id = self.SetupClient(0)
p1 = rdf_client.Process(
pid=2, ppid=1, cmdline=["test_img.dd"], ctime=1333718907167083)
p2 = rdf_client.Process(
pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083,
connections=[
rdf_client_network.NetworkConnection(
family="INET", state="ESTABLISHED")
])
client_mock = action_mocks.ListProcessesMock([p1, p2])
session_id = flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
fetch_binaries=True,
client_id=client_id,
connection_states=["LISTEN"],
creator=self.test_username)
# No output matched.
processes = flow_test_lib.GetFlowResults(client_id, session_id)
self.assertEmpty(processes)
def testWhenFetchingIgnoresMissingFiles(self):
process1 = rdf_client.Process(pid=2,
ppid=1,
cmdline=["test_img.dd"],
exe=os.path.join(self.base_path,
"test_img.dd"),
ctime=long(1333718907.167083 * 1e6))
process2 = rdf_client.Process(pid=2,
ppid=1,
cmdline=["file_that_does_not_exist"],
exe=os.path.join(
self.base_path,
"file_that_does_not_exist"),
ctime=long(1333718907.167083 * 1e6))
client_mock = action_mocks.ListProcessesMock([process1, process2])
for s in flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
client_id=test_lib.TEST_CLIENT_ID,
fetch_binaries=True,
token=self.token,
check_flow_errors=False):
session_id = s
results = flow.GRRFlow.ResultCollectionForFID(session_id)
binaries = list(results)
self.assertEqual(len(binaries), 1)
self.assertEqual(binaries[0].pathspec.path, process1.exe)
def testProcessListingWithFilter(self):
"""Test that the ListProcesses flow works with filter."""
client_id = self.SetupClient(0)
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(
pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=long(1333718907.167083 * 1e6)),
rdf_client.Process(
pid=3,
ppid=1,
cmdline=["cmd2.exe"],
exe="c:\\windows\\cmd2.exe",
ctime=long(1333718907.167083 * 1e6)),
rdf_client.Process(
pid=4,
ppid=1,
cmdline=["missing_exe.exe"],
ctime=long(1333718907.167083 * 1e6)),
rdf_client.Process(
pid=5,
ppid=1,
cmdline=["missing2_exe.exe"],
ctime=long(1333718907.167083 * 1e6))
])
flow_urn = flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name=flow_processes.ListProcesses.__name__,
filename_regex=r".*cmd2.exe",
token=self.token)
session_id = flow_test_lib.TestFlowHelper(
flow_urn, client_mock, client_id=client_id, token=self.token)
# Expect one result that matches regex
processes = flow.GRRFlow.ResultCollectionForFID(session_id)
self.assertEqual(len(processes), 1)
self.assertEqual(processes[0].ctime, 1333718907167083L)
self.assertEqual(processes[0].cmdline, ["cmd2.exe"])
# Expect two skipped results
logs = flow.GRRFlow.LogCollectionForFID(flow_urn)
for log in logs:
if "Skipped 2" in log.log_message:
return
raise RuntimeError("Skipped process not mentioned in logs")
def testProcessListingOnlyFleetspeak(self):
"""Test that the ListProcesses flow works with Fleetspeak."""
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe=r"c:\windows\cmd.exe",
ctime=1333718907167083)
])
client_mock.mock_task_queue = []
def SendCallback(fs_msg):
pb_msg = jobs_pb2.GrrMessage()
fs_msg.data.Unpack(pb_msg)
msg = rdf_flows.GrrMessage.FromSerializedString(
pb_msg.SerializeToString())
client_mock.mock_task_queue.append(msg)
service_name = "GRR"
fake_service_client = _FakeGRPCServiceClient(
service_name, send_callback=SendCallback)
fleetspeak_connector.Reset()
fleetspeak_connector.Init(service_client=fake_service_client)
with mock.patch.object(
fake_service_client.outgoing,
"InsertMessage",
wraps=fake_service_client.outgoing.InsertMessage):
flow_urn = flow.StartFlow(
client_id=self.client_id,
flow_name=flow_processes.ListProcesses.__name__,
token=self.token)
session_id = flow_test_lib.TestFlowHelper(flow_urn,
client_mock,
client_id=self.client_id,
token=self.token)
fleetspeak_connector.CONN.outgoing.InsertMessage.assert_called()
# Check the output collection
processes = flow.GRRFlow.ResultCollectionForFID(session_id)
self.assertEqual(len(processes), 1)
process, = processes
self.assertEqual(process.ctime, 1333718907167083)
self.assertEqual(process.cmdline, ["cmd.exe"])
def Run(self):
client_id = self.SetupClient(0)
with test_lib.FakeTime(42):
flow_id = flow_test_lib.StartFlow(
processes.ListProcesses, client_id, creator=self.test_username)
test_process = rdf_client.Process(name="test_process")
mock = flow_test_lib.MockClient(
client_id, action_mocks.ListProcessesMock([test_process]))
mock.Next()
replace = api_regression_test_lib.GetFlowTestReplaceDict(client_id, flow_id)
self.Check(
"ListFlowRequests",
args=flow_plugin.ApiListFlowRequestsArgs(
client_id=client_id, flow_id=flow_id),
replace=replace)
def testProcessListingWithFilter(self):
"""Test that the ListProcesses flow works with filter."""
client_id = self.SetupClient(0)
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083),
rdf_client.Process(pid=3,
ppid=1,
cmdline=["cmd2.exe"],
exe="c:\\windows\\cmd2.exe",
ctime=1333718907167083),
rdf_client.Process(pid=4,
ppid=1,
cmdline=["missing_exe.exe"],
ctime=1333718907167083),
rdf_client.Process(pid=5,
ppid=1,
cmdline=["missing2_exe.exe"],
ctime=1333718907167083)
])
session_id = flow_test_lib.TestFlowHelper(compatibility.GetName(
flow_processes.ListProcesses),
client_mock,
client_id=client_id,
creator=self.test_username,
filename_regex=r".*cmd2.exe")
# Expect one result that matches regex
processes = flow_test_lib.GetFlowResults(client_id, session_id)
self.assertLen(processes, 1)
self.assertEqual(processes[0].ctime, 1333718907167083)
self.assertEqual(processes[0].cmdline, ["cmd2.exe"])
# Expect two skipped results
logs = data_store.REL_DB.ReadFlowLogEntries(client_id, session_id, 0,
100)
for log in logs:
if "Skipped 2" in log.message:
return
raise RuntimeError("Skipped process not mentioned in logs")
def testProcessListingFilterConnectionState(self):
client_id = self.SetupClient(0)
p1 = rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083,
connections=[
rdf_client_network.NetworkConnection(
family="INET", state="CLOSED")
])
p2 = rdf_client.Process(pid=3,
ppid=1,
cmdline=["cmd2.exe"],
exe="c:\\windows\\cmd2.exe",
ctime=1333718907167083,
connections=[
rdf_client_network.NetworkConnection(
family="INET", state="LISTEN")
])
p3 = rdf_client.Process(pid=4,
ppid=1,
cmdline=["missing_exe.exe"],
ctime=1333718907167083,
connections=[
rdf_client_network.NetworkConnection(
family="INET", state="ESTABLISHED")
])
client_mock = action_mocks.ListProcessesMock([p1, p2, p3])
flow_urn = flow.StartFlow(
client_id=client_id,
flow_name=flow_processes.ListProcesses.__name__,
connection_states=["ESTABLISHED", "LISTEN"],
token=self.token)
session_id = flow_test_lib.TestFlowHelper(flow_urn,
client_mock,
client_id=client_id,
token=self.token)
processes = flow.GRRFlow.ResultCollectionForFID(session_id)
self.assertEqual(len(processes), 2)
states = set()
for process in processes:
states.add(str(process.connections[0].state))
self.assertItemsEqual(states, ["ESTABLISHED", "LISTEN"])
def testProcessListingFilterConnectionState(self):
client_id = self.SetupClient(0)
p1 = rdf_client.Process(
pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083,
connections=[
rdf_client_network.NetworkConnection(family="INET", state="CLOSED")
])
p2 = rdf_client.Process(
pid=3,
ppid=1,
cmdline=["cmd2.exe"],
exe="c:\\windows\\cmd2.exe",
ctime=1333718907167083,
connections=[
rdf_client_network.NetworkConnection(family="INET", state="LISTEN")
])
p3 = rdf_client.Process(
pid=4,
ppid=1,
cmdline=["missing_exe.exe"],
ctime=1333718907167083,
connections=[
rdf_client_network.NetworkConnection(
family="INET", state="ESTABLISHED")
])
client_mock = action_mocks.ListProcessesMock([p1, p2, p3])
session_id = flow_test_lib.TestFlowHelper(
compatibility.GetName(flow_processes.ListProcesses),
client_mock,
client_id=client_id,
creator=self.test_username,
connection_states=["ESTABLISHED", "LISTEN"])
processes = flow_test_lib.GetFlowResults(client_id, session_id)
self.assertLen(processes, 2)
states = set()
for process in processes:
states.add(str(process.connections[0].state))
self.assertCountEqual(states, ["ESTABLISHED", "LISTEN"])
def testProcessListingOnlyFleetspeak(self):
"""Test that the ListProcesses flow works with Fleetspeak."""
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe=r"c:\windows\cmd.exe",
ctime=1333718907167083)
])
client_mock.mock_task_queue = []
def SendCallback(fs_msg):
pb_msg = jobs_pb2.GrrMessage()
fs_msg.data.Unpack(pb_msg)
msg = rdf_flows.GrrMessage.FromSerializedString(
pb_msg.SerializeToString())
client_mock.mock_task_queue.append(msg)
fake_conn = _FakeGRPCServiceClient(FS_SERVICE_NAME,
send_callback=SendCallback)
with fleetspeak_test_lib.ConnectionOverrider(fake_conn):
with mock.patch.object(fake_conn.outgoing,
"InsertMessage",
wraps=fake_conn.outgoing.InsertMessage):
session_id = flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
client_id=self.client_id,
token=self.token)
fleetspeak_connector.CONN.outgoing.InsertMessage.assert_called(
)
# Check the output collection
processes = flow_test_lib.GetFlowResults(self.client_id,
session_id)
self.assertLen(processes, 1)
process, = processes
self.assertEqual(process.ctime, 1333718907167083)
self.assertEqual(process.cmdline, ["cmd.exe"])
def testFetchesAndStoresBinary(self):
process = rdf_client.Process(pid=2,
ppid=1,
cmdline=["test_img.dd"],
exe=os.path.join(self.base_path,
"test_img.dd"),
ctime=1333718907167083)
client_mock = action_mocks.ListProcessesMock([process])
session_id = flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
client_id=self.SetupClient(0),
fetch_binaries=True,
token=self.token)
results = flow.GRRFlow.ResultCollectionForFID(session_id)
binaries = list(results)
self.assertEqual(len(binaries), 1)
self.assertEqual(binaries[0].pathspec.path, process.exe)
self.assertEqual(binaries[0].st_size, os.stat(process.exe).st_size)
def testProcessListingOnly(self):
"""Test that the ListProcesses flow works."""
client_id = self.SetupClient(0)
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe="c:\\windows\\cmd.exe",
ctime=1333718907167083)
])
session_id = flow_test_lib.TestFlowHelper(compatibility.GetName(
flow_processes.ListProcesses),
client_mock,
client_id=client_id,
creator=self.test_username)
processes = flow_test_lib.GetFlowResults(client_id, session_id)
self.assertLen(processes, 1)
self.assertEqual(processes[0].ctime, 1333718907167083)
self.assertEqual(processes[0].cmdline, ["cmd.exe"])
def testFetchesAndStoresBinary(self):
client_id = self.SetupClient(0)
process = rdf_client.Process(
pid=2,
ppid=1,
cmdline=["test_img.dd"],
exe=os.path.join(self.base_path, "test_img.dd"),
ctime=1333718907167083)
client_mock = action_mocks.ListProcessesMock([process])
session_id = flow_test_lib.TestFlowHelper(
flow_processes.ListProcesses.__name__,
client_mock,
client_id=client_id,
fetch_binaries=True,
creator=self.test_username)
binaries = flow_test_lib.GetFlowResults(client_id, session_id)
self.assertLen(binaries, 1)
self.assertEqual(binaries[0].pathspec.path, process.exe)
self.assertEqual(binaries[0].st_size, os.stat(process.exe).st_size)
def _StartFlow(self, client_id, flow_cls, **kw):
if data_store.RelationalDBFlowsEnabled():
flow_id = flow.StartFlow(flow_cls=flow_cls,
client_id=client_id,
**kw)
# Lease the client message.
data_store.REL_DB.LeaseClientMessages(
client_id, lease_time=rdfvalue.Duration("10000s"))
# Write some responses.
response = rdf_flow_objects.FlowResponse(
client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1,
payload=rdf_client.Process(name="test_process"))
status = rdf_flow_objects.FlowStatus(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=2)
data_store.REL_DB.WriteFlowResponses([response, status])
return flow_id
else:
flow_id = flow.StartAFF4Flow(
flow_name=compatibility.GetName(flow_cls),
client_id=client_id,
token=self.token,
**kw).Basename()
# Have the client write some responses.
test_process = rdf_client.Process(name="test_process")
mock = flow_test_lib.MockClient(client_id,
action_mocks.ListProcessesMock(
[test_process]),
token=self.token)
mock.Next()
return flow_id
def RunOnClients(client_ids, num_processes):
client_mock = action_mocks.ListProcessesMock(
[rdf_client.Process(pid=1, exe="a.exe")] * num_processes)
self.AssignTasksToClients(client_ids)
hunt_test_lib.TestHuntHelper(
client_mock, client_ids, check_flow_errors=False, token=self.token)
def ProcessFlow():
time.sleep(1)
client_mock = action_mocks.ListProcessesMock([])
flow_test_lib.FinishAllFlowsOnClient(client_urn,
client_mock=client_mock)
