def testSerialization(self):
test_dict = dict(
key1=1, # Integer.
key2="foo", # String.
key3=u"\u4f60\u597d", # Unicode.
key5=rdfvalue.RDFDatetime("2012/12/11"), # RDFValue.
key6=None, # Support None Encoding.
key7=structs.Enum(5, name="Test"), # Integer like objects.
)
# Initialize through keywords.
sample = rdfvalue.Dict(**test_dict)
self.CheckTestDict(test_dict, sample)
# Initialize through dict.
sample = rdfvalue.Dict(test_dict)
self.CheckTestDict(test_dict, sample)
# Initialize through a serialized form.
serialized = sample.SerializeToString()
self.assertIsInstance(serialized, str)
sample = rdfvalue.Dict(serialized)
self.CheckTestDict(test_dict, sample)
# Convert to a dict.
self.CheckTestDict(test_dict, sample.ToDict())
def testNestedDictsOpaqueTypes(self):
class UnSerializable(object):
pass
test_dict = dict(key1={"A": 1},
key2=rdfvalue.Dict({"A": 1}),
key3=[1, UnSerializable(), 3, [1, 2, [3]]],
key4=[[], None, ["abc"]],
key5=UnSerializable(),
key6=["a", UnSerializable(), "b"])
self.assertRaises(TypeError, rdfvalue.Dict, **test_dict)
sample = rdfvalue.Dict()
for key, value in test_dict.iteritems():
sample.SetItem(key, value, raise_on_error=False)
# Need to do some manual checking here since this is a lossy conversion.
self.assertEqual(test_dict["key1"], sample["key1"])
self.assertEqual(test_dict["key2"], sample["key2"])
self.assertEqual(1, sample["key3"][0])
self.assertTrue("Unsupported type" in sample["key3"][1])
self.assertItemsEqual(test_dict["key3"][2:], sample["key3"][2:])
self.assertEqual(test_dict["key4"], sample["key4"])
self.assertTrue("Unsupported type" in sample["key5"])
self.assertEqual("a", sample["key6"][0])
self.assertTrue("Unsupported type" in sample["key6"][1])
self.assertEqual("b", sample["key6"][2])
def testNestedDicts(self):
test_dict = dict(
key1={"A": 1},
key2=rdfvalue.Dict({"A": 1}),
)
sample = rdfvalue.Dict(**test_dict)
self.CheckTestDict(test_dict, sample)
self.CheckTestDict(test_dict, sample.ToDict())
def testNestedDictsMultipleTypes(self):
test_dict = dict(key1={"A": 1},
key2=rdfvalue.Dict({"A": 1}),
key3=[1, 2, 3, [1, 2, [3]]],
key4=[[], None, ["abc"]])
sample = rdfvalue.Dict(**test_dict)
self.CheckTestDict(test_dict, sample)
self.CheckTestDict(test_dict, sample.ToDict())
def testOverwriting(self):
req = rdfvalue.Iterator(client_state=rdfvalue.Dict({"A": 1}))
# There should be one element now.
self.assertEqual(len(list(req.client_state.items())), 1)
req.client_state = rdfvalue.Dict({"B": 2})
# Still one element.
self.assertEqual(len(list(req.client_state.items())), 1)
req.client_state = rdfvalue.Dict({})
# And now it's gone.
self.assertEqual(len(list(req.client_state.items())), 0)
def testUpdateConfig(self):
"""Ensure we can retrieve and update the config."""
# Only mock the pieces we care about.
client_mock = action_mocks.ActionMock("GetConfiguration",
"UpdateConfiguration")
loc = "http://www.example.com"
new_config = rdfvalue.Dict({
"Client.control_urls": [loc],
"Client.foreman_check_frequency": 3600,
"Client.poll_min": 1
})
# Write the config.
for _ in test_lib.TestFlowHelper("UpdateConfiguration",
client_mock,
client_id=self.client_id,
token=self.token,
config=new_config):
pass
# Now retrieve it again to see if it got written.
for _ in test_lib.TestFlowHelper("Interrogate",
client_mock,
token=self.token,
client_id=self.client_id):
pass
fd = aff4.FACTORY.Open(self.client_id, token=self.token)
config_dat = fd.Get(fd.Schema.GRR_CONFIGURATION)
self.assertEqual(config_dat["Client.control_urls"], [loc])
self.assertEqual(config_dat["Client.poll_min"], 1)
def __init__(self, initializer=None, age=None, **kwargs):
if isinstance(initializer, dict):
conf = initializer
initializer = None
else:
conf = kwargs
super(Method, self).__init__(initializer=initializer, age=age)
probe = conf.get("probe", {})
resource = conf.get("resource", {})
hint = conf.get("hint", {})
target = conf.get("target", {})
if hint:
# Add the hint to children.
for cfg in probe:
cfg["hint"] = hints.Overlay(child=cfg.get("hint", {}),
parent=hint)
self.probe = [Probe(**cfg) for cfg in probe]
self.hint = Hint(hint, reformat=False)
self.match = MatchStrToList(kwargs.get("match"))
self.matcher = checks.Matcher(self.match, self.hint)
self.resource = [rdfvalue.Dict(**r) for r in resource]
self.target = Target(**target)
self.triggers = triggers.Triggers()
for p in self.probe:
# If the probe has a target, use it. Otherwise, use the method's target.
target = p.target or self.target
self.triggers.Add(p.artifact, target, p)
def ParseMultiple(self, stats, file_objects, knowledge_base):
"""Parse the found release files."""
_ = knowledge_base
# Collate files into path: contents dictionary.
found_files = self._Combine(stats, file_objects)
# Determine collected files and apply weighting.
weights = [w for w in self.WEIGHTS if w.path in found_files]
weights = sorted(weights, key=lambda x: x.weight)
for _, path, handler in weights:
contents = found_files[path]
obj = handler(contents)
complete, result = obj.Parse()
if result is None:
continue
elif complete:
yield rdfvalue.Dict({
'os_release': result.release,
'os_major_version': result.major,
'os_minor_version': result.minor
})
break
else:
# No successful parse.
yield rdfvalue.Anomaly(type='PARSER_ANOMALY',
symptom='Unable to determine distribution.')
def Parse(self, stat, knowledge_base):
_ = stat, knowledge_base
test_dict = {
"environ_temp": rdfvalue.RDFString("tempvalue"),
"environ_path": rdfvalue.RDFString("pathvalue")
}
yield rdfvalue.Dict(test_dict)
def testVolatilityModules(self):
"""Tests the end to end volatility memory analysis."""
image_path = os.path.join(self.base_path, "win7_trial_64bit.raw")
if not os.access(image_path, os.R_OK):
logging.warning(
"Unable to locate test memory image. Skipping test.")
return
self.CreateClient()
self.CreateSignedDriver()
class ClientMock(test_lib.MemoryClientMock):
"""A mock which returns the image as the driver path."""
def GetMemoryInformation(self, _):
"""Mock out the driver loading code to pass the memory image."""
reply = rdfvalue.MemoryInformation(device=rdfvalue.PathSpec(
path=image_path, pathtype=rdfvalue.PathSpec.PathType.OS))
reply.runs.Append(offset=0, length=1000000000)
return [reply]
request = rdfvalue.VolatilityRequest()
request.args["pslist"] = {}
request.args["modules"] = {}
# To speed up the test we provide these values. In real life these values
# will be provided by the kernel driver.
request.session = rdfvalue.Dict(dtb=0x187000, kdbg=0xF80002803070)
# Allow the real VolatilityAction to run against the image.
for _ in test_lib.TestFlowHelper("AnalyzeClientMemory",
ClientMock("VolatilityAction"),
token=self.token,
client_id=self.client_id,
request=request,
output="analysis/memory"):
pass
fd = aff4.FACTORY.Open(self.client_id.Add("analysis/memory/pslist"),
token=self.token)
result = fd.Get(fd.Schema.RESULT)
# Pslist should have 32 rows.
self.assertEqual(len(result.sections[0].table.rows), 32)
# And should include the DumpIt binary.
self.assert_("DumpIt.exe" in str(result))
fd = aff4.FACTORY.Open(self.client_id.Add("analysis/memory/modules"),
token=self.token)
result = fd.Get(fd.Schema.RESULT)
# Modules should have 133 lines.
self.assertEqual(len(result.sections[0].table.rows), 133)
# And should include the DumpIt kernel driver.
self.assert_("DumpIt.sys" in str(result))
def GetConfiguration(self, _):
return [
rdfvalue.Dict({
"Client.control_urls": ["http://localhost:8001/control"],
"Client.poll_min":
1.0
})
]
def Run(self, args):
"""Run the WMI query and return the data."""
query = args.query
base_object = args.base_object or r"winmgmts:\root\cimv2"
if not query.upper().startswith("SELECT "):
raise RuntimeError("Only SELECT WMI queries allowed.")
# Now return the data to the server
for response_dict in RunWMIQuery(query, baseobj=base_object):
self.SendReply(rdfvalue.Dict(response_dict))
def testDictBehaviour(self):
tested = rdfvalue.Dict(a=1)
now = rdfvalue.RDFDatetime().Now()
tested["b"] = now
self.assertEqual(tested["b"], now)
self.assertEqual(tested["a"], 1)
tested["b"] = rdfvalue.RDFURN("aff4:/users/")
self.assertEqual(len(tested), 2)
self.assertEqual(tested["b"].SerializeToString(), "aff4:/users")
def testGetConfig(self):
"""Check GetConfig client action works."""
# Use UpdateConfig to generate a config.
location = ["http://example.com"]
request = rdfvalue.Dict()
request["Client.control_urls"] = location
request["Client.foreman_check_frequency"] = 3600
self.RunAction("UpdateConfiguration", request)
# Check that our GetConfig actually gets the real data.
self.RunAction("GetConfiguration")
self.assertEqual(config_lib.CONFIG["Client.foreman_check_frequency"], 3600)
self.assertEqual(config_lib.CONFIG["Client.control_urls"], location)
def Run(self, request):
"""Munge the iterator to the server and abstract it away."""
# Pass the client_state as a dict to the action. This is often more
# efficient than manipulating a protobuf.
client_state = request.iterator.client_state.ToDict()
# Derived classes should implement this.
self.Iterate(request, client_state)
# Update the iterator client_state from the dict.
request.iterator.client_state = rdfvalue.Dict(client_state)
# Return the iterator
self.SendReply(request.iterator,
message_type=rdfvalue.GrrMessage.Type.ITERATOR)
def testArgs(self):
"""Test passing arguments."""
utils.TEST_VAL = "original"
python_code = """
magic_return_str = py_args['test']
utils.TEST_VAL = py_args[43]
"""
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
pdict = rdfvalue.Dict({"test": "dict_arg", 43: "dict_arg2"})
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob,
py_args=pdict)
result = self.RunAction("ExecutePython", request)[0]
self.assertEqual(result.return_val, "dict_arg")
self.assertEqual(utils.TEST_VAL, "dict_arg2")
def testInterfaceParsing(self):
parser = wmi_parser.WMIInterfacesParser()
rdf_dict = rdfvalue.Dict()
wmi_properties = (
client_fixture.WMIWin32NetworkAdapterConfigurationMock.__dict__.
iteritems())
for key, value in wmi_properties:
if not key.startswith("__"):
try:
rdf_dict[key] = value
except TypeError:
rdf_dict[key] = "Failed to encode: %s" % value
result_list = list(parser.Parse(None, rdf_dict, None))
self.assertEqual(len(result_list), 2)
for result in result_list:
if isinstance(result, rdfvalue.Interface):
self.assertEqual(len(result.addresses), 4)
self.assertItemsEqual(
[x.human_readable_address for x in result.addresses], [
"192.168.1.20", "ffff::ffff:aaaa:1111:aaaa",
"dddd:0:8888:6666:bbbb:aaaa:eeee:bbbb",
"dddd:0:8888:6666:bbbb:aaaa:ffff:bbbb"
])
self.assertItemsEqual([
x.human_readable_address for x in result.dhcp_server_list
], ["192.168.1.1"])
self.assertEqual(
result.dhcp_lease_expires.AsMicroSecondsFromEpoch(),
1409008979123456)
self.assertEqual(
result.dhcp_lease_obtained.AsMicroSecondsFromEpoch(),
1408994579123456)
elif isinstance(result, rdfvalue.DNSClientConfiguration):
self.assertItemsEqual(
result.dns_server,
["192.168.1.1", "192.168.255.81", "192.168.128.88"])
self.assertItemsEqual(result.dns_suffix, [
"blah.example.com", "ad.example.com",
"internal.example.com", "example.com"
])
def testUpdateConfigBlacklist(self):
"""Tests that disallowed fields are not getting updated."""
config_lib.CONFIG.Set("Client.control_urls", ["http://something.com/"])
config_lib.CONFIG.Set("Client.server_serial_number", 1)
location = ["http://www.example.com"]
request = rdfvalue.Dict()
request["Client.control_urls"] = location
request["Client.server_serial_number"] = 10
self.RunAction("UpdateConfiguration", request)
# Location can be set.
self.assertEqual(config_lib.CONFIG["Client.control_urls"], location)
# But the server serial number can not be updated.
self.assertEqual(config_lib.CONFIG["Client.server_serial_number"], 1)
def ProcessFileStats(self, responses):
"""Extract DataBlob from Stat response."""
if not responses.success:
return
system_root_paths = ["Windows", "WinNT", "WINNT35", "WTSRV"]
for response in responses:
if response.pathspec.path[4:] in system_root_paths:
systemdrive = response.pathspec.path[1:3]
systemroot = "%s\\%s" % (systemdrive,
response.pathspec.path[4:])
self.SendReply(
rdfvalue.Dict({
"environ_systemroot": systemroot,
"environ_systemdrive": systemdrive
}))
self.state.bootstrap_initialized = True
break
def testRdfFormatterFanOut(self):
rdf = rdfvalue.Dict()
user1 = rdfvalue.KnowledgeBaseUser(username="******")
user2 = rdfvalue.KnowledgeBaseUser(username="******")
rdf["cataclysm"] = "GreyGoo"
rdf["thinkers"] = [user1, user2]
rdf["reference"] = {
"ecophage": ["bots", ["nanobots", ["picobots"]]],
"doomsday": {
"books": ["cats cradle", "prey"]
}
}
template = ("{cataclysm}; {thinkers.username}; {reference.ecophage}; "
"{reference.doomsday}")
hinter = hints.Hinter(template=template)
expected = ("GreyGoo; drexler,joy; bots,nanobots,picobots; "
"books:cats cradle,prey")
result = hinter.Render(rdf)
self.assertEqual(expected, result)
def testUpdateConfiguration(self):
"""Test that we can update the config."""
# A unique name on the filesystem for the writeback.
self.config_file = os.path.join(self.temp_dir, "ConfigActionTest.yaml")
# In a real client, the writeback location should be set to something real,
# but for this test we make it the same as the config file..
config_lib.CONFIG.SetWriteBack(self.config_file)
# Make sure the file is gone
self.assertRaises(IOError, open, self.config_file)
location = ["http://www.example1.com/", "http://www.example2.com/"]
request = rdfvalue.Dict()
request["Client.control_urls"] = location
request["Client.foreman_check_frequency"] = 3600
result = self.RunAction("UpdateConfiguration", request)
self.assertEqual(result, [])
self.assertEqual(config_lib.CONFIG["Client.foreman_check_frequency"],
3600)
# Test the config file got written.
data = open(self.config_file).read()
self.assertTrue("control_urls: {0}".format(",".join(location)) in data)
self.urls = []
# Now test that our location was actually updated.
def FakeUrlOpen(req, timeout=10):
_ = timeout
self.urls.append(req.get_full_url())
return StringIO.StringIO()
comms.urllib2.urlopen = FakeUrlOpen
client_context = comms.GRRHTTPClient(worker=MockClientWorker)
client_context.MakeRequest("", comms.Status())
self.assertTrue(location[0] in self.urls[0])
self.assertTrue(location[1] in self.urls[1])
def WmiQuery(self, query):
if query.query == u"SELECT * FROM Win32_LogicalDisk":
self.response_count += 1
return client_fixture.WMI_SAMPLE
elif query.query.startswith("Select * "
"from Win32_NetworkAdapterConfiguration"):
self.response_count += 1
rdf_dict = rdfvalue.Dict()
wmi_properties = (
client_fixture.WMIWin32NetworkAdapterConfigurationMock.
__dict__.iteritems())
for key, value in wmi_properties:
if not key.startswith("__"):
try:
rdf_dict[key] = value
except TypeError:
rdf_dict[key] = "Failed to encode: %s" % value
return [rdf_dict]
else:
return None
def RunWMIQuery(query, baseobj=r"winmgmts:\root\cimv2"):
"""Run a WMI query and return a result.
Args:
query: the WMI query to run.
baseobj: the base object for the WMI query.
Yields:
rdfvalue.Dicts containing key value pairs from the resulting COM objects.
"""
pythoncom.CoInitialize() # Needs to be called if using com from a thread.
wmi_obj = win32com.client.GetObject(baseobj)
# This allows our WMI to do some extra things, in particular
# it gives it access to find the executable path for all processes.
wmi_obj.Security_.Privileges.AddAsString("SeDebugPrivilege")
# Run query
try:
query_results = wmi_obj.ExecQuery(query)
except pythoncom.com_error as e:
raise RuntimeError("Failed to run WMI query \'%s\' err was %s" %
(query, e))
# Extract results from the returned COMObject and return dicts.
try:
for result in query_results:
response = rdfvalue.Dict()
for prop in result.Properties_:
if prop.Name not in IGNORE_PROPS:
# Protodict can handle most of the types we care about, but we may
# get some objects that we don't know how to serialize, so we tell the
# dict to set the value to an error message and keep going
response.SetItem(prop.Name,
prop.Value,
raise_on_error=False)
yield response
except pythoncom.com_error as e:
raise RuntimeError("WMI query data error on query \'%s\' err was %s" %
(e, query))
def SetupRequest(self, plugin):
# Only run this test if the image file is found.
image_path = os.path.join(self.base_path, "win7_trial_64bit.raw")
if not os.access(image_path, os.R_OK):
logging.warning(
"Unable to locate test memory image. Skipping test.")
return False
self.request = rdfvalue.VolatilityRequest(
device=rdfvalue.PathSpec(path=image_path,
pathtype=rdfvalue.PathSpec.PathType.OS),
# To speed up the test we provide these values. In real life these
# values will be provided by the kernel driver.
session=rdfvalue.Dict(dtb=0x187000, kdbg=0xF80002803070))
# In this test we explicitly do not set the profile to use so we can see if
# the profile autodetection works.
# Add the plugin to the request.
self.request.args[plugin] = None
return True
def testParameters(self):
if not self.SetupRequest("pslist"):
return
args = {"pslist": {"pid": 2860}}
self.request.args = rdfvalue.Dict(args)
result = self.RunAction("VolatilityAction", self.request)
# There should be 1 result back.
self.assertEqual(len(result), 1)
# There should be one section.
self.assertEqual(len(result[0].sections), 1)
rows = result[0].sections[0].table.rows
# Pslist should now have 1 result.
self.assertEqual(len(rows), 1)
name = rows[0].values[1].GetValue()
self.assertTrue("DumpIt.exe" in name)
def ProcessRegStat(self, responses):
"""Check SystemRoot registry value."""
if responses.success:
systemroot = responses.First().registry_data.GetValue()
if systemroot:
systemdrive = systemroot[0:2]
if re.match(r"^[A-Za-z]:$", systemdrive):
self.SendReply(
rdfvalue.Dict({
"environ_systemroot": systemroot,
"environ_systemdrive": systemdrive
}))
self.state.bootstrap_initialized = True
return
# If registry querying didn't work, we try to guess common paths instead.
system_drive_opts = ["C:", "D:"]
for drive in system_drive_opts:
pathspec = rdfvalue.PathSpec(
path=drive, pathtype=rdfvalue.PathSpec.PathType.OS)
self.CallClient("ListDirectory",
pathspec=pathspec,
request_data={"bootstrap_var": "system_root"},
next_state="ProcessFileStats")
from grr.lib import config_lib
from grr.lib import flags
from grr.lib import flow
from grr.lib import parsers
from grr.lib import rdfvalue
from grr.lib import test_lib
from grr.lib import utils
# pylint: mode=test
WMI_SAMPLE = [
rdfvalue.Dict({
u"Version": u"65.61.49216",
u"InstallDate2": u"",
u"Name": u"Google Chrome",
u"Vendor": u"Google, Inc.",
u"Description": u"Google Chrome",
u"IdentifyingNumber": u"{35790B21-ACFE-33F5-B320-9DA320D96682}",
u"InstallDate": u"20130710"
}),
rdfvalue.Dict({
u"Version": u"7.0.1",
u"InstallDate2": u"",
u"Name": u"Parity Agent",
u"Vendor": u"Bit9, Inc.",
u"Description": u"Parity Agent",
u"IdentifyingNumber": u"{ADC7EB41-4CC2-4FBA-8FBE-9338A9FB7666}",
u"InstallDate": u"20130710"
}),
rdfvalue.Dict({
u"Version": u"8.0.61000",
def CallState(self,
messages=None,
next_state="",
client_id=None,
request_data=None,
start_time=None):
"""This method is used to asynchronously schedule a new hunt state.
The state will be invoked in a later time and receive all the messages
we send.
Args:
messages: A list of rdfvalues to send. If the last one is not a
GrrStatus, we append an OK Status.
next_state: The state in this hunt to be invoked with the responses.
client_id: ClientURN to use in scheduled requests.
request_data: Any dict provided here will be available in the
RequestState protobuf. The Responses object maintains a
reference to this protobuf for use in the execution of the
state method. (so you can access this data by
responses.request).
start_time: Schedule the state at this time. This delays notification
and messages for processing into the future.
Raises:
ValueError: on arguments error.
"""
if messages is None:
messages = []
if not next_state:
raise ValueError("next_state can't be empty.")
# Now we construct a special response which will be sent to the hunt
# flow. Randomize the request_id so we do not overwrite other messages in
# the queue.
request_state = rdfvalue.RequestState(
id=utils.PRNG.GetULong(),
session_id=self.context.session_id,
client_id=client_id,
next_state=next_state)
if request_data:
request_state.data = rdfvalue.Dict().FromDict(request_data)
self.QueueRequest(request_state, timestamp=start_time)
# Add the status message if needed.
if not messages or not isinstance(messages[-1], rdfvalue.GrrStatus):
messages.append(rdfvalue.GrrStatus())
# Send all the messages
for i, payload in enumerate(messages):
if isinstance(payload, rdfvalue.RDFValue):
msg = rdfvalue.GrrMessage(
session_id=self.session_id,
request_id=request_state.id,
response_id=1 + i,
auth_state=rdfvalue.GrrMessage.AuthorizationState.
AUTHENTICATED,
payload=payload,
type=rdfvalue.GrrMessage.Type.MESSAGE)
if isinstance(payload, rdfvalue.GrrStatus):
msg.type = rdfvalue.GrrMessage.Type.STATUS
else:
raise flow_runner.FlowRunnerError(
"Bad message %s of type %s." % (payload, type(payload)))
self.QueueResponse(msg, timestamp=start_time)
# Add the status message if needed.
if not messages or not isinstance(messages[-1], rdfvalue.GrrStatus):
messages.append(rdfvalue.GrrStatus())
# Notify the worker about it.
self.QueueNotification(session_id=self.session_id,
timestamp=start_time)
def CallFlow(self,
flow_name=None,
next_state=None,
sync=True,
request_data=None,
client_id=None,
base_session_id=None,
output=None,
**kwargs):
"""Creates a new flow and send its responses to a state.
This creates a new flow. The flow may send back many responses which will be
queued by the framework until the flow terminates. The final status message
will cause the entire transaction to be committed to the specified state.
Args:
flow_name: The name of the flow to invoke.
next_state: The state in this flow, that responses to this
message should go to.
sync: If True start the flow inline on the calling thread, else schedule
a worker to actually start the child flow.
request_data: Any dict provided here will be available in the
RequestState protobuf. The Responses object maintains a reference
to this protobuf for use in the execution of the state method. (so
you can access this data by responses.request). There is no
format mandated on this data but it may be a serialized protobuf.
client_id: If given, the flow is started for this client.
base_session_id: A URN which will be used to build a URN.
output: A relative output name for the child collection. Normally
subflows do not write their own collections, but this can be specified
to change this behaviour.
**kwargs: Arguments for the child flow.
Raises:
FlowRunnerError: If next_state is not one of the allowed next states.
Returns:
The URN of the child flow which was created.
"""
if self.process_requests_in_order:
# Check that the next state is allowed
if next_state and next_state not in self.context.next_states:
raise FlowRunnerError(
"Flow %s: State '%s' called to '%s' which is "
"not declared in decorator." %
(self.__class__.__name__, self.context.current_state,
next_state))
client_id = client_id or self.args.client_id
# This looks very much like CallClient() above - we prepare a request state,
# and add it to our queue - any responses from the child flow will return to
# the request state and the stated next_state. Note however, that there is
# no client_id or actual request message here because we directly invoke the
# child flow rather than queue anything for it.
state = rdfvalue.RequestState(id=self.GetNextOutboundId(),
session_id=utils.SmartUnicode(
self.session_id),
client_id=client_id,
next_state=next_state,
response_count=0)
if request_data:
state.data = rdfvalue.Dict().FromDict(request_data)
# If the urn is passed explicitly (e.g. from the hunt runner) use that,
# otherwise use the urn from the flow_runner args. If both are None, create
# a new collection and give the urn to the flow object.
logs_urn = self._GetLogsCollectionURN(
kwargs.pop("logs_collection_urn", None)
or self.args.logs_collection_urn)
# If we were called with write_intermediate_results, propagate down to
# child flows. This allows write_intermediate_results to be set to True
# either at the top level parent, or somewhere in the middle of
# the call chain.
write_intermediate = (kwargs.pop("write_intermediate_results", False)
or getattr(self.args,
"write_intermediate_results", False))
# Create the new child flow but do not notify the user about it.
child_urn = self.flow_obj.StartFlow(
client_id=client_id,
flow_name=flow_name,
base_session_id=base_session_id or self.session_id,
event_id=self.context.get("event_id"),
request_state=state,
token=self.token,
notify_to_user=False,
parent_flow=self.flow_obj,
_store=self.data_store,
sync=sync,
output=output,
queue=self.args.queue,
write_intermediate_results=write_intermediate,
logs_collection_urn=logs_urn,
creator=self.context.creator,
**kwargs)
self.QueueRequest(state)
return child_urn
def CallClient(self,
action_name,
request=None,
next_state=None,
client_id=None,
request_data=None,
start_time=None,
**kwargs):
"""Calls the client asynchronously.
This sends a message to the client to invoke an Action. The run
action may send back many responses. These will be queued by the
framework until a status message is sent by the client. The status
message will cause the entire transaction to be committed to the
specified state.
Args:
action_name: The function to call on the client.
request: The request to send to the client. If not specified (Or None) we
create a new RDFValue using the kwargs.
next_state: The state in this flow, that responses to this
message should go to.
client_id: rdfvalue.ClientURN to send the request to.
request_data: A dict which will be available in the RequestState
protobuf. The Responses object maintains a reference to this
protobuf for use in the execution of the state method. (so you can
access this data by responses.request). Valid values are
strings, unicode and protobufs.
start_time: Call the client at this time. This Delays the client request
for into the future.
**kwargs: These args will be used to construct the client action semantic
protobuf.
Raises:
FlowRunnerError: If next_state is not one of the allowed next states.
RuntimeError: The request passed to the client does not have the correct
type.
"""
if client_id is None:
client_id = self.args.client_id
if client_id is None:
raise FlowRunnerError(
"CallClient() is used on a flow which was not "
"started with a client.")
if not isinstance(client_id, rdfvalue.ClientURN):
# Try turning it into a ClientURN
client_id = rdfvalue.ClientURN(client_id)
# Retrieve the correct rdfvalue to use for this client action.
try:
action = actions.ActionPlugin.classes[action_name]
except KeyError:
raise RuntimeError("Client action %s not found." % action_name)
if action.in_rdfvalue is None:
if request:
raise RuntimeError("Client action %s does not expect args." %
action_name)
else:
if request is None:
# Create a new rdf request.
request = action.in_rdfvalue(**kwargs)
else:
# Verify that the request type matches the client action requirements.
if not isinstance(request, action.in_rdfvalue):
raise RuntimeError("Client action expected %s but got %s" %
(action.in_rdfvalue, type(request)))
outbound_id = self.GetNextOutboundId()
# Create a new request state
state = rdfvalue.RequestState(id=outbound_id,
session_id=self.session_id,
next_state=next_state,
client_id=client_id)
if request_data is not None:
state.data = rdfvalue.Dict(request_data)
# Send the message with the request state
msg = rdfvalue.GrrMessage(session_id=utils.SmartUnicode(
self.session_id),
name=action_name,
request_id=outbound_id,
priority=self.args.priority,
require_fastpoll=self.args.require_fastpoll,
queue=client_id.Queue(),
payload=request)
if self.context.remaining_cpu_quota:
msg.cpu_limit = int(self.context.remaining_cpu_quota)
cpu_usage = self.context.client_resources.cpu_usage
if self.context.args.cpu_limit:
msg.cpu_limit = max(
self.context.args.cpu_limit - cpu_usage.user_cpu_time -
cpu_usage.system_cpu_time, 0)
if msg.cpu_limit == 0:
raise FlowRunnerError("CPU limit exceeded.")
if self.context.args.network_bytes_limit:
msg.network_bytes_limit = max(
self.context.args.network_bytes_limit -
self.context.network_bytes_sent, 0)
if msg.network_bytes_limit == 0:
raise FlowRunnerError("Network limit exceeded.")
state.request = msg
self.QueueRequest(state, timestamp=start_time)
