def testProvidesMultiple(self):
"""Test provides values are updated from a dictionary."""
provides = ["domain", "current_control_set"]
self.request = self.InitializeRequest(provides=provides)
self.response = rdf_protodict.Dict(
domain="MICROSOFT",
current_control_set="HKEY_LOCAL_MACHINE\\SYSTEM\\ControlSet001",
environ_systemdrive="C:")
knowledge_base = self.GetUpdatedKnowledgeBase()
self.assertEqual(knowledge_base.domain, "MICROSOFT")
self.assertEqual(knowledge_base.current_control_set,
"HKEY_LOCAL_MACHINE\\SYSTEM\\ControlSet001")
self.assertEqual(knowledge_base.environ_systemdrive, "")
def ParseFiles(
self,
knowledge_base: rdf_client.KnowledgeBase,
pathspecs: Collection[rdf_paths.PathSpec],
filedescs: Collection[file_store.BlobStream],
) -> Iterable[rdf_protodict.Dict]:
results = []
for pathspec, filedesc in zip(pathspecs, filedescs):
result = rdf_protodict.Dict()
result["path"] = pathspec.path
result["content"] = filedesc.Read()
results.append(result)
return results
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.CONFIG.SetWriteBack(self.config_file)
# Make sure the file is gone
self.assertRaises(IOError, open, self.config_file)
location = [u"http://www.example1.com/", u"http://www.example2.com/"]
request = rdf_protodict.Dict()
request["Client.server_urls"] = location
request["Client.foreman_check_frequency"] = 3600
result = self.RunAction(admin.UpdateConfiguration, request)
self.assertEqual(result, [])
self.assertEqual(config.CONFIG["Client.foreman_check_frequency"], 3600)
# Test the config file got written.
with io.open(self.config_file, "r") as filedesc:
data = filedesc.read()
server_urls = """
Client.server_urls:
- http://www.example1.com/
- http://www.example2.com/
"""
self.assertIn(server_urls, data)
self.urls = []
# Now test that our location was actually updated.
def FakeUrlOpen(url=None, data=None, **_):
self.urls.append(url)
response = requests.Response()
response.status_code = 200
response._content = data
return response
with utils.Stubber(requests, "request", FakeUrlOpen):
client_context = comms.GRRHTTPClient(worker_cls=MockClientWorker)
client_context.MakeRequest("")
# Since the request is successful we only connect to one location.
self.assertIn(location[0], self.urls[0])
def testToPrimitiveDict(self):
# See rationale for using serialized non-attributed dict above.
pdict = rdf_protodict.Dict()
pdict[b"foo"] = 42
pdict[b"bar"] = b"quux"
pdict[b"baz"] = [4, 8, 15, 16, 23, 42]
serialized = pdict.SerializeToBytes()
adict = rdf_protodict.AttributedDict.FromSerializedBytes(serialized)
dct = adict.ToDict()
self.assertEqual(dct["foo"], 42)
self.assertEqual(dct["bar"], b"quux")
self.assertEqual(dct["baz"], [4, 8, 15, 16, 23, 42])
def testFromSerializedProtoDict(self):
# In this test we use a non-attributed dict to force a serialization with
# byte keys and then we deserialize it as attributed dict that should have
# these attributes properly to unicode string keys.
pdict = rdf_protodict.Dict()
pdict[b"foo"] = 42
pdict[b"bar"] = b"quux"
pdict[b"baz"] = [4, 8, 15, 16, 23, 42]
serialized = pdict.SerializeToBytes()
adict = rdf_protodict.AttributedDict.FromSerializedBytes(serialized)
self.assertEqual(adict.foo, 42)
self.assertEqual(adict.bar, b"quux")
self.assertEqual(adict.baz, [4, 8, 15, 16, 23, 42])
def testUpdateConfig(self):
"""Ensure we can retrieve and update the config."""
# Write a client without a proper system so we don't need to
# provide the os specific artifacts in the interrogate flow below.
client_id = self.SetupClient(0, system="")
# Only mock the pieces we care about.
client_mock = action_mocks.ActionMock(admin.GetConfiguration,
admin.UpdateConfiguration)
loc = "http://www.example.com/"
new_config = rdf_protodict.Dict({
"Client.server_urls": [loc],
"Client.foreman_check_frequency": 3600,
"Client.poll_min": 1
})
# Setting config options is disallowed in tests so we need to temporarily
# revert this.
with utils.Stubber(config.CONFIG, "Set", config.CONFIG.Set.old_target):
# Write the config.
flow_test_lib.TestFlowHelper(
administrative.UpdateConfiguration.__name__,
client_mock,
client_id=client_id,
token=self.token,
config=new_config)
# Now retrieve it again to see if it got written.
flow_test_lib.TestFlowHelper(discovery.Interrogate.__name__,
client_mock,
token=self.token,
client_id=client_id)
if data_store.RelationalDBReadEnabled():
client = data_store.REL_DB.ReadClientSnapshot(client_id.Basename())
config_dat = {
item.key: item.value
for item in client.grr_configuration
}
# The grr_configuration only contains strings.
self.assertEqual(config_dat["Client.server_urls"],
"[u'http://www.example.com/']")
self.assertEqual(config_dat["Client.poll_min"], "1.0")
else:
fd = aff4.FACTORY.Open(client_id, token=self.token)
config_dat = fd.Get(fd.Schema.GRR_CONFIGURATION)
self.assertEqual(config_dat["Client.server_urls"], [loc])
self.assertEqual(config_dat["Client.poll_min"], 1)
def testGetConfig(self):
"""Check GetConfig client action works."""
# Use UpdateConfig to generate a config.
location = ["http://example.com/"]
request = rdf_protodict.Dict()
request["Client.server_urls"] = location
request["Client.foreman_check_frequency"] = 3600
self.RunAction(admin.UpdateConfiguration, request)
# Check that our GetConfig actually gets the real data.
self.RunAction(admin.GetConfiguration)
self.assertEqual(config.CONFIG["Client.foreman_check_frequency"], 3600)
self.assertEqual(config.CONFIG["Client.server_urls"], location)
def testUpdateConfig(self):
"""Ensure we can retrieve and update the config."""
# Write a client without a proper system so we don't need to
# provide the os specific artifacts in the interrogate flow below.
client_id = self.SetupClient(0, system="")
# Only mock the pieces we care about.
client_mock = action_mocks.ActionMock(admin.GetConfiguration,
admin.UpdateConfiguration)
loc = "http://www.example.com/"
new_config = rdf_protodict.Dict({
"Client.server_urls": [loc],
"Client.foreman_check_frequency": 3600,
"Client.poll_min": 1
})
# Setting config options is disallowed in tests so we need to temporarily
# revert this.
with utils.Stubber(config.CONFIG, "Set", config.CONFIG.Set.old_target):
# Write the config.
flow_test_lib.TestFlowHelper(
administrative.UpdateConfiguration.__name__,
client_mock,
client_id=client_id,
token=self.token,
config=new_config)
# Now retrieve it again to see if it got written.
flow_test_lib.TestFlowHelper(discovery.Interrogate.__name__,
client_mock,
token=self.token,
client_id=client_id)
client = data_store.REL_DB.ReadClientSnapshot(client_id)
config_dat = {
item.key: item.value
for item in client.grr_configuration
}
# The grr_configuration only contains strings.
# TODO: We use eval here, because in Python 2 these server URLs
# are a stringified list and have leading 'u' characters (to denote unicode
# literals). To overcome this difference in Python versions, we use `eval`
# to evaluate these lists. Once support for Python 2 is dropped, this can be
# again made an equality check for raw string contents.
self.assertEqual(
eval(config_dat["Client.server_urls"]), # pylint: disable=eval-used
["http://www.example.com/"])
self.assertEqual(config_dat["Client.poll_min"], "1.0")
def CallFlow(self,
flow_name: Optional[str] = None,
next_state: Optional[str] = None,
request_data: Optional[Mapping[str, Any]] = None,
**kwargs: Any) -> str:
"""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.
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.
**kwargs: Arguments for the child flow.
Returns:
The flow_id of the child flow which was created.
Raises:
ValueError: The requested next state does not exist.
"""
if not getattr(self, next_state):
raise ValueError("Next state %s is invalid." % next_state)
flow_request = rdf_flow_objects.FlowRequest(
client_id=self.rdf_flow.client_id,
flow_id=self.rdf_flow.flow_id,
request_id=self.GetNextOutboundId(),
next_state=next_state)
if request_data is not None:
flow_request.request_data = rdf_protodict.Dict().FromDict(
request_data)
self.flow_requests.append(flow_request)
flow_cls = FlowRegistry.FlowClassByName(flow_name)
return flow.StartFlow(client_id=self.rdf_flow.client_id,
flow_cls=flow_cls,
parent=flow.FlowParent.FromFlow(self),
**kwargs)
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 = rdf_crypto.SignedBlob()
signed_blob.Sign(python_code.encode("utf-8"), self.signing_key)
pdict = rdf_protodict.Dict({"test": "dict_arg", 43: "dict_arg2"})
request = rdf_client_action.ExecutePythonRequest(
python_code=signed_blob, py_args=pdict)
result = self.RunAction(standard.ExecutePython, request)[0]
self.assertEqual(result.return_val, "dict_arg")
self.assertEqual(utils.TEST_VAL, "dict_arg2")
def testEndToEndGoogleCOS(self):
parser = linux_release_parser.LinuxReleaseParser()
test_data = [
("/etc/os-release",
os.path.join(self.parser_test_dir, "google-cos-os-release")),
]
pathspecs, file_objects = self._CreateTestData(test_data)
actual_result = list(parser.ParseFiles(None, pathspecs, file_objects))
expected_result = [
rdf_protodict.Dict({
"os_release": "Container-Optimized OS",
"os_major_version": 69,
"os_minor_version": 0,
})
]
self.assertCountEqual(actual_result, expected_result)
def testEndToEndAmazon(self):
parser = linux_release_parser.LinuxReleaseParser()
test_data = [
("/etc/system-release",
os.path.join(self.parser_test_dir, "amazon-system-release")),
]
pathspecs, file_objects = self._CreateTestData(test_data)
actual_result = list(parser.ParseFiles(None, pathspecs, file_objects))
expected_result = [
rdf_protodict.Dict({
"os_release": "AmazonLinuxAMI",
"os_major_version": 2018,
"os_minor_version": 3,
})
]
self.assertCountEqual(actual_result, expected_result)
def testWMICommandLineEventConsumerParser(self):
parser = wmi_parser.WMICommandLineEventConsumerParser()
rdf_dict = rdf_protodict.Dict()
rdf_dict["CommandLineTemplate"] = "cscript KernCap.vbs"
rdf_dict["CreateNewConsole"] = False
rdf_dict["CreateNewProcessGroup"] = False
rdf_dict["CreateSeparateWowVdm"] = False
rdf_dict["CreateSharedWowVdm"] = False
rdf_dict["CreatorSID"] = [
1, 5, 0, 0, 0, 0, 0, 5, 21, 0, 0, 0, 133, 116, 119, 185, 124, 13,
122, 150, 111, 189, 41, 154, 244, 1, 0, 0
]
rdf_dict["DesktopName"] = None
rdf_dict["ExecutablePath"] = None
rdf_dict["FillAttribute"] = None
rdf_dict["ForceOffFeedback"] = False
rdf_dict["ForceOnFeedback"] = False
rdf_dict["KillTimeout"] = 0
rdf_dict["MachineName"] = None
rdf_dict["MaximumQueueSize"] = None
rdf_dict["Name"] = "BVTConsumer"
rdf_dict["Priority"] = 32
rdf_dict["RunInteractively"] = False
rdf_dict["ShowWindowCommand"] = None
rdf_dict["UseDefaultErrorMode"] = False
rdf_dict["WindowTitle"] = None
rdf_dict["WorkingDirectory"] = "C:\\tools\\kernrate"
rdf_dict["XCoordinate"] = None
rdf_dict["XNumCharacters"] = None
rdf_dict["XSize"] = None
rdf_dict["YCoordinate"] = None
rdf_dict["YNumCharacters"] = None
rdf_dict["YSize"] = None
result_list = list(parser.ParseMultiple([rdf_dict]))
self.assertLen(result_list, 1)
result = result_list[0]
self.assertEqual(result.CreatorSID,
"S-1-5-21-3111613573-2524581244-2586426735-500")
self.assertEqual(result.CommandLineTemplate, "cscript KernCap.vbs")
self.assertEqual(result.Name, "BVTConsumer")
self.assertEqual(result.KillTimeout, 0)
self.assertEqual(result.FillAttribute, 0)
self.assertEqual(result.FillAttributes, 0)
self.assertFalse(result.ForceOffFeedback)
self.assertFalse(result.ForceOnFeedback)
def testWMIEventConsumerParserDoesntFailOnMalformedSIDs(self):
parser = wmi_parser.WMIActiveScriptEventConsumerParser()
rdf_dict = rdf_protodict.Dict()
tests = [
[1, 5, 0, 0, 0, 0, 0, 5, 21, 0, 0],
"(1, 2, 3)", # Older clients (3.0.0.3) return a the SID like this
1,
{
1: 2
},
(1, 2)
]
for test in tests:
rdf_dict["CreatorSID"] = test
result_list = list(parser.Parse(None, rdf_dict, None))
self.assertEqual(len(result_list), 1)
def testToDictNestedLists(self):
dct = {
"foo": [
[42],
[1, 3, 3, 7],
],
"bar": [
{
"quux": [4, 8, 15],
"thud": [16, 23],
},
{
"norf": [42],
},
],
}
self.assertEqual(rdf_protodict.Dict(dct).ToDict(), dct)
def testEndToEndCoreOS(self):
parser = linux_release_parser.LinuxReleaseParser()
test_data = [
("/etc/os-release",
os.path.join(self.parser_test_dir, "coreos-os-release")),
]
stat_entries, file_objects = self._CreateTestData(test_data)
actual_result = list(
parser.ParseMultiple(stat_entries, file_objects, None))
expected_result = [
rdf_protodict.Dict({
"os_release": "Container Linux by CoreOS",
"os_major_version": 2023,
"os_minor_version": 4,
})
]
self.assertCountEqual(actual_result, expected_result)
def testConvertsDictWithPrimitiveValues(self):
source = rdf_protodict.Dict()
source["foo"] = "bar"
source["bar"] = 42
# Serialize/unserialize to make sure we deal with the object that is
# similar to what we may get from the datastore.
source = rdf_protodict.Dict.FromSerializedBytes(source.SerializeToBytes())
converted = list(self.converter.Convert(self.metadata, source))
self.assertLen(converted, 2)
# Output should be stable sorted by dict's keys.
self.assertEqual(converted[0].key, "bar")
self.assertEqual(converted[0].value, "42")
self.assertEqual(converted[1].key, "foo")
self.assertEqual(converted[1].value, "bar")
def testOnlyUpdatableFieldsAreUpdated(self):
with test_lib.ConfigOverrider({
"Client.server_urls": [u"http://something.com/"],
"Client.server_serial_number": 1
}):
location = [u"http://www.example.com"]
request = rdf_protodict.Dict()
request["Client.server_urls"] = location
request["Client.server_serial_number"] = 10
with self.assertRaises(ValueError):
self.RunAction(admin.UpdateConfiguration, request)
# Nothing was updated.
self.assertEqual(config.CONFIG["Client.server_urls"],
[u"http://something.com/"])
self.assertEqual(config.CONFIG["Client.server_serial_number"], 1)
def testInterfaceParsing(self):
parser = wmi_parser.WMIInterfacesParser()
rdf_dict = rdf_protodict.Dict()
mock_config = client_test_lib.WMIWin32NetworkAdapterConfigurationMock
wmi_properties = mock_config.__dict__.items()
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.ParseMultiple([rdf_dict]))
self.assertLen(result_list, 2)
for result in result_list:
if isinstance(result, rdf_client_network.Interface):
self.assertLen(result.addresses, 4)
self.assertCountEqual(
[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.assertCountEqual([
x.human_readable_address for x in result.dhcp_server_list
], ["192.168.1.1"])
self.assertEqual(
result.dhcp_lease_expires.AsMicrosecondsSinceEpoch(),
1409008979123456)
self.assertEqual(
result.dhcp_lease_obtained.AsMicrosecondsSinceEpoch(),
1408994579123456)
elif isinstance(result, rdf_client_network.DNSClientConfiguration):
self.assertCountEqual(
result.dns_server,
["192.168.1.1", "192.168.255.81", "192.168.128.88"])
self.assertCountEqual(result.dns_suffix, [
"blah.example.com", "ad.example.com",
"internal.example.com", "example.com"
])
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:
rdf_protodict.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 = rdf_protodict.Dict()
properties = (list(result.Properties_) +
list(getattr(result, "SystemProperties_", [])))
for prop in 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 testUpdateConfigBlacklist(self):
"""Tests that disallowed fields are not getting updated."""
with test_lib.ConfigOverrider({
"Client.server_urls": ["http://something.com/"],
"Client.server_serial_number": 1
}):
location = ["http://www.example.com"]
request = rdf_protodict.Dict()
request["Client.server_urls"] = location
request["Client.server_serial_number"] = 10
with self.assertRaises(ValueError):
self.RunAction(admin.UpdateConfiguration, request)
# Nothing was updated.
self.assertEqual(config.CONFIG["Client.server_urls"],
["http://something.com/"])
self.assertEqual(config.CONFIG["Client.server_serial_number"], 1)
def testRdfFormatterFanOut(self):
rdf = rdf_protodict.Dict()
user1 = rdf_client.User(username="******")
user2 = rdf_client.User(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}\n")
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 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 = rdf_protodict.Dict()
mock = client_test_lib.WMIWin32NetworkAdapterConfigurationMock
wmi_properties = iteritems(mock.__dict__)
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 []
def _ParseOSReleaseFile(self, matches_dict):
# The spec for the os-release file is given at
# https://www.freedesktop.org/software/systemd/man/os-release.html.
try:
os_release_contents = matches_dict['/etc/os-release']
except KeyError:
return None
os_release_name = None
os_major_version = None
os_minor_version = None
for entry in os_release_contents.splitlines():
entry_parts = entry.split('=', 1)
if len(entry_parts) != 2:
continue
key = entry_parts[0].strip()
# Remove whitespace and quotes from the value (leading and trailing).
value = entry_parts[1].strip('\t \'"')
if key == _SYSTEMD_OS_RELEASE_NAME:
os_release_name = value
elif key == _SYSTEMD_OS_RELEASE_VERSION:
match_object = re.search(r'(?P<major>\d+)\.?(?P<minor>\d+)?',
value)
if match_object is not None:
os_major_version = int(match_object.group('major'))
minor_match = match_object.group('minor')
# Some platforms (e.g. Google's Container-Optimized OS) do not have
# multi-part version numbers so we use a default minor version of
# zero.
os_minor_version = 0 if minor_match is None else int(
minor_match)
if (os_release_name and os_major_version is not None
and os_minor_version is not None):
return rdf_protodict.Dict({
'os_release': os_release_name,
'os_major_version': os_major_version,
'os_minor_version': os_minor_version,
})
return None
def testConvertsDictWithNestedSetListOrTuple(self):
# Note that set's contents will be sorted on export.
variants = [set([43, 42, 44]), (42, 43, 44), [42, 43, 44]]
for variant in variants:
source = rdf_protodict.Dict()
source["foo"] = "bar"
source["bar"] = variant
# Serialize/unserialize to make sure we deal with the object that is
# similar to what we may get from the datastore.
source = rdf_protodict.Dict.FromSerializedBytes(source.SerializeToBytes())
converted = list(self.converter.Convert(self.metadata, source))
self.assertLen(converted, 4)
self.assertEqual(converted[0].key, "bar[0]")
self.assertEqual(converted[0].value, "42")
self.assertEqual(converted[1].key, "bar[1]")
self.assertEqual(converted[1].value, "43")
self.assertEqual(converted[2].key, "bar[2]")
self.assertEqual(converted[2].value, "44")
self.assertEqual(converted[3].key, "foo")
self.assertEqual(converted[3].value, "bar")
def testConvertsDictWithNestedDictAndIterables(self):
source = rdf_protodict.Dict()
source["foo"] = "bar"
# pyformat: disable
source["bar"] = {
"a": {
"c": [42, 43, 44, {"x": "y"}],
"d": "oh"
},
"b": 43
}
# pyformat: enable
# Serialize/unserialize to make sure we deal with the object that is
# similar to what we may get from the datastore.
source = rdf_protodict.Dict.FromSerializedBytes(source.SerializeToBytes())
converted = list(self.converter.Convert(self.metadata, source))
self.assertLen(converted, 7)
# Output should be stable sorted by dict's keys.
self.assertEqual(converted[0].key, "bar.a.c[0]")
self.assertEqual(converted[0].value, "42")
self.assertEqual(converted[1].key, "bar.a.c[1]")
self.assertEqual(converted[1].value, "43")
self.assertEqual(converted[2].key, "bar.a.c[2]")
self.assertEqual(converted[2].value, "44")
self.assertEqual(converted[3].key, "bar.a.c[3].x")
self.assertEqual(converted[3].value, "y")
self.assertEqual(converted[4].key, "bar.a.d")
self.assertEqual(converted[4].value, "oh")
self.assertEqual(converted[5].key, "bar.b")
self.assertEqual(converted[5].value, "43")
self.assertEqual(converted[6].key, "foo")
self.assertEqual(converted[6].value, "bar")
def WmiQuery(self, query):
expected_query = "SELECT * FROM Win32_ShadowCopy"
if query.query != expected_query:
raise RuntimeError("Received unexpected query.")
return [rdf_protodict.Dict(**self._RESPONSES)]
def CallFlow(self,
flow_name=None,
next_state=None,
request_data=None,
client_id=None,
base_session_id=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.
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.
**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.
"""
client_id = client_id or self.runner_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 = rdf_flow_runner.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 = rdf_protodict.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._GetLogCollectionURN(
kwargs.pop("logs_collection_urn", None) or
self.runner_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
self.runner_args.write_intermediate_results)
# Create the new child flow but do not notify the user about it.
child_urn = self.flow_obj.StartAFF4Flow(
client_id=client_id,
flow_name=flow_name,
base_session_id=base_session_id or self.session_id,
request_state=state,
token=self.token,
notify_to_user=False,
parent_flow=self.flow_obj,
queue=self.runner_args.queue,
write_intermediate_results=write_intermediate,
logs_collection_urn=logs_urn,
sync=True,
**kwargs)
self.QueueRequest(state)
return child_urn
def CallClient(self,
action_cls,
request=None,
next_state=None,
request_data=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_cls: 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.
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.
**kwargs: These args will be used to construct the client action semantic
protobuf.
Raises:
FlowRunnerError: If called on a flow that doesn't run on a single client.
ValueError: The request passed to the client does not have the correct
type.
"""
client_id = self.runner_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, rdf_client.ClientURN):
# Try turning it into a ClientURN
client_id = rdf_client.ClientURN(client_id)
if action_cls.in_rdfvalue is None:
if request:
raise ValueError(
"Client action %s does not expect args." % action_cls.__name__)
else:
if request is None:
# Create a new rdf request.
request = action_cls.in_rdfvalue(**kwargs)
else:
# Verify that the request type matches the client action requirements.
if not isinstance(request, action_cls.in_rdfvalue):
raise ValueError("Client action expected %s but got %s" %
(action_cls.in_rdfvalue, type(request)))
outbound_id = self.GetNextOutboundId()
# Create a new request state
state = rdf_flow_runner.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 = rdf_protodict.Dict(request_data)
# Send the message with the request state
msg = rdf_flows.GrrMessage(
session_id=utils.SmartUnicode(self.session_id),
name=action_cls.__name__,
request_id=outbound_id,
require_fastpoll=self.runner_args.require_fastpoll,
queue=client_id.Queue(),
payload=request,
generate_task_id=True)
cpu_usage = self.context.client_resources.cpu_usage
if self.runner_args.cpu_limit:
msg.cpu_limit = max(
self.runner_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.runner_args.network_bytes_limit:
msg.network_bytes_limit = max(
self.runner_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)
from grr.test_lib import action_mocks
from grr.test_lib import artifact_test_lib
from grr.test_lib import client_test_lib
from grr.test_lib import flow_test_lib
from grr.test_lib import parser_test_lib
from grr.test_lib import test_lib
from grr.test_lib import vfs_test_lib
# pylint: mode=test
WMI_SAMPLE = [
rdf_protodict.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"
}),
rdf_protodict.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"
}),
rdf_protodict.Dict({
u"Version": u"8.0.61000",
