def testStartVariableHuntRaisesIfMoreThanOneFlowPerClient(self): client_id = self.SetupClients(1)[0] hunt_obj = rdf_hunt_objects.Hunt(client_rate=0) hunt_obj.args.hunt_type = hunt_obj.args.HuntType.VARIABLE for index in range(2): hunt_obj.args.variable.flow_groups.append( rdf_hunt_objects.VariableHuntFlowGroup( client_ids=[client_id.Basename()], flow_name=compatibility.GetName(transfer.GetFile), flow_args=transfer.GetFileArgs(pathspec=rdf_paths.PathSpec( path="/tmp/evil_%d.txt" % index, pathtype=rdf_paths.PathSpec.PathType.OS, )))) data_store.REL_DB.WriteHuntObject(hunt_obj) with self.assertRaises(hunt.CanStartAtMostOneFlowPerClientError): hunt.StartHunt(hunt_obj.hunt_id) # Check that no flows were scheduled on the client. flows = data_store.REL_DB.ReadAllFlowObjects(client_id.Basename()) self.assertEmpty(flows)
def CreateAndStartHunt(flow_name, flow_args, creator, **kwargs): """Creates and starts a new hunt.""" # This interface takes a time when the hunt expires. However, the legacy hunt # starting interface took an rdfvalue.DurationSeconds object which was then # added to the current time to get the expiry. This check exists to make sure # we don't confuse the two. if "duration" in kwargs: precondition.AssertType(kwargs["duration"], rdfvalue.Duration) hunt_args = rdf_hunt_objects.HuntArguments.Standard( flow_name=flow_name, flow_args=rdf_structs.AnyValue.Pack(flow_args)) hunt_obj = rdf_hunt_objects.Hunt( creator=creator, args=hunt_args, create_time=rdfvalue.RDFDatetime.Now(), **kwargs) CreateHunt(hunt_obj) StartHunt(hunt_obj.hunt_id) return hunt_obj.hunt_id
def testReadHuntResultsCorrectlyAppliesWithSubstringFilter(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_obj.hunt_id) sample_results = self._SampleSingleTypeHuntResults( client_id=client_id, flow_id=flow_id, hunt_id=hunt_obj.hunt_id) self._WriteHuntResults(sample_results) results = self.db.ReadHuntResults( hunt_obj.hunt_id, 0, 100, with_substring="blah") self.assertFalse(results) results = self.db.ReadHuntResults( hunt_obj.hunt_id, 0, 100, with_substring="manufacturer") self.assertEqual( [i.payload for i in results], [i.payload for i in sample_results], ) results = self.db.ReadHuntResults( hunt_obj.hunt_id, 0, 100, with_substring="manufacturer_1") self.assertEqual([i.payload for i in results], [sample_results[1].payload])
def CreateHunt(self, flow_runner_args=None, flow_args=None, client_rule_set=None, original_object=None, client_rate=0, duration=None, creator=None, **kwargs): # Only initialize default flow_args value if default flow_runner_args value # is to be used. if not flow_runner_args: flow_args = (flow_args or transfer.GetFileArgs(pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS))) flow_runner_args = (flow_runner_args or rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__)) client_rule_set = (client_rule_set or self._CreateForemanClientRuleSet()) hunt_args = rdf_hunt_objects.HuntArguments.Standard( flow_name=flow_runner_args.flow_name, flow_args=rdf_structs.AnyValue.Pack(flow_args)) hunt_obj = rdf_hunt_objects.Hunt(creator=creator, client_rule_set=client_rule_set, original_object=original_object, client_rate=client_rate, duration=duration, args=hunt_args, **kwargs) hunt.CreateHunt(hunt_obj) return hunt_obj.hunt_id
def testVariableHuntSchedulesAllFlowsOnStart(self): client_ids = self.SetupClients(10) hunt_obj = rdf_hunt_objects.Hunt(client_rate=0) hunt_obj.args.hunt_type = hunt_obj.args.HuntType.VARIABLE for index, pair in enumerate(collection.Batch(client_ids, 2)): hunt_obj.args.variable.flow_groups.append( rdf_hunt_objects.VariableHuntFlowGroup( client_ids=[c.Basename() for c in pair], flow_name=compatibility.GetName(transfer.GetFile), flow_args=transfer.GetFileArgs(pathspec=rdf_paths.PathSpec( path="/tmp/evil_%d.txt" % index, pathtype=rdf_paths.PathSpec.PathType.OS, )))) data_store.REL_DB.WriteHuntObject(hunt_obj) hunt.StartHunt(hunt_obj.hunt_id) hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_obj.hunt_id) self.assertEqual(hunt_counters.num_clients, 10) all_flows = data_store.REL_DB.ReadHuntFlows(hunt_obj.hunt_id, 0, sys.maxsize) self.assertItemsEqual([c.Basename() for c in client_ids], [f.client_id for f in all_flows]) for index, pair in enumerate(collection.Batch(client_ids, 2)): for client_id in pair: all_flows = data_store.REL_DB.ReadAllFlowObjects( client_id.Basename()) self.assertLen(all_flows, 1) self.assertEqual(all_flows[0].flow_class_name, compatibility.GetName(transfer.GetFile)) self.assertEqual(all_flows[0].args.pathspec.path, "/tmp/evil_%d.txt" % index)
def _WriteHuntOutputPluginLogEntries(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) output_plugin_id = "1" for i in range(10): client_id, flow_id = self._SetupHuntClientAndFlow( client_id="C.12345678901234a%d" % i, hunt_id=hunt_obj.hunt_id) enum = rdf_flow_objects.FlowOutputPluginLogEntry.LogEntryType if i % 3 == 0: log_entry_type = enum.ERROR else: log_entry_type = enum.LOG self.db.WriteFlowOutputPluginLogEntries([ rdf_flow_objects.FlowOutputPluginLogEntry( client_id=client_id, flow_id=flow_id, hunt_id=hunt_obj.hunt_id, output_plugin_id=output_plugin_id, log_entry_type=log_entry_type, message="blah%d" % i) ]) return hunt_obj
def testReadHuntFlowsCorrectlyAppliesOffsetAndCountFilters(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) expectations = self._BuildFilterConditionExpectations(hunt_obj) for filter_condition, _ in expectations.items(): full_results = self.db.ReadHuntFlows( hunt_obj.hunt_id, 0, 1024, filter_condition=filter_condition) full_results_ids = [r.flow_id for r in full_results] for index in range(0, 2): for count in range(1, 3): results = self.db.ReadHuntFlows( hunt_obj.hunt_id, index, count, filter_condition=filter_condition) results_ids = [r.flow_id for r in results] expected_ids = full_results_ids[index:index + count] self.assertCountEqual( results_ids, expected_ids, "Result items do not match for " "(filter_condition=%d, index=%d, count=%d): %s vs %s" % (filter_condition, index, count, expected_ids, results_ids))
def CreateAndStartHunt(flow_name, flow_args, creator, **kwargs): """Creates and starts a new hunt.""" # This interface takes a time when the hunt expires. However, the legacy hunt # starting interface took an rdfvalue.Duration object which was then added to # the current time to get the expiry. This check exists to make sure we don't # confuse the two. if "expiry_time" in kwargs: precondition.AssertType(kwargs["expiry_time"], rdfvalue.RDFDatetime) hunt_args = rdf_hunt_objects.HuntArguments( hunt_type=rdf_hunt_objects.HuntArguments.HuntType.STANDARD, standard=rdf_hunt_objects.HuntArgumentsStandard(flow_name=flow_name, flow_args=flow_args)) hunt_obj = rdf_hunt_objects.Hunt(creator=creator, args=hunt_args, create_time=rdfvalue.RDFDatetime.Now(), **kwargs) CreateHunt(hunt_obj) StartHunt(hunt_obj.hunt_id) return hunt_obj.hunt_id
def testReadHuntResultsCorrectlyAppliesWithTypeFilter(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) sample_results = [] for i in range(10): client_id, flow_id = self._SetupHuntClientAndFlow( hunt_id=hunt_obj.hunt_id) results = self._SampleTwoTypeHuntResults( client_id=client_id, flow_id=flow_id, hunt_id=hunt_obj.hunt_id, count_per_type=1) sample_results.extend(results) self._WriteHuntResults(results) results = self.db.ReadHuntResults( hunt_obj.hunt_id, 0, 100, with_type=compatibility.GetName(rdf_client.ClientInformation)) self.assertFalse(results) results = self.db.ReadHuntResults( hunt_obj.hunt_id, 0, 100, with_type=compatibility.GetName(rdf_client.ClientSummary)) self.assertCountEqual( [i.payload for i in results], [ i.payload for i in sample_results if isinstance(i.payload, rdf_client.ClientSummary) ], )
def testReadHuntResultsCorrectlyAppliedOffsetAndCountFilters(self): hunt_obj = rdf_hunt_objects.Hunt( hunt_id=hunt.RandomHuntId(), description="foo") self.db.WriteHuntObject(hunt_obj) sample_results = [] for i in range(10): client_id, flow_id = self._SetupHuntClientAndFlow( hunt_id=hunt_obj.hunt_id) results = self._SampleSingleTypeHuntResults(client_id=client_id, count=1) sample_results.extend(results) self._WriteHuntResults({(client_id, flow_id): results}) for l in range(1, 11): for i in range(10): results = self.db.ReadHuntResults(hunt_obj.hunt_id, i, l) expected = sample_results[i:i + l] result_payloads = [x.payload for x in results] expected_payloads = [x.payload for x in expected] self.assertEqual( result_payloads, expected_payloads, "Results differ from expected (from %d, size %d): %s vs %s" % (i, l, result_payloads, expected_payloads))
def testReadHuntResultsReturnsPayloadWithMissingTypeAsSpecialValue(self): hunt_obj = rdf_hunt_objects.Hunt(hunt_id=hunt.RandomHuntId(), description="foo") self.db.WriteHuntObject(hunt_obj) client_id, flow_id = self._SetupHuntClientAndFlow( hunt_id=hunt_obj.hunt_id) sample_results = self._SampleSingleTypeHuntResults() self._WriteHuntResults({(client_id, flow_id): sample_results}) type_name = compatibility.GetName(rdf_client.ClientSummary) try: cls = rdfvalue.RDFValue.classes.pop(type_name) results = self.db.ReadHuntResults(hunt_obj.hunt_id, 0, 100) finally: rdfvalue.RDFValue.classes[type_name] = cls self.assertEqual(len(sample_results), len(results)) for r in results: self.assertTrue( isinstance(r.payload, rdf_objects.SerializedValueOfUnrecognizedType)) self.assertEqual(r.payload.type_name, type_name)
def testForemanRulesAreCorrectlyPropagatedWhenHuntStarts(self): client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.REGEX, regex=foreman_rules.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="HUNT")), foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.INTEGER, integer=foreman_rules.ForemanIntegerClientRule( field="CLIENT_CLOCK", operator=foreman_rules.ForemanIntegerClientRule.Operator. GREATER_THAN, value=1336650631137737)) ]) self.assertEmpty(data_store.REL_DB.ReadAllForemanRules()) hunt_obj = rdf_hunt_objects.Hunt(client_rule_set=client_rule_set) hunt_obj.args.hunt_type = hunt_obj.args.HuntType.STANDARD data_store.REL_DB.WriteHuntObject(hunt_obj) hunt_obj = hunt.StartHunt(hunt_obj.hunt_id) rules = data_store.REL_DB.ReadAllForemanRules() self.assertLen(rules, 1) rule = rules[0] self.assertEqual(rule.client_rule_set, client_rule_set) self.assertEqual(rule.hunt_id, hunt_obj.hunt_id) self.assertEqual(rule.expiration_time, hunt_obj.init_start_time + hunt_obj.duration) # Running a second time should not change the rules any more. with self.assertRaises(hunt.OnlyPausedHuntCanBeStartedError): hunt.StartHunt(hunt_obj.hunt_id) rules = data_store.REL_DB.ReadAllForemanRules() self.assertLen(rules, 1)
def testGetCollectedTimelinesGzchunked(self): client_id = db_test_utils.InitializeClient(data_store.REL_DB) fqdn = "foo.bar.baz" snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id snapshot.knowledge_base.fqdn = fqdn data_store.REL_DB.WriteClientSnapshot(snapshot) hunt_id = "A0B1D2C3" flow_id = "0A1B2D3C" hunt_obj = rdf_hunt_objects.Hunt() hunt_obj.hunt_id = hunt_id hunt_obj.args.standard.client_ids = [client_id] hunt_obj.args.standard.flow_name = timeline.TimelineFlow.__name__ hunt_obj.hunt_state = rdf_hunt_objects.Hunt.HuntState.PAUSED data_store.REL_DB.WriteHuntObject(hunt_obj) flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = timeline.TimelineFlow.__name__ flow_obj.create_time = rdfvalue.RDFDatetime.Now() flow_obj.parent_hunt_id = hunt_id data_store.REL_DB.WriteFlowObject(flow_obj) entry_1 = rdf_timeline.TimelineEntry() entry_1.path = "/foo/bar".encode("utf-8") entry_1.ino = 7890178901 entry_1.size = 4815162342 entry_1.atime_ns = 123 * 10**9 entry_1.mtime_ns = 234 * 10**9 entry_1.ctime_ns = 567 * 10**9 entry_1.mode = 0o654 entry_2 = rdf_timeline.TimelineEntry() entry_2.path = "/foo/baz".encode("utf-8") entry_1.ino = 8765487654 entry_2.size = 1337 entry_1.atime_ns = 987 * 10**9 entry_1.mtime_ns = 876 * 10**9 entry_1.ctime_ns = 765 * 10**9 entry_2.mode = 0o757 entries = [entry_1, entry_2] blobs = list(rdf_timeline.TimelineEntry.SerializeStream(iter(entries))) blob_ids = data_store.BLOBS.WriteBlobsWithUnknownHashes(blobs) result = rdf_timeline.TimelineResult() result.entry_batch_blob_ids = [ blob_id.AsBytes() for blob_id in blob_ids ] flow_result = rdf_flow_objects.FlowResult() flow_result.client_id = client_id flow_result.flow_id = flow_id flow_result.payload = result data_store.REL_DB.WriteFlowResults([flow_result]) buffer = io.BytesIO() fmt = timeline_pb2.ApiGetCollectedTimelineArgs.Format.RAW_GZCHUNKED self.api.Hunt(hunt_id).GetCollectedTimelines(fmt).WriteToStream(buffer) with zipfile.ZipFile(buffer, mode="r") as archive: with archive.open(f"{client_id}_{fqdn}.gzchunked", mode="r") as file: chunks = chunked.ReadAll(file) entries = list( rdf_timeline.TimelineEntry.DeserializeStream(chunks)) self.assertEqual(entries, [entry_1, entry_2])
def testGetCollectedTimelinesBody(self): client_id = db_test_utils.InitializeClient(data_store.REL_DB) fqdn = "foo.bar.quux" snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id snapshot.knowledge_base.fqdn = fqdn data_store.REL_DB.WriteClientSnapshot(snapshot) hunt_id = "B1C2E3D4" flow_id = "1B2C3E4D" hunt_obj = rdf_hunt_objects.Hunt() hunt_obj.hunt_id = hunt_id hunt_obj.args.standard.client_ids = [client_id] hunt_obj.args.standard.flow_name = timeline.TimelineFlow.__name__ hunt_obj.hunt_state = rdf_hunt_objects.Hunt.HuntState.PAUSED data_store.REL_DB.WriteHuntObject(hunt_obj) flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = timeline.TimelineFlow.__name__ flow_obj.create_time = rdfvalue.RDFDatetime.Now() flow_obj.parent_hunt_id = hunt_id data_store.REL_DB.WriteFlowObject(flow_obj) entry_1 = rdf_timeline.TimelineEntry() entry_1.path = "/bar/baz/quux".encode("utf-8") entry_1.ino = 5926273453 entry_1.size = 13373 entry_1.atime_ns = 111 * 10**9 entry_1.mtime_ns = 222 * 10**9 entry_1.ctime_ns = 333 * 10**9 entry_1.mode = 0o664 entry_2 = rdf_timeline.TimelineEntry() entry_2.path = "/bar/baz/quuz".encode("utf-8") entry_2.ino = 6037384564 entry_2.size = 13374 entry_2.atime_ns = 777 * 10**9 entry_2.mtime_ns = 888 * 10**9 entry_2.ctime_ns = 999 * 10**9 entry_2.mode = 0o777 entries = [entry_1, entry_2] blobs = list(rdf_timeline.TimelineEntry.SerializeStream(iter(entries))) blob_ids = data_store.BLOBS.WriteBlobsWithUnknownHashes(blobs) result = rdf_timeline.TimelineResult() result.entry_batch_blob_ids = [ blob_id.AsBytes() for blob_id in blob_ids ] flow_result = rdf_flow_objects.FlowResult() flow_result.client_id = client_id flow_result.flow_id = flow_id flow_result.payload = result data_store.REL_DB.WriteFlowResults([flow_result]) buffer = io.BytesIO() self.api.Hunt(hunt_id).GetCollectedTimelines( timeline_pb2.ApiGetCollectedTimelineArgs.Format.BODY ).WriteToStream(buffer) with zipfile.ZipFile(buffer, mode="r") as archive: with archive.open(f"{client_id}_{fqdn}.body", mode="r") as file: content_file = file.read().decode("utf-8") rows = list( csv.reader(io.StringIO(content_file), delimiter="|")) self.assertLen(rows, 2) self.assertEqual(rows[0][1], "/bar/baz/quux") self.assertEqual(rows[0][2], "5926273453") self.assertEqual(rows[0][3], stat.filemode(0o664)) self.assertEqual(rows[0][6], "13373") self.assertEqual(rows[0][7], "111") self.assertEqual(rows[0][8], "222") self.assertEqual(rows[0][9], "333") self.assertEqual(rows[1][1], "/bar/baz/quuz") self.assertEqual(rows[1][2], "6037384564") self.assertEqual(rows[1][3], stat.filemode(0o777)) self.assertEqual(rows[1][6], "13374") self.assertEqual(rows[1][7], "777") self.assertEqual(rows[1][8], "888") self.assertEqual(rows[1][9], "999")
def testReadHuntClientResourcesStatsCorrectlyAggregatesData(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) flow_data = [] expected_user_cpu_histogram = rdf_stats.StatsHistogram.FromBins( rdf_stats.ClientResourcesStats.CPU_STATS_BINS) expected_system_cpu_histogram = rdf_stats.StatsHistogram.FromBins( rdf_stats.ClientResourcesStats.CPU_STATS_BINS) expected_network_histogram = rdf_stats.StatsHistogram.FromBins( rdf_stats.ClientResourcesStats.NETWORK_STATS_BINS) for i in range(10): user_cpu_time = 4.5 + i system_cpu_time = 10 + i * 2 network_bytes_sent = 42 + i * 3 client_id, flow_id = self._SetupHuntClientAndFlow( flow_state=rdf_flow_objects.Flow.FlowState.FINISHED, cpu_time_used=rdf_client_stats.CpuSeconds( user_cpu_time=user_cpu_time, system_cpu_time=system_cpu_time), network_bytes_sent=network_bytes_sent, hunt_id=hunt_obj.hunt_id) expected_user_cpu_histogram.RegisterValue(user_cpu_time) expected_system_cpu_histogram.RegisterValue(system_cpu_time) expected_network_histogram.RegisterValue(network_bytes_sent) flow_data.append((client_id, flow_id, (user_cpu_time, system_cpu_time, network_bytes_sent))) usage_stats = self.db.ReadHuntClientResourcesStats(hunt_obj.hunt_id) self.assertEqual(usage_stats.user_cpu_stats.num, 10) self.assertAlmostEqual(usage_stats.user_cpu_stats.mean, 9) self.assertAlmostEqual(usage_stats.user_cpu_stats.std, 2.8722813232690143) self.assertLen(usage_stats.user_cpu_stats.histogram.bins, len(expected_user_cpu_histogram.bins)) for b, model_b in zip(usage_stats.user_cpu_stats.histogram.bins, expected_user_cpu_histogram.bins): self.assertAlmostEqual(b.range_max_value, model_b.range_max_value) self.assertEqual(b.num, model_b.num) self.assertEqual(usage_stats.system_cpu_stats.num, 10) self.assertAlmostEqual(usage_stats.system_cpu_stats.mean, 19) self.assertAlmostEqual(usage_stats.system_cpu_stats.std, 5.744562646538029) self.assertLen(usage_stats.system_cpu_stats.histogram.bins, len(expected_system_cpu_histogram.bins)) for b, model_b in zip(usage_stats.system_cpu_stats.histogram.bins, expected_system_cpu_histogram.bins): self.assertAlmostEqual(b.range_max_value, model_b.range_max_value) self.assertEqual(b.num, model_b.num) self.assertEqual(usage_stats.network_bytes_sent_stats.num, 10) self.assertAlmostEqual(usage_stats.network_bytes_sent_stats.mean, 55.5) self.assertAlmostEqual(usage_stats.network_bytes_sent_stats.std, 8.616843969807043) self.assertLen(usage_stats.network_bytes_sent_stats.histogram.bins, len(expected_network_histogram.bins)) for b, model_b in zip(usage_stats.network_bytes_sent_stats.histogram.bins, expected_network_histogram.bins): self.assertAlmostEqual(b.range_max_value, model_b.range_max_value) self.assertEqual(b.num, model_b.num) self.assertLen(usage_stats.worst_performers, 10) for worst_performer, flow_d in zip(usage_stats.worst_performers, reversed(flow_data)): client_id, flow_id, (user_cpu_time, system_cpu_time, network_bytes_sent) = flow_d self.assertEqual(worst_performer.client_id.Basename(), client_id) self.assertAlmostEqual(worst_performer.cpu_usage.user_cpu_time, user_cpu_time) self.assertAlmostEqual(worst_performer.cpu_usage.system_cpu_time, system_cpu_time) self.assertEqual(worst_performer.network_bytes_sent, network_bytes_sent) self.assertEqual(worst_performer.session_id.Path(), "/%s/%s" % (client_id, flow_id))
def testCountHuntFlowsReturnsEmptyListWhenNoFlows(self): hunt_obj = rdf_hunt_objects.Hunt(hunt_id=hunt.RandomHuntId(), description="foo") self.db.WriteHuntObject(hunt_obj) self.assertEqual(self.db.CountHuntFlows(hunt_obj.hunt_id), 0)
def testReadHuntFlowsReturnsEmptyListWhenNoFlows(self): hunt_obj = rdf_hunt_objects.Hunt(hunt_id=hunt.RandomHuntId(), description="foo") self.db.WriteHuntObject(hunt_obj) self.assertEmpty(self.db.ReadHuntFlows(hunt_obj.hunt_id, 0, 10))
def _CreateHunt(self, **kwargs): hunt_obj = rdf_hunt_objects.Hunt(creator=self.test_username, **kwargs) hunt.CreateHunt(hunt_obj) hunt_obj = hunt.StartHunt(hunt_obj.hunt_id) return hunt_obj.hunt_id
def testRawGzchunkedMultipleClients(self): client_id_1 = db_test_utils.InitializeClient(data_store.REL_DB) client_id_2 = db_test_utils.InitializeClient(data_store.REL_DB) snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id_1 snapshot.knowledge_base.fqdn = "foo.quux.com" data_store.REL_DB.WriteClientSnapshot(snapshot) snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id_2 snapshot.knowledge_base.fqdn = "foo.norf.com" data_store.REL_DB.WriteClientSnapshot(snapshot) hunt_id = "A0B1D2C3E4" hunt_obj = rdf_hunt_objects.Hunt() hunt_obj.hunt_id = hunt_id hunt_obj.args.standard.client_ids = [client_id_1, client_id_2] hunt_obj.args.standard.flow_name = timeline.TimelineFlow.__name__ hunt_obj.hunt_state = rdf_hunt_objects.Hunt.HuntState.PAUSED data_store.REL_DB.WriteHuntObject(hunt_obj) entry_1 = rdf_timeline.TimelineEntry() entry_1.path = "foo_1".encode("utf-8") entry_1.ino = 5432154321 entry_1.size = 13371 entry_1.atime_ns = 122 * 10**9 entry_1.mtime_ns = 233 * 10**9 entry_1.ctime_ns = 344 * 10**9 entry_1.mode = 0o663 entry_2 = rdf_timeline.TimelineEntry() entry_2.path = "foo_2".encode("utf-8") entry_1.ino = 7654376543 entry_2.size = 13372 entry_1.atime_ns = 788 * 10**9 entry_1.mtime_ns = 899 * 10**9 entry_1.ctime_ns = 900 * 10**9 entry_1.mode = 0o763 _WriteTimeline(client_id_1, [entry_1], hunt_id=hunt_id) _WriteTimeline(client_id_2, [entry_2], hunt_id=hunt_id) args = api_timeline.ApiGetCollectedHuntTimelinesArgs() args.hunt_id = hunt_id args.format = api_timeline.ApiGetCollectedTimelineArgs.Format.RAW_GZCHUNKED content = b"".join(self.handler.Handle(args).GenerateContent()) buffer = io.BytesIO(content) with zipfile.ZipFile(buffer, mode="r") as archive: client_filename_1 = f"{client_id_1}_foo.quux.com.gzchunked" with archive.open(client_filename_1, mode="r") as file: chunks = chunked.ReadAll(file) entries = list( rdf_timeline.TimelineEntry.DeserializeStream(chunks)) self.assertEqual(entries, [entry_1]) client_filename_2 = f"{client_id_2}_foo.norf.com.gzchunked" with archive.open(client_filename_2, mode="r") as file: chunks = chunked.ReadAll(file) entries = list( rdf_timeline.TimelineEntry.DeserializeStream(chunks)) self.assertEqual(entries, [entry_2])
def testBodyMultipleClients(self): client_id_1 = db_test_utils.InitializeClient(data_store.REL_DB) client_id_2 = db_test_utils.InitializeClient(data_store.REL_DB) snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id_1 snapshot.knowledge_base.fqdn = "bar.quux.com" data_store.REL_DB.WriteClientSnapshot(snapshot) snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id_2 snapshot.knowledge_base.fqdn = "bar.quuz.com" data_store.REL_DB.WriteClientSnapshot(snapshot) hunt_id = "B1C2E3D4F5" hunt_obj = rdf_hunt_objects.Hunt() hunt_obj.hunt_id = hunt_id hunt_obj.args.standard.client_ids = [client_id_1, client_id_2] hunt_obj.args.standard.flow_name = timeline.TimelineFlow.__name__ hunt_obj.hunt_state = rdf_hunt_objects.Hunt.HuntState.PAUSED data_store.REL_DB.WriteHuntObject(hunt_obj) entry_1 = rdf_timeline.TimelineEntry() entry_1.path = "/bar/baz/quux".encode("utf-8") entry_1.ino = 5926273453 entry_1.size = 13373 entry_1.atime_ns = 111 * 10**9 entry_1.mtime_ns = 222 * 10**9 entry_1.ctime_ns = 333 * 10**9 entry_1.mode = 0o664 entry_2 = rdf_timeline.TimelineEntry() entry_2.path = "/bar/baz/quuz".encode("utf-8") entry_2.ino = 6037384564 entry_2.size = 13374 entry_2.atime_ns = 777 * 10**9 entry_2.mtime_ns = 888 * 10**9 entry_2.ctime_ns = 999 * 10**9 entry_2.mode = 0o777 _WriteTimeline(client_id_1, [entry_1], hunt_id=hunt_id) _WriteTimeline(client_id_2, [entry_2], hunt_id=hunt_id) args = api_timeline.ApiGetCollectedHuntTimelinesArgs() args.hunt_id = hunt_id args.format = api_timeline.ApiGetCollectedTimelineArgs.Format.BODY content = b"".join(self.handler.Handle(args).GenerateContent()) buffer = io.BytesIO(content) with zipfile.ZipFile(buffer, mode="r") as archive: client_filename_1 = f"{client_id_1}_bar.quux.com.body" with archive.open(client_filename_1, mode="r") as file: content_file = file.read().decode("utf-8") rows = list( csv.reader(io.StringIO(content_file), delimiter="|")) self.assertLen(rows, 1) self.assertEqual(rows[0][1], "/bar/baz/quux") self.assertEqual(rows[0][2], "5926273453") self.assertEqual(rows[0][3], stat.filemode(0o664)) self.assertEqual(rows[0][6], "13373") self.assertEqual(rows[0][7], "111") self.assertEqual(rows[0][8], "222") self.assertEqual(rows[0][9], "333") client_filename_2 = f"{client_id_2}_bar.quuz.com.body" with archive.open(client_filename_2, mode="r") as file: content_file = file.read().decode("utf-8") rows = list( csv.reader(io.StringIO(content_file), delimiter="|")) self.assertLen(rows, 1) self.assertEqual(rows[0][1], "/bar/baz/quuz") self.assertEqual(rows[0][2], "6037384564") self.assertEqual(rows[0][3], stat.filemode(0o777)) self.assertEqual(rows[0][6], "13374") self.assertEqual(rows[0][7], "777") self.assertEqual(rows[0][8], "888") self.assertEqual(rows[0][9], "999")
def _CreateHunt(self, **kwargs): hunt_obj = rdf_hunt_objects.Hunt(creator=self.token.username, **kwargs) data_store.REL_DB.WriteHuntObject(hunt_obj) hunt_obj = hunt.StartHunt(hunt_obj.hunt_id) return hunt_obj.hunt_id