def _GetSignedBlobsRoots():
return {
ApiGrrBinary.Type.PYTHON_HACK:
signed_binary_utils.GetAFF4PythonHackRoot(),
ApiGrrBinary.Type.EXECUTABLE:
signed_binary_utils.GetAFF4ExecutablesRoot()
}
def _GetBinaryRootUrn(binary_type):
if binary_type == api_config.ApiGrrBinary.Type.PYTHON_HACK:
return signed_binary_utils.GetAFF4PythonHackRoot()
elif binary_type == api_config.ApiGrrBinary.Type.EXECUTABLE:
return signed_binary_utils.GetAFF4ExecutablesRoot()
else:
raise ValueError("Invalid binary type: %s" % binary_type)
def testScheduleLaunchExecutePythonHackFlow(self):
self._SetUpAdminUser()
maintenance_utils.UploadSignedConfigBlob(
b'foo',
aff4_path=signed_binary_utils.GetAFF4PythonHackRoot().Add(
'windows/a.py'))
maintenance_utils.UploadSignedConfigBlob(
b'foo',
aff4_path=signed_binary_utils.GetAFF4PythonHackRoot().Add(
'windows/test.py'))
self.Open(f'/v2/clients/{self.client_id}')
self.WaitUntilContains('No access', self.GetText,
'css=client-overview')
self.Type('css=flow-form input[name=flowSearchBox]',
'python',
end_with_enter=True)
self.Type('css=flow-args-form input[name=hackName]',
'test',
end_with_enter=True)
self.Click('css=flow-args-form button:contains("Add argument")')
self.Type('css=flow-args-form .key-input input', 'fookey')
self.Type('css=flow-args-form .value-input input', 'foovalue')
self.Click('css=flow-form button:contains("Schedule")')
def GetFirstScheduledFlow():
scheduled_flows = _ListScheduledFlows(self.client_id,
self.test_username)
return scheduled_flows[0] if len(scheduled_flows) == 1 else None
scheduled_flow = self.WaitUntil(GetFirstScheduledFlow)
self.assertEqual(scheduled_flow.flow_name,
administrative.ExecutePythonHack.__name__)
self.assertEqual(scheduled_flow.flow_args.hack_name, 'windows/test.py')
self.assertEqual(scheduled_flow.flow_args.py_args['fookey'],
'foovalue')
def SetUpBinaries(self):
with test_lib.FakeTime(42):
code = "I am a binary file"
upload_path = signed_binary_utils.GetAFF4ExecutablesRoot().Add(
"windows/test.exe")
maintenance_utils.UploadSignedConfigBlob(
code.encode("utf-8"), aff4_path=upload_path)
with test_lib.FakeTime(43):
code = "I'm a python hack"
upload_path = signed_binary_utils.GetAFF4PythonHackRoot().Add("test")
maintenance_utils.UploadSignedConfigBlob(
code.encode("utf-8"), aff4_path=upload_path)
def Start(self):
"""The start method."""
python_hack_urn = signed_binary_utils.GetAFF4PythonHackRoot().Add(
self.args.hack_name)
try:
blob_iterator, _ = signed_binary_utils.FetchBlobsForSignedBinaryByURN(
python_hack_urn)
except signed_binary_utils.SignedBinaryNotFoundError:
raise flow_base.FlowError("Python hack %s not found." %
self.args.hack_name)
# TODO(amoser): This will break if someone wants to execute lots of Python.
for python_blob in blob_iterator:
self.CallClient(server_stubs.ExecutePython,
python_code=python_blob,
py_args=self.args.py_args,
next_state=compatibility.GetName(self.Done))
def UploadSignedBinary(source_path,
binary_type,
platform,
upload_subdirectory="",
token=None):
"""Signs a binary and uploads it to the datastore.
Args:
source_path: Path to the binary to upload.
binary_type: Type of the binary, e.g python-hack or executable.
platform: Client platform where the binary is intended to be run.
upload_subdirectory: Path of a subdirectory to upload the binary to,
relative to the canonical path for binaries of the given type and
platform.
token: ACL token to use for uploading.
Raises:
BinaryTooLargeError: If the binary to upload is too large.
"""
if binary_type == rdf_objects.SignedBinaryID.BinaryType.PYTHON_HACK:
root_urn = signed_binary_utils.GetAFF4PythonHackRoot()
elif binary_type == rdf_objects.SignedBinaryID.BinaryType.EXECUTABLE:
root_urn = signed_binary_utils.GetAFF4ExecutablesRoot()
else:
raise ValueError("Unknown binary type %s." % binary_type)
file_size = os.path.getsize(source_path)
if file_size > _MAX_SIGNED_BINARY_BYTES:
raise BinaryTooLargeError(
"File [%s] is of size %d (bytes), which exceeds the allowed maximum "
"of %d bytes." %
(source_path, file_size, _MAX_SIGNED_BINARY_BYTES))
binary_urn = root_urn.Add(platform.lower()).Add(upload_subdirectory).Add(
os.path.basename(source_path))
context = ["Platform:%s" % platform.title(), "Client Context"]
with open(source_path, "rb") as f:
file_content = f.read()
maintenance_utils.UploadSignedConfigBlob(file_content,
aff4_path=binary_urn,
client_context=context,
token=token)
print("Uploaded to %s" % binary_urn)
