def testGetFileIndicatesCollectedSizeAfterCollection(self):
# Find the file with FileFinder stat action, so that we can reference it
# and trigger "Collect" operation on it.
client_ref = self.api.Client(client_id=self.client_id)
args = rdf_file_finder.FileFinderArgs(
paths=[os.path.join(self.base_path, "numbers.txt")],
action=rdf_file_finder.FileFinderAction.Stat()).AsPrimitiveProto()
client_ref.CreateFlow(name=file_finder.FileFinder.__name__, args=args)
client_mock = action_mocks.FileFinderClientMock()
flow_test_lib.FinishAllFlowsOnClient(
self.client_id, client_mock=client_mock)
f = client_ref.File("fs/os" + os.path.join(self.base_path, "numbers.txt"))
with flow_test_lib.TestWorker():
operation = f.Collect()
self.assertEqual(operation.GetState(), operation.STATE_RUNNING)
flow_test_lib.FinishAllFlowsOnClient(
self.client_id, client_mock=client_mock)
self.assertEqual(operation.GetState(), operation.STATE_FINISHED)
f = f.Get()
self.assertNotEqual(f.data.hash.sha256, b"")
self.assertGreater(f.data.hash.num_bytes, 0)
self.assertGreater(f.data.last_collected, 0)
self.assertGreater(f.data.last_collected_size, 0)
def testClickingOnInterrogateStartsInterrogateFlow(self):
self.Open("/#/clients/%s" % self.client_id)
# A click on the Interrogate button starts a flow, disables the button and
# shows a loading icon within the button.
self.Click("css=button:contains('Interrogate'):not([disabled])")
self.WaitUntil(self.IsElementPresent,
"css=button:contains('Interrogate')[disabled]")
self.WaitUntil(self.IsElementPresent,
"css=button:contains('Interrogate') i")
# Get the started flow and finish it, this will re-enable the button.
flow_test_lib.FinishAllFlowsOnClient(self.client_id,
check_flow_errors=False)
self.WaitUntilNot(self.IsElementPresent,
"css=button:contains('Interrogate')[disabled]")
# Check if an Interrogate flow was started.
self.Click("css=a[grrtarget='client.flows']")
self.Click("css=td:contains('Interrogate')")
self.WaitUntilContains(
discovery.Interrogate.__name__, self.GetText,
"css=table td.proto_key:contains('Flow name') "
"~ td.proto_value")
def _RunListNamedPipesFlow(
self,
args: pipes.ListNamedPipesFlowArgs,
pipe_results: Iterable[rdf_client.NamedPipe] = (),
proc_results: Iterable[rdf_client.Process] = (),
) -> Sequence[pipes.ListNamedPipesFlowResult]:
"""Runs the flow listing named pipes with the given fake action results."""
class ActionMock(action_mocks.ActionMock):
def ListNamedPipes(
self,
args: None,
) -> Iterator[rdf_client.NamedPipe]:
del args # Unused.
for pipe in pipe_results:
yield pipe
def ListProcesses(
self,
args: None,
) -> Iterator[rdf_client.Process]:
del args # Unused.
for proc in proc_results:
yield proc
flow_id = flow_test_lib.StartAndRunFlow(
pipes.ListNamedPipesFlow,
client_mock=ActionMock(),
client_id=self.client_id,
flow_args=args,
)
flow_test_lib.FinishAllFlowsOnClient(self.client_id)
return flow_test_lib.GetFlowResults(self.client_id, flow_id)
def testValidatesParsersWereNotApplied(self, db: abstract_db.Database):
token = _CreateToken(db)
client_id = db_test_utils.InitializeClient(db)
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(self.ECHO1337_ARTIFACT)
flow_args = rdf_artifacts.ArtifactCollectorFlowArgs()
flow_args.artifact_list = [self.ECHO1337_ARTIFACT.name]
flow_args.apply_parsers = True
flow_id = flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
self.FakeExecuteCommand(),
client_id=client_id,
args=flow_args,
token=token)
flow_test_lib.FinishAllFlowsOnClient(client_id)
args = flow_plugin.ApiListParsedFlowResultsArgs()
args.client_id = client_id
args.flow_id = flow_id
with self.assertRaisesRegex(ValueError, "already parsed"):
self.handler.Handle(args, token=token)
def testArtifactCollectorIndicatesCollectedSizeAfterCollection(self):
registry_stub = artifact_registry.ArtifactRegistry()
source = rdf_artifacts.ArtifactSource(
type=rdf_artifacts.ArtifactSource.SourceType.FILE,
attributes={
"paths": [os.path.join(self.base_path, "numbers.txt")],
})
artifact = rdf_artifacts.Artifact(
name="FakeArtifact", sources=[source], doc="fake artifact doc")
registry_stub.RegisterArtifact(artifact)
client_ref = self.api.Client(client_id=self.client_id)
with mock.patch.object(artifact_registry, "REGISTRY", registry_stub):
args = rdf_artifacts.ArtifactCollectorFlowArgs(
artifact_list=["FakeArtifact"]).AsPrimitiveProto()
client_ref.CreateFlow(
name=collectors.ArtifactCollectorFlow.__name__, args=args)
client_mock = action_mocks.FileFinderClientMock()
flow_test_lib.FinishAllFlowsOnClient(
self.client_id, client_mock=client_mock)
f = client_ref.File("fs/os" +
os.path.join(self.base_path, "numbers.txt")).Get()
self.assertNotEqual(f.data.hash.sha256, b"")
self.assertGreater(f.data.hash.num_bytes, 0)
self.assertGreater(f.data.last_collected, 0)
self.assertGreater(f.data.last_collected_size, 0)
def testEmptyResults(self, db: abstract_db.Database):
token = _CreateToken(db)
client_id = db_test_utils.InitializeClient(db)
fake_artifact = rdf_artifacts.Artifact(name="FakeArtifact",
doc="Lorem ipsum.",
sources=[])
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(fake_artifact)
flow_args = rdf_artifacts.ArtifactCollectorFlowArgs()
flow_args.artifact_list = [fake_artifact.name]
flow_args.apply_parsers = False
flow_id = flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
self.FakeExecuteCommand(),
client_id=client_id,
args=flow_args,
token=token)
flow_test_lib.FinishAllFlowsOnClient(client_id)
args = flow_plugin.ApiListParsedFlowResultsArgs(client_id=client_id,
flow_id=flow_id,
offset=0,
count=1024)
result = self.handler.Handle(args, token=token)
self.assertEmpty(result.errors)
self.assertEmpty(result.items)
def testLogsWarningIfBtimeNotSupported(self, db: abstract_db.Database):
client_id = self.client_id
db.WriteClientMetadata(client_id, fleetspeak_enabled=True)
snapshot = rdf_objects.ClientSnapshot()
snapshot.client_id = client_id
snapshot.knowledge_base.os = "Linux"
snapshot.startup_info.client_info.timeline_btime_support = False
db.WriteClientSnapshot(snapshot)
with temp.AutoTempDirPath() as tempdir:
args = rdf_timeline.TimelineArgs(root=tempdir.encode("utf-8"))
flow_id = flow_test_lib.TestFlowHelper(
timeline_flow.TimelineFlow.__name__,
action_mocks.ActionMock(timeline_action.Timeline),
client_id=client_id,
token=self.token,
args=args)
flow_test_lib.FinishAllFlowsOnClient(client_id)
log_entries = db.ReadFlowLogEntries(client_id,
flow_id,
offset=0,
count=1)
self.assertLen(log_entries, 1)
self.assertRegex(log_entries[0].message, "birth time is not supported")
def testExecutePythonHackWithResult(self):
client_id = db_test_utils.InitializeClient(data_store.REL_DB)
code = """
magic_return_str = str(py_args["foobar"])
"""
maintenance_utils.UploadSignedConfigBlob(
content=code.encode("utf-8"),
aff4_path="aff4:/config/python_hacks/quux")
flow_id = flow_test_lib.TestFlowHelper(
administrative.ExecutePythonHack.__name__,
client_mock=action_mocks.ActionMock(standard.ExecutePython),
client_id=client_id,
hack_name="quux",
py_args={"foobar": 42},
token=self.token)
flow_test_lib.FinishAllFlowsOnClient(client_id=client_id)
results = flow_test_lib.GetFlowResults(client_id=client_id,
flow_id=flow_id)
self.assertLen(results, 1)
self.assertIsInstance(results[0],
administrative.ExecutePythonHackResult)
self.assertEqual(results[0].result_string, "42")
def _Collect(self, path: str, signed_url: str) -> str:
"""Runs the large file collection flow.
Args:
path: A path to the file to collect.
signed_url: A signed URL to the where the file should be sent to.
Returns:
An identifier of the flow that was created.
"""
args = large_file.CollectLargeFileFlowArgs()
args.signed_url = signed_url
args.path_spec.pathtype = rdf_paths.PathSpec.PathType.OS
args.path_spec.path = path
action_mock = action_mocks.ActionMock.With({
"CollectLargeFile":
large_file_action.CollectLargeFileAction,
})
flow_id = flow_test_lib.TestFlowHelper(
large_file.CollectLargeFileFlow.__name__,
action_mock,
client_id=self.client_id,
creator=self.test_username,
args=args)
flow_test_lib.FinishAllFlowsOnClient(self.client_id)
return flow_id
def testNoLogsIfBtimeSupported(self, db: abstract_db.Database):
client_id = self.client_id
db.WriteClientMetadata(client_id, fleetspeak_enabled=True)
snapshot = rdf_objects.ClientSnapshot()
snapshot.client_id = client_id
snapshot.knowledge_base.os = "Linux"
snapshot.startup_info.client_info.timeline_btime_support = True
db.WriteClientSnapshot(snapshot)
with temp.AutoTempDirPath() as tempdir:
args = rdf_timeline.TimelineArgs(root=tempdir.encode("utf-8"))
flow_id = flow_test_lib.TestFlowHelper(
timeline_flow.TimelineFlow.__name__,
action_mocks.ActionMock(timeline_action.Timeline),
client_id=client_id,
creator=self.test_username,
args=args)
flow_test_lib.FinishAllFlowsOnClient(client_id)
log_entries = db.ReadFlowLogEntries(client_id,
flow_id,
offset=0,
count=1)
self.assertEmpty(log_entries)
def testValidatesFlowName(self, db: abstract_db.Database):
context = _CreateContext(db)
class FakeFlow(flow_base.FlowBase):
def Start(self):
self.CallState("End")
def End(self, responses: flow_responses.Responses) -> None:
del responses # Unused.
client_id = db_test_utils.InitializeClient(db)
flow_id = flow_test_lib.TestFlowHelper(
FakeFlow.__name__,
client_id=client_id,
token=access_control.ACLToken(username=context.username))
flow_test_lib.FinishAllFlowsOnClient(client_id)
args = flow_plugin.ApiListParsedFlowResultsArgs()
args.client_id = client_id
args.flow_id = flow_id
with self.assertRaisesRegex(ValueError, "artifact-collector"):
self.handler.Handle(args, context=context)
def testEdrAgentCollection(self):
client_id = db_test_utils.InitializeClient(data_store.REL_DB)
artifact_source = rdf_artifacts.ArtifactSource()
artifact_source.type = rdf_artifacts.ArtifactSource.SourceType.COMMAND
artifact_source.attributes = {"cmd": "/bin/echo", "args": ["1337"]}
artifact = rdf_artifacts.Artifact()
artifact.name = "Foo"
artifact.doc = "Lorem ipsum."
artifact.sources = [artifact_source]
class FooParser(parsers.SingleResponseParser):
supported_artifacts = ["Foo"]
def ParseResponse(
self,
knowledge_base: rdf_client.KnowledgeBase,
response: rdf_client_action.ExecuteResponse,
) -> Iterator[rdf_client.EdrAgent]:
edr_agent = rdf_client.EdrAgent()
edr_agent.name = "echo"
edr_agent.agent_id = response.stdout.decode("utf-8")
yield edr_agent
class EchoActionMock(action_mocks.InterrogatedClient):
def ExecuteCommand(
self,
args: rdf_client_action.ExecuteRequest,
) -> Iterable[rdf_client_action.ExecuteResponse]:
response = rdf_client_action.ExecuteResponse()
response.stdout = " ".join(args.args).encode("utf-8")
response.exit_status = 0
return [response]
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(artifact)
with test_lib.ConfigOverrider({"Artifacts.edr_agents": ["Foo"]}):
with parser_test_lib._ParserContext("Foo", FooParser):
flow_test_lib.TestFlowHelper(
discovery.Interrogate.__name__,
client_mock=EchoActionMock(),
client_id=client_id,
creator=self.test_username)
flow_test_lib.FinishAllFlowsOnClient(client_id)
snapshot = data_store.REL_DB.ReadClientSnapshot(client_id)
self.assertLen(snapshot.edr_agents, 1)
self.assertEqual(snapshot.edr_agents[0].name, "echo")
self.assertEqual(snapshot.edr_agents[0].agent_id, "1337")
def testParsesArtifactCollectionResults(self, db: abstract_db.Database):
context = _CreateContext(db)
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(self.ECHO1337_ARTIFACT)
flow_args = rdf_artifacts.ArtifactCollectorFlowArgs()
flow_args.artifact_list = [self.ECHO1337_ARTIFACT.name]
flow_args.apply_parsers = False
client_id = db_test_utils.InitializeClient(db)
flow_id = flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
self.FakeExecuteCommand(),
client_id=client_id,
args=flow_args,
creator=context.username)
flow_test_lib.FinishAllFlowsOnClient(client_id)
class FakeParser(
abstract_parser.SingleResponseParser[
rdf_client_action.ExecuteResponse], ):
supported_artifacts = [self.ECHO1337_ARTIFACT.name]
def ParseResponse(
self,
knowledge_base: rdf_client.KnowledgeBase,
response: rdf_client_action.ExecuteResponse,
) -> Iterable[rdf_client_action.ExecuteResponse]:
precondition.AssertType(response,
rdf_client_action.ExecuteResponse)
parsed_response = rdf_client_action.ExecuteResponse()
parsed_response.stdout = response.stdout
parsed_response.stderr = b"4815162342"
return [parsed_response]
with parser_test_lib._ParserContext("Fake", FakeParser):
args = flow_plugin.ApiListParsedFlowResultsArgs(
client_id=client_id, flow_id=flow_id, offset=0, count=1024)
result = self.handler.Handle(args, context=context)
self.assertEmpty(result.errors)
self.assertLen(result.items, 1)
response = result.items[0].payload
self.assertIsInstance(response, rdf_client_action.ExecuteResponse)
self.assertEqual(response.stdout, b"1337")
self.assertEqual(response.stderr, b"4815162342")
def _RunUpdateFlow(self, client_id):
gui_test_lib.CreateFileVersion(
client_id,
"fs/os/c/a.txt",
"Hello World".encode("utf-8"),
timestamp=gui_test_lib.TIME_0)
gui_test_lib.CreateFolder(
client_id, "fs/os/c/TestFolder", timestamp=gui_test_lib.TIME_0)
gui_test_lib.CreateFolder(
client_id, "fs/os/c/bin/TestBinFolder", timestamp=gui_test_lib.TIME_0)
flow_test_lib.FinishAllFlowsOnClient(client_id)
def _Collect(self, root: bytes) -> Iterator[timeline_pb2.TimelineEntry]:
args = rdf_timeline.TimelineArgs(root=root)
flow_id = flow_test_lib.TestFlowHelper(
timeline_flow.TimelineFlow.__name__,
action_mocks.ActionMock(timeline_action.Timeline),
client_id=self.client_id,
creator=self.test_username,
args=args)
flow_test_lib.FinishAllFlowsOnClient(self.client_id)
return timeline_flow.ProtoEntries(client_id=self.client_id, flow_id=flow_id)
def _Collect(self, root):
args = rdf_timeline.TimelineArgs(root=root)
flow_id = flow_test_lib.TestFlowHelper(
timeline_flow.TimelineFlow.__name__,
action_mocks.ActionMock(timeline_action.Timeline),
client_id=self.client_id,
token=self.token,
args=args)
flow_test_lib.FinishAllFlowsOnClient(self.client_id)
return timeline_flow.Entries(client_id=self.client_id, flow_id=flow_id)
def _YaraProcessScan(
self,
args: rdf_memory.YaraProcessScanRequest,
action_mock: Optional[action_mocks.ActionMock] = None,
) -> None:
if action_mock is None:
action_mock = action_mocks.ActionMock()
flow_test_lib.TestFlowHelper(memory.YaraProcessScan.__name__,
action_mock,
client_id=self.client_id,
creator=self.test_username,
args=args)
flow_test_lib.FinishAllFlowsOnClient(self.client_id)
def _YaraProcessScan(
self,
args,
action_mock=None,
):
if action_mock is None:
action_mock = action_mocks.ActionMock()
flow_test_lib.TestFlowHelper(memory.YaraProcessScan.__name__,
action_mock,
client_id=self.client_id,
token=self.token,
args=args)
flow_test_lib.FinishAllFlowsOnClient(self.client_id)
def _RunUpdateFlow(self, client_id):
gui_test_lib.CreateFileVersion(
client_id,
"fs/os/c/a.txt",
"Hello World".encode("utf-8"),
timestamp=gui_test_lib.TIME_0)
gui_test_lib.CreateFolder(
client_id, "fs/os/c/TestFolder", timestamp=gui_test_lib.TIME_0)
gui_test_lib.CreateFolder(
client_id, "fs/os/c/bin/TestBinFolder", timestamp=gui_test_lib.TIME_0)
flow_test_lib.FinishAllFlowsOnClient(
client_id,
client_mock=action_mocks.MultiGetFileClientMock(),
check_flow_errors=False)
def testClientFileFinderIndicatesCollectedSizeAfterCollection(self):
client_ref = self.api.Client(client_id=self.client_id)
args = rdf_file_finder.FileFinderArgs(
paths=[os.path.join(self.base_path, "numbers.txt")],
action=rdf_file_finder.FileFinderAction.Download(),
follow_links=True).AsPrimitiveProto()
client_ref.CreateFlow(name=file_finder.ClientFileFinder.__name__, args=args)
flow_test_lib.FinishAllFlowsOnClient(
self.client_id, client_mock=action_mocks.ClientFileFinderClientMock())
f = client_ref.File("fs/os" +
os.path.join(self.base_path, "numbers.txt")).Get()
self.assertNotEqual(f.data.hash.sha256, b"")
self.assertGreater(f.data.hash.num_bytes, 0)
self.assertGreater(f.data.last_collected, 0)
self.assertGreater(f.data.last_collected_size, 0)
def testReportsArtifactCollectionErrors(self, db: abstract_db.Database):
context = _CreateContext(db)
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(self.ECHO1337_ARTIFACT)
flow_args = rdf_artifacts.ArtifactCollectorFlowArgs()
flow_args.artifact_list = [self.ECHO1337_ARTIFACT.name]
flow_args.apply_parsers = False
client_id = db_test_utils.InitializeClient(db)
flow_id = flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
self.FakeExecuteCommand(),
client_id=client_id,
args=flow_args,
creator=context.username)
flow_test_lib.FinishAllFlowsOnClient(client_id)
class FakeParser(
abstract_parser.SingleResponseParser[
rdf_client_action.ExecuteResponse], ):
supported_artifacts = [self.ECHO1337_ARTIFACT.name]
def ParseResponse(
self, knowledge_base: rdf_client.KnowledgeBase,
response: rdf_client_action.ExecuteResponse
) -> Iterable[rdf_client_action.ExecuteResponse]:
del knowledge_base, response # Unused.
raise abstract_parser.ParseError("Lorem ipsum.")
with parser_test_lib._ParserContext("Fake", FakeParser):
args = flow_plugin.ApiListParsedFlowResultsArgs(
client_id=client_id, flow_id=flow_id, offset=0, count=1024)
result = self.handler.Handle(args, context=context)
self.assertEmpty(result.items)
self.assertLen(result.errors, 1)
self.assertEqual(result.errors[0], "Lorem ipsum.")
def testFileFinderIndicatesCollectedSizeAfterCollection(self):
client_ref = self.api.Client(client_id=self.client_id)
# TODO(user): for symlink-related test scenarions, this should require
# follow_links to be True. However, unlike the ClientFileFinder test
# below, this one doesn't care about this setting. Fix the
# FileFinder/ClientFileFinder behavior to match each other.
args = rdf_file_finder.FileFinderArgs(
paths=[os.path.join(self.base_path, "numbers.txt")],
action=rdf_file_finder.FileFinderAction.Download()).AsPrimitiveProto()
client_ref.CreateFlow(name=file_finder.FileFinder.__name__, args=args)
flow_test_lib.FinishAllFlowsOnClient(
self.client_id, client_mock=action_mocks.FileFinderClientMock())
f = client_ref.File("fs/os" +
os.path.join(self.base_path, "numbers.txt")).Get()
self.assertNotEqual(f.data.hash.sha256, b"")
self.assertGreater(f.data.hash.num_bytes, 0)
self.assertGreater(f.data.last_collected, 0)
self.assertGreater(f.data.last_collected_size, 0)
def ProcessFlow():
time.sleep(1)
client_mock = action_mocks.ListProcessesMock([])
flow_test_lib.FinishAllFlowsOnClient(client_urn,
client_mock=client_mock)
def ProcessOperation():
time.sleep(1)
flow_test_lib.FinishAllFlowsOnClient(self.client_id)
def testRefreshFileStartsFlow(self):
self.Open("/#/clients/C.0000000000000001/vfs/fs/os/c/Downloads/")
# Select a file and start a flow by requesting a newer version.
self.Click("css=tr:contains(\"a.txt\")")
self.Click("css=li[heading=Download]")
self.Click("css=button:contains(\"Collect from the client\")")
# Create a new file version (that would have been created by the flow
# otherwise) and finish the flow.
# Make sure that the flow has started (when button is clicked, the HTTP
# API request is sent asynchronously).
def MultiGetFileStarted():
return compatibility.GetName(transfer.MultiGetFile) in [
f.flow_class_name
for f in data_store.REL_DB.ReadAllFlowObjects(
client_id=self.client_id)
]
self.WaitUntil(MultiGetFileStarted)
flow_test_lib.FinishAllFlowsOnClient(self.client_id,
check_flow_errors=False)
time_in_future = rdfvalue.RDFDatetime.Now() + rdfvalue.DurationSeconds(
"1h")
# We have to make sure that the new version will not be within a second
# from the current one, otherwise the previous one and the new one will
# be indistinguishable in the UI (as it has a 1s precision when
# displaying versions).
gui_test_lib.CreateFileVersion(self.client_id,
"fs/os/c/Downloads/a.txt",
"The newest version!".encode("utf-8"),
timestamp=time_in_future)
# Once the flow has finished, the file view should update and add the
# newly created, latest version of the file to the list. The selected
# option should still be "HEAD".
self.WaitUntilContains("HEAD", self.GetText,
"css=.version-dropdown > option[selected]")
# The file table should also update and display the new timestamp.
self.WaitUntilContains(gui_test_lib.DateTimeString(time_in_future),
self.GetText,
"css=.version-dropdown > option:nth(1)")
# The file table should also update and display the new timestamp.
self.WaitUntil(
self.IsElementPresent,
"css=grr-file-table tbody > tr td:contains(\"%s\")" %
(gui_test_lib.DateTimeString(time_in_future)))
# Make sure the file content has changed.
self.Click("css=li[heading=TextView]")
self.WaitUntilContains("The newest version!", self.GetText,
"css=div.monospace pre")
# Go to the flow management screen and check that there was a new flow.
self.Click("css=a:contains('Manage launched flows')")
self.Click("css=grr-flows-list tr:contains('MultiGetFile')")
self.WaitUntilContains(transfer.MultiGetFile.__name__, self.GetText,
"css=#main_bottomPane")
self.WaitUntilContains(
"c/Downloads/a.txt", self.GetText,
"css=#main_bottomPane table > tbody td.proto_key:contains(\"Path\") "
"~ td.proto_value")
def testUsesCollectionTimeFiles(self, db: abstract_db.Database):
token = _CreateToken(db)
client_id = db_test_utils.InitializeClient(db)
snapshot = rdf_objects.ClientSnapshot()
snapshot.client_id = client_id
snapshot.knowledge_base.os = "redox"
db.WriteClientSnapshot(snapshot)
with temp.AutoTempFilePath() as temp_filepath:
fake_artifact_source = rdf_artifacts.ArtifactSource(
type=rdf_artifacts.ArtifactSource.SourceType.FILE,
attributes={
"paths": [temp_filepath],
})
fake_artifact = rdf_artifacts.Artifact(
name="FakeArtifact",
doc="Lorem ipsum.",
sources=[fake_artifact_source])
flow_args = rdf_artifacts.ArtifactCollectorFlowArgs()
flow_args.artifact_list = [fake_artifact.name]
flow_args.apply_parsers = False
with io.open(temp_filepath, mode="wb") as temp_filedesc:
temp_filedesc.write(b"OLD")
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(fake_artifact)
# First, we run the artifact collector to collect the old file and save
# the flow id to parse the results later.
flow_id = flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
action_mocks.FileFinderClientMock(),
client_id=client_id,
args=flow_args,
token=token)
flow_test_lib.FinishAllFlowsOnClient(client_id)
with io.open(temp_filepath, mode="wb") as temp_filedesc:
temp_filedesc.write(b"NEW")
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(fake_artifact)
# Now, we run the artifact collector again to collect the new file to
# update to this version on the server. The parsing should be performed
# against the previous flow.
flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
action_mocks.FileFinderClientMock(),
client_id=client_id,
args=flow_args,
token=token)
flow_test_lib.FinishAllFlowsOnClient(client_id)
class FakeFileParser(abstract_parser.SingleFileParser):
supported_artifacts = [fake_artifact.name]
def ParseFile(
self,
knowledge_base: rdf_client.KnowledgeBase,
pathspec: rdf_paths.PathSpec,
filedesc: file_store.BlobStream,
) -> Iterable[rdfvalue.RDFBytes]:
del knowledge_base, pathspec # Unused.
return [rdfvalue.RDFBytes(filedesc.Read())]
with parser_test_lib._ParserContext("FakeFile", FakeFileParser):
args = flow_plugin.ApiListParsedFlowResultsArgs(
client_id=client_id, flow_id=flow_id, offset=0, count=1024)
result = self.handler.Handle(args, token=token)
self.assertEmpty(result.errors)
self.assertLen(result.items, 1)
response = result.items[0].payload
self.assertEqual(response, b"OLD")
