Esempio n. 1
0
 def Wrapper(*args, **kwargs):  # pylint: disable=missing-docstring
     with mock.patch.object(parsers, "SINGLE_RESPONSE_PARSER_FACTORY",
                            factory.Factory(parsers.SingleResponseParser)),\
          mock.patch.object(parsers, "SINGLE_FILE_PARSER_FACTORY",
                            factory.Factory(parsers.SingleFileParser)),\
          mock.patch.object(parsers, "MULTI_RESPONSE_PARSER_FACTORY",
                            factory.Factory(parsers.MultiResponseParser)),\
          mock.patch.object(parsers, "MULTI_FILE_PARSER_FACTORY",
                            factory.Factory(parsers.MultiFileParser)):
         all_parsers.Register()
         func(*args, **kwargs)
Esempio n. 2
0
  def testRegisterDuplicateThrows(self):
    obj_factory = factory.Factory(object)
    obj_factory.Register("foo", object)
    obj_factory.Register("bar", object)

    with self.assertRaisesRegex(ValueError, "foo"):
      obj_factory.Register("foo", object)
Esempio n. 3
0
    def setUp(self):
        super(DecodersTestMixin, self).setUp()
        self.client_id = self.SetupClient(0)

        self.decoders_patcher = mock.patch.object(
            decoders, "FACTORY", factory.Factory(decoders.AbstractDecoder))
        self.decoders_patcher.start()
Esempio n. 4
0
    def testCreateAllSome(self):
        int_factory = factory.Factory(int)
        int_factory.Register("foo", lambda: 1337)
        int_factory.Register("bar", lambda: 101)
        int_factory.Register("baz", lambda: 108)

        self.assertCountEqual(list(int_factory.CreateAll()), [1337, 101, 108])
Esempio n. 5
0
File: vfs_test.py Progetto: avmi/grr
  def setUp(self):
    super().setUp()
    self.client_id = self.SetupClient(0)

    decoders_patcher = mock.patch.object(
        decoders, "FACTORY", factory.Factory(decoders.AbstractDecoder))
    decoders_patcher.start()
    self.addCleanup(decoders_patcher.stop)
Esempio n. 6
0
  def testCreateString(self):
    str_factory = factory.Factory(str)
    str_factory.Register("foo", str, lambda: "FOO")
    str_factory.Register("bar", str, lambda: "BAR")
    str_factory.Register("baz", str, lambda: "BAZ")

    self.assertEqual(str_factory.Create("foo"), "FOO")
    self.assertEqual(str_factory.Create("bar"), "BAR")
    self.assertEqual(str_factory.Create("baz"), "BAZ")
Esempio n. 7
0
    def testGetTypesReturnsAllTypes(self):
        class Foo(object):
            pass

        class Bar(object):
            pass

        int_factory = factory.Factory(object)
        int_factory.Register("foo", Foo)
        int_factory.Register("bar", Bar)

        self.assertCountEqual(list(int_factory.GetTypes()), [Foo, Bar])
Esempio n. 8
0
    def testCreateClass(self):
        class Foo(object):
            pass

        class Bar(object):
            pass

        cls_factory = factory.Factory(object)
        cls_factory.Register("Foo", Foo)
        cls_factory.Register("Bar", Bar)

        self.assertIsInstance(cls_factory.Create("Foo"), Foo)
        self.assertIsInstance(cls_factory.Create("Bar"), Bar)
Esempio n. 9
0
  def testRegisterAndUnregister(self):
    del self  # Unused.

    obj_factory = factory.Factory(object)

    # First, we check whether registering works.
    obj_factory.Register("foo", object)
    obj_factory.Register("bar", object)

    # Now, we should be able to unregister these constructors.
    obj_factory.Unregister("foo")
    obj_factory.Unregister("bar")

    # Once they are unregistered, names are free to be bound again.
    obj_factory.Register("foo", object)
    obj_factory.Register("bar", object)
Esempio n. 10
0
from typing import Text

from grr_response_core.lib import factory
from grr_response_core.lib.parsers import abstract
from grr_response_core.lib.util import collection
from grr_response_core.lib.util import precondition

ParseError = abstract.ParseError

Parser = abstract.Parser
SingleResponseParser = abstract.SingleResponseParser
SingleFileParser = abstract.SingleFileParser
MultiResponseParser = abstract.MultiResponseParser
MultiFileParser = abstract.MultiFileParser

SINGLE_RESPONSE_PARSER_FACTORY = factory.Factory(SingleResponseParser)
MULTI_RESPONSE_PARSER_FACTORY = factory.Factory(MultiResponseParser)
SINGLE_FILE_PARSER_FACTORY = factory.Factory(SingleFileParser)
MULTI_FILE_PARSER_FACTORY = factory.Factory(MultiFileParser)


class ArtifactParserFactory(object):
    """A factory wrapper class that yields parsers for specific artifact."""
    def __init__(self, artifact_name):
        """Initializes the artifact parser factory.

    Args:
      artifact_name: A name of the artifact this factory is supposed to provide
        parser instances for.
    """
        precondition.AssertType(artifact_name, Text)
Esempio n. 11
0
class ArtifactParserFactoryTest(absltest.TestCase):
    @mock.patch.object(parsers, "SINGLE_RESPONSE_PARSER_FACTORY",
                       factory.Factory(parser.SingleResponseParser))
    def testSingleResponseParsers(self):
        class FooParser(parser.SingleResponseParser):

            supported_artifacts = ["Quux", "Norf"]

            def ParseResponse(self, knowledge_base, response, path_type):
                raise NotImplementedError()

        class BarParser(parser.SingleResponseParser):

            supported_artifacts = ["Norf", "Thud"]

            def ParseResponse(self, knowledge_base, response, path_type):
                raise NotImplementedError()

        class BazParser(parser.SingleResponseParser):

            supported_artifacts = ["Thud", "Quux"]

            def ParseResponse(self, knowledge_base, response, path_type):
                raise NotImplementedError()

        parsers.SINGLE_RESPONSE_PARSER_FACTORY.Register("Foo", FooParser)
        parsers.SINGLE_RESPONSE_PARSER_FACTORY.Register("Bar", BarParser)
        parsers.SINGLE_RESPONSE_PARSER_FACTORY.Register("Baz", BazParser)

        quux_factory = parsers.ArtifactParserFactory("Quux")
        quux_parsers = quux_factory.SingleResponseParsers()
        self.assertCountEqual(map(type, quux_parsers), [FooParser, BazParser])

        norf_factory = parsers.ArtifactParserFactory("Norf")
        norf_parsers = norf_factory.SingleResponseParsers()
        self.assertCountEqual(map(type, norf_parsers), [FooParser, BarParser])

        thud_factory = parsers.ArtifactParserFactory("Thud")
        thud_parsers = thud_factory.SingleResponseParsers()
        self.assertCountEqual(map(type, thud_parsers), [BarParser, BazParser])

    @mock.patch.object(parsers, "MULTI_RESPONSE_PARSER_FACTORY",
                       factory.Factory(parser.MultiResponseParser))
    def testMultiResponseParsers(self):
        class FooParser(parser.MultiResponseParser):

            supported_artifacts = ["Foo"]

            def ParseResponses(self, knowledge_base, responses):
                raise NotImplementedError()

        class BarParser(parser.MultiResponseParser):

            supported_artifacts = ["Bar"]

            def ParseResponses(self, knowledge_base, responses):
                raise NotImplementedError()

        parsers.MULTI_RESPONSE_PARSER_FACTORY.Register("Foo", FooParser)
        parsers.MULTI_RESPONSE_PARSER_FACTORY.Register("Bar", BarParser)

        foo_factory = parsers.ArtifactParserFactory("Foo")
        foo_parsers = foo_factory.MultiResponseParsers()
        self.assertCountEqual(map(type, foo_parsers), [FooParser])

        bar_factory = parsers.ArtifactParserFactory("Bar")
        bar_parsers = bar_factory.MultiResponseParsers()
        self.assertCountEqual(map(type, bar_parsers), [BarParser])

    @mock.patch.object(parsers, "SINGLE_FILE_PARSER_FACTORY",
                       factory.Factory(parser.SingleFileParser))
    def testSingleFileParsers(self):
        class FooParser(parser.SingleFileParser):

            supported_artifacts = ["Bar"]

            def ParseFile(self, knowledge_base, pathspec, filedesc):
                raise NotImplementedError()

        parsers.SINGLE_FILE_PARSER_FACTORY.Register("Foo", FooParser)

        bar_factory = parsers.ArtifactParserFactory("Bar")
        bar_parsers = bar_factory.SingleFileParsers()
        self.assertCountEqual(map(type, bar_parsers), [FooParser])

        baz_factory = parsers.ArtifactParserFactory("Baz")
        baz_parsers = baz_factory.SingleFileParsers()
        self.assertCountEqual(map(type, baz_parsers), [])

    @mock.patch.object(parsers, "MULTI_FILE_PARSER_FACTORY",
                       factory.Factory(parser.MultiFileParser))
    def testMultiFileParsers(self):
        class FooParser(parser.MultiFileParser):

            supported_artifacts = ["Quux", "Norf"]

            def ParseFiles(self, knowledge_base, pathspecs, filedescs):
                raise NotImplementedError()

        class BarParser(parser.MultiFileParser):

            supported_artifacts = ["Quux", "Thud"]

            def ParseFiles(self, knowledge_base, pathspecs, filedescs):
                raise NotImplementedError()

        parsers.MULTI_FILE_PARSER_FACTORY.Register("Foo", FooParser)
        parsers.MULTI_FILE_PARSER_FACTORY.Register("Bar", BarParser)

        quux_factory = parsers.ArtifactParserFactory("Quux")
        quux_parsers = quux_factory.MultiFileParsers()
        self.assertCountEqual(map(type, quux_parsers), [FooParser, BarParser])

        norf_factory = parsers.ArtifactParserFactory("Norf")
        norf_parsers = norf_factory.MultiFileParsers()
        self.assertCountEqual(map(type, norf_parsers), [FooParser])

        thud_factory = parsers.ArtifactParserFactory("Thud")
        thud_parsers = thud_factory.MultiFileParsers()
        self.assertCountEqual(map(type, thud_parsers), [BarParser])
Esempio n. 12
0
class ParseResponsesTest(client_test_lib.EmptyActionTest):
    @mock.patch.object(parsers, "SINGLE_RESPONSE_PARSER_FACTORY",
                       factory.Factory(parser.SingleResponseParser))
    def testCmdArtifactAction(self):
        """Test the actual client action with parsers."""
        parsers.SINGLE_RESPONSE_PARSER_FACTORY.Register(
            "Cmd", TestEchoCmdParser)

        client_test_lib.Command("/bin/echo", args=["1"])

        source = rdf_artifact.ArtifactSource(
            type=rdf_artifact.ArtifactSource.SourceType.COMMAND,
            attributes={
                "cmd": "/bin/echo",
                "args": ["1"]
            })
        ext_src = rdf_artifact.ExpandedSource(base_source=source)
        ext_art = rdf_artifact.ExpandedArtifact(name="TestEchoCmdArtifact",
                                                sources=[ext_src])
        request = rdf_artifact.ClientArtifactCollectorArgs(
            artifacts=[ext_art],
            knowledge_base=None,
            ignore_interpolation_errors=True,
            apply_parsers=True)
        result = self.RunAction(artifact_collector.ArtifactCollector,
                                request)[0]
        self.assertIsInstance(result,
                              rdf_artifact.ClientArtifactCollectorResult)
        self.assertLen(result.collected_artifacts, 1)
        res = result.collected_artifacts[0].action_results[0].value
        self.assertIsInstance(res, rdf_client.SoftwarePackage)
        self.assertEqual(res.description, "1\n")

    @mock.patch.object(parsers, "SINGLE_FILE_PARSER_FACTORY",
                       factory.Factory(parser.SingleFileParser))
    def testFakeFileArtifactAction(self):
        """Test collecting a file artifact and parsing the response."""
        parsers.SINGLE_FILE_PARSER_FACTORY.Register("Fake", FakeFileParser)

        file_path = os.path.join(self.base_path, "numbers.txt")
        source = rdf_artifact.ArtifactSource(
            type=rdf_artifact.ArtifactSource.SourceType.FILE,
            attributes={"paths": [file_path]})

        ext_src = rdf_artifact.ExpandedSource(base_source=source)
        ext_art = rdf_artifact.ExpandedArtifact(name="FakeFileArtifact",
                                                sources=[ext_src])
        request = rdf_artifact.ClientArtifactCollectorArgs(
            artifacts=[ext_art],
            knowledge_base=None,
            ignore_interpolation_errors=True,
            apply_parsers=True)
        result = self.RunAction(artifact_collector.ArtifactCollector,
                                request)[0]
        self.assertLen(result.collected_artifacts[0].action_results, 1)
        res = result.collected_artifacts[0].action_results[0].value
        self.assertIsInstance(res, rdf_protodict.AttributedDict)
        self.assertLen(res.users, 1000)
        self.assertEqual(res.filename, file_path)

    @mock.patch.object(parsers, "MULTI_FILE_PARSER_FACTORY",
                       factory.Factory(parser.MultiFileParser))
    def testFakeFileArtifactActionProcessTogether(self):
        """Test collecting a file artifact and parsing the responses together."""
        parsers.MULTI_FILE_PARSER_FACTORY.Register("Fake", FakeFileMultiParser)

        file_path = os.path.join(self.base_path, "numbers.txt")
        source = rdf_artifact.ArtifactSource(
            type=rdf_artifact.ArtifactSource.SourceType.FILE,
            attributes={"paths": [file_path]})

        ext_src = rdf_artifact.ExpandedSource(base_source=source)
        ext_art = rdf_artifact.ExpandedArtifact(name="FakeFileArtifact2",
                                                sources=[ext_src])
        request = rdf_artifact.ClientArtifactCollectorArgs(
            artifacts=[ext_art],
            knowledge_base=None,
            ignore_interpolation_errors=True,
            apply_parsers=True)
        result = self.RunAction(artifact_collector.ArtifactCollector,
                                request)[0]
        self.assertLen(result.collected_artifacts[0].action_results, 1)
        res = result.collected_artifacts[0].action_results[0].value
        self.assertIsInstance(res, rdf_protodict.AttributedDict)
        self.assertLen(res.users, 1000)
        self.assertEqual(res.filename, file_path)
Esempio n. 13
0
class ClientArtifactCollectorFlowTest(flow_test_lib.FlowTestsBaseclass):
    """Test the client side artifact collection test artifacts."""
    def setUp(self):
        super(ClientArtifactCollectorFlowTest, self).setUp()
        self.cleanup = None
        InitGRRWithTestArtifacts()

        self.client_id = self.SetupClient(0)

    def tearDown(self):
        super(ClientArtifactCollectorFlowTest, self).tearDown()
        if self.cleanup:
            self.cleanup()

        artifact_registry.REGISTRY.ClearSources()
        artifact_registry.REGISTRY.ClearRegistry()
        artifact_registry.REGISTRY.AddDefaultSources()

    def _RunFlow(self, flow_cls, action, artifact_list, apply_parsers):
        session_id = flow_test_lib.TestFlowHelper(
            flow_cls.__name__,
            action_mocks.ActionMock(action),
            artifact_list=artifact_list,
            token=self.token,
            apply_parsers=apply_parsers,
            client_id=self.client_id)
        return flow_test_lib.GetFlowResults(self.client_id, session_id)

    def InitializeTestFileArtifact(self, with_pathspec_attribute=False):
        file_path = os.path.join(self.base_path, "numbers.txt")
        source = rdf_artifacts.ArtifactSource(
            type=rdf_artifacts.ArtifactSource.SourceType.FILE,
            attributes={"paths": [file_path]})
        if with_pathspec_attribute:
            source.attributes = {
                "paths": [file_path],
                "pathspec_attribute": "pathspec"
            }
        artifact_obj = artifact_registry.REGISTRY.GetArtifact(
            "TestFileArtifact")
        artifact_obj.sources.append(source)
        return file_path

    def testClientArtifactCollector(self):
        """Test artifact collector flow with a single artifact."""

        client_test_lib.Command("/usr/bin/dpkg",
                                args=["--list"],
                                system="Linux")

        artifact_list = ["TestCmdArtifact"]

        results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                artifact_collector.ArtifactCollector,
                                artifact_list,
                                apply_parsers=False)
        self.assertLen(results, 1)

        artifact_response = results[0]
        self.assertIsInstance(artifact_response,
                              rdf_client_action.ExecuteResponse)
        self.assertGreater(artifact_response.time_used, 0)

    def testClientArtifactCollectorWithMultipleArtifacts(self):
        """Test artifact collector flow with a single artifact."""

        client_test_lib.Command("/usr/bin/dpkg",
                                args=["--list"],
                                system="Linux")

        artifact_list = ["TestCmdArtifact", "TestOSAgnostic"]

        results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                artifact_collector.ArtifactCollector,
                                artifact_list,
                                apply_parsers=False)
        self.assertLen(results, 2)

        artifact_response = results[0]
        self.assertIsInstance(artifact_response,
                              rdf_client_action.ExecuteResponse)
        self.assertGreater(artifact_response.time_used, 0)

        artifact_response = results[1]
        self.assertTrue(artifact_response.string)

    def testLinuxMountCmdArtifact(self):
        """Test that LinuxMountCmd artifact can be collected."""

        artifact_list = ["LinuxMountCmd"]

        self.cleanup = InitGRRWithTestSources("""
name: LinuxMountCmd
doc: Linux output of mount.
sources:
- type: COMMAND
  attributes:
    cmd: '/bin/mount'
    args: []
labels: [System]
supported_os: [Linux]
""")

        self.assertTrue(
            artifact_registry.REGISTRY.GetArtifact("LinuxMountCmd"))

        # Run the ArtifactCollector to get the expected result.
        expected = self._RunFlow(aff4_flows.ArtifactCollectorFlow,
                                 standard.ExecuteCommand,
                                 artifact_list,
                                 apply_parsers=False)
        expected = expected[0]
        self.assertIsInstance(expected, rdf_client_action.ExecuteResponse)

        # Run the ClientArtifactCollector to get the actual result.
        results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                artifact_collector.ArtifactCollector,
                                artifact_list,
                                apply_parsers=False)
        artifact_response = results[0]
        self.assertIsInstance(artifact_response,
                              rdf_client_action.ExecuteResponse)

        self.assertEqual(artifact_response, expected)

    def testBasicRegistryKeyArtifact(self):
        """Test that a registry key artifact can be collected."""

        artifact_list = ["TestRegistryKey"]

        self.cleanup = InitGRRWithTestSources(r"""
name: TestRegistryKey
doc: A sample registry key artifact.
sources:
- type: REGISTRY_KEY
  attributes:
    keys: [
      'HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Control\Session Manager'
    ]
""")

        with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.REGISTRY,
                                       vfs_test_lib.FakeRegistryVFSHandler):
            with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
                                           vfs_test_lib.FakeFullVFSHandler):

                # Run the ArtifactCollector to get the expected result.
                session_id = flow_test_lib.TestFlowHelper(
                    aff4_flows.ArtifactCollectorFlow.__name__,
                    action_mocks.FileFinderClientMock(),
                    artifact_list=artifact_list,
                    token=self.token,
                    client_id=self.client_id,
                    apply_parsers=False)
                results = flow_test_lib.GetFlowResults(self.client_id,
                                                       session_id)
                expected = results[0]
                self.assertIsInstance(expected, rdf_client_fs.StatEntry)

                # Run the ClientArtifactCollector to get the actual result.
                cac_results = self._RunFlow(
                    aff4_flows.ClientArtifactCollector,
                    artifact_collector.ArtifactCollector,
                    artifact_list,
                    apply_parsers=False)
                artifact_response = cac_results[0]
                self.assertIsInstance(artifact_response,
                                      rdf_client_fs.StatEntry)

                self.assertEqual(results, cac_results)

    def testRegistryKeyArtifactWithWildcard(self):
        """Test that a registry key artifact can be collected."""

        artifact_list = ["TestRegistryKey"]

        self.cleanup = InitGRRWithTestSources(r"""
name: TestRegistryKey
doc: A sample registry key artifact.
sources:
- type: REGISTRY_KEY
  attributes:
    keys: [
      'HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Control\Session Manager\*'
    ]
""")

        with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.REGISTRY,
                                       vfs_test_lib.FakeRegistryVFSHandler):
            with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
                                           vfs_test_lib.FakeFullVFSHandler):

                # Run the ArtifactCollector to get the expected result.
                session_id = flow_test_lib.TestFlowHelper(
                    aff4_flows.ArtifactCollectorFlow.__name__,
                    action_mocks.FileFinderClientMock(),
                    artifact_list=artifact_list,
                    token=self.token,
                    client_id=self.client_id,
                    apply_parsers=False)
                results = flow_test_lib.GetFlowResults(self.client_id,
                                                       session_id)
                self.assertIsInstance(results[0], rdf_client_fs.StatEntry)

                # Run the ClientArtifactCollector to get the actual result.
                cac_results = self._RunFlow(
                    aff4_flows.ClientArtifactCollector,
                    artifact_collector.ArtifactCollector,
                    artifact_list,
                    apply_parsers=False)
                artifact_response = cac_results[0]
                self.assertIsInstance(artifact_response,
                                      rdf_client_fs.StatEntry)

                self.assertEqual(cac_results, results)

    def testRegistryKeyArtifactWithPathRecursion(self):
        """Test that a registry key artifact can be collected."""

        artifact_list = ["TestRegistryKey"]

        self.cleanup = InitGRRWithTestSources(r"""
name: TestRegistryKey
doc: A sample registry key artifact.
sources:
- type: REGISTRY_KEY
  attributes:
    keys: [
      'HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\**\Session Manager\*'
    ]
""")

        with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.REGISTRY,
                                       vfs_test_lib.FakeRegistryVFSHandler):
            with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
                                           vfs_test_lib.FakeFullVFSHandler):

                # Run the ArtifactCollector to get the expected result.
                session_id = flow_test_lib.TestFlowHelper(
                    aff4_flows.ArtifactCollectorFlow.__name__,
                    action_mocks.FileFinderClientMock(),
                    artifact_list=artifact_list,
                    token=self.token,
                    client_id=self.client_id,
                    apply_parsers=False)
                expected = flow_test_lib.GetFlowResults(
                    self.client_id, session_id)[0]
                self.assertIsInstance(expected, rdf_client_fs.StatEntry)

                # Run the ClientArtifactCollector to get the actual result.
                results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                        artifact_collector.ArtifactCollector,
                                        artifact_list,
                                        apply_parsers=False)
                artifact_response = results[0]
                self.assertIsInstance(artifact_response,
                                      rdf_client_fs.StatEntry)

                self.assertEqual(artifact_response, expected)

    @mock.patch.object(parsers, "SINGLE_RESPONSE_PARSER_FACTORY",
                       factory.Factory(parser.SingleResponseParser))
    def testCmdArtifactWithParser(self):
        """Test a command artifact and parsing the response."""

        client_test_lib.Command("/bin/echo", args=["1"])

        parsers.SINGLE_RESPONSE_PARSER_FACTORY.Register(
            "TestCmd", TestCmdParser)
        try:
            artifact_list = ["TestEchoArtifact"]

            # Run the ArtifactCollector to get the expected result.
            expected = self._RunFlow(aff4_flows.ArtifactCollectorFlow,
                                     standard.ExecuteCommand,
                                     artifact_list,
                                     apply_parsers=True)
            self.assertTrue(expected)
            expected = expected[0]
            self.assertIsInstance(expected, rdf_client.SoftwarePackage)

            # Run the ClientArtifactCollector to get the actual result.
            results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                    artifact_collector.ArtifactCollector,
                                    artifact_list,
                                    apply_parsers=True)
            self.assertLen(results, 1)
            artifact_response = results[0]
            self.assertIsInstance(artifact_response,
                                  rdf_client.SoftwarePackage)

            self.assertEqual(artifact_response, expected)
        finally:
            parsers.SINGLE_RESPONSE_PARSER_FACTORY.Unregister("TestCmd")

    @mock.patch.object(parsers, "SINGLE_FILE_PARSER_FACTORY",
                       factory.Factory(parser.SingleFileParser))
    def testFileArtifactWithParser(self):
        """Test collecting a file artifact and parsing the response."""
        parsers.SINGLE_FILE_PARSER_FACTORY.Register("TestFile", TestFileParser)
        try:
            artifact_list = ["TestFileArtifact"]

            file_path = self.InitializeTestFileArtifact()

            # Run the ArtifactCollector to get the expected result.
            session_id = flow_test_lib.TestFlowHelper(
                aff4_flows.ArtifactCollectorFlow.__name__,
                action_mocks.FileFinderClientMock(),
                artifact_list=artifact_list,
                token=self.token,
                apply_parsers=True,
                client_id=self.client_id)
            results = flow_test_lib.GetFlowResults(self.client_id, session_id)
            expected = results[0]
            self.assertIsInstance(expected, rdf_protodict.AttributedDict)
            self.assertEqual(expected.filename, file_path)
            self.assertLen(expected.users, 1000)

            # Run the ClientArtifactCollector to get the actual result.
            cac_results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                        artifact_collector.ArtifactCollector,
                                        artifact_list,
                                        apply_parsers=True)
            self.assertLen(cac_results, 1)

            self.assertEqual(results, cac_results)
        finally:
            parsers.SINGLE_FILE_PARSER_FACTORY.Unregister("TestFile")

    def testAggregatedArtifact(self):
        """Test we can collect an ARTIFACT_GROUP."""

        client_test_lib.Command("/bin/echo", args=["1"])

        artifact_list = ["TestArtifactGroup"]

        self.InitializeTestFileArtifact()

        results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                artifact_collector.ArtifactCollector,
                                artifact_list,
                                apply_parsers=False)
        self.assertLen(results, 2)

        artifact_response = results[0]
        self.assertIsInstance(artifact_response, rdf_client_fs.StatEntry)

        artifact_response = results[1]
        self.assertIsInstance(artifact_response,
                              rdf_client_action.ExecuteResponse)
        self.assertEqual(artifact_response.stdout, "1\n")

    def testArtifactFiles(self):
        """Test collecting an ArtifactFiles artifact."""

        artifact_list = ["TestArtifactFilesArtifact"]

        self.InitializeTestFileArtifact()

        # Run the ArtifactCollector to get the expected result.
        session_id = flow_test_lib.TestFlowHelper(
            aff4_flows.ArtifactCollectorFlow.__name__,
            action_mocks.FileFinderClientMock(),
            artifact_list=artifact_list,
            token=self.token,
            apply_parsers=False,
            client_id=self.client_id)
        results = flow_test_lib.GetFlowResults(self.client_id, session_id)
        expected = results[0]

        self.assertIsInstance(expected, rdf_client_fs.StatEntry)

        # Run the ClientArtifactCollector to get the actual result.
        cac_results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                    artifact_collector.ArtifactCollector,
                                    artifact_list,
                                    apply_parsers=False)
        self.assertLen(cac_results, 1)
        artifact_response = cac_results[0]
        self.assertEqual(artifact_response.pathspec.path,
                         expected.pathspec.path)

    def testArtifactFilesWithPathspecAttribute(self):
        """Test collecting ArtifactFiles with specified pathspec attribute."""

        artifact_list = ["TestArtifactFilesArtifact"]

        self.InitializeTestFileArtifact(with_pathspec_attribute=True)

        # Run the ArtifactCollector to get the expected result.
        session_id = flow_test_lib.TestFlowHelper(
            aff4_flows.ArtifactCollectorFlow.__name__,
            action_mocks.FileFinderClientMock(),
            artifact_list=artifact_list,
            token=self.token,
            apply_parsers=False,
            client_id=self.client_id)
        results = flow_test_lib.GetFlowResults(self.client_id, session_id)
        expected = results[0]

        self.assertIsInstance(expected, rdf_client_fs.StatEntry)

        # Run the ClientArtifactCollector to get the actual result.
        cac_results = self._RunFlow(aff4_flows.ClientArtifactCollector,
                                    artifact_collector.ArtifactCollector,
                                    artifact_list,
                                    apply_parsers=False)
        self.assertLen(cac_results, 1)
        artifact_response = cac_results[0]

        self.assertEqual(artifact_response.pathspec.path,
                         expected.pathspec.path)
Esempio n. 14
0
class ArtifactParserFactoryTest(absltest.TestCase):

  @mock.patch.object(parsers, "SINGLE_RESPONSE_PARSER_FACTORY",
                     factory.Factory(parsers.SingleResponseParser))
  def testSingleResponseParsers(self):

    class FooParser(parsers.SingleResponseParser[None]):

      supported_artifacts = ["Quux", "Norf"]

      def ParseResponse(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          response: rdfvalue.RDFValue,
      ) -> Iterator[None]:
        raise NotImplementedError()

    class BarParser(parsers.SingleResponseParser[None]):

      supported_artifacts = ["Norf", "Thud"]

      def ParseResponse(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          response: rdfvalue.RDFValue,
      ) -> Iterator[None]:
        raise NotImplementedError()

    class BazParser(parsers.SingleResponseParser[None]):

      supported_artifacts = ["Thud", "Quux"]

      def ParseResponse(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          response: rdfvalue.RDFValue,
      ) -> Iterator[None]:
        raise NotImplementedError()

    parsers.SINGLE_RESPONSE_PARSER_FACTORY.Register("Foo", FooParser)
    parsers.SINGLE_RESPONSE_PARSER_FACTORY.Register("Bar", BarParser)
    parsers.SINGLE_RESPONSE_PARSER_FACTORY.Register("Baz", BazParser)

    quux_factory = parsers.ArtifactParserFactory("Quux")
    quux_parsers = quux_factory.SingleResponseParsers()
    self.assertCountEqual(map(type, quux_parsers), [FooParser, BazParser])

    norf_factory = parsers.ArtifactParserFactory("Norf")
    norf_parsers = norf_factory.SingleResponseParsers()
    self.assertCountEqual(map(type, norf_parsers), [FooParser, BarParser])

    thud_factory = parsers.ArtifactParserFactory("Thud")
    thud_parsers = thud_factory.SingleResponseParsers()
    self.assertCountEqual(map(type, thud_parsers), [BarParser, BazParser])

  @mock.patch.object(parsers, "MULTI_RESPONSE_PARSER_FACTORY",
                     factory.Factory(parsers.MultiResponseParser))
  def testMultiResponseParsers(self):

    class FooParser(parsers.MultiResponseParser[None]):

      supported_artifacts = ["Foo"]

      def ParseResponses(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          responses: Iterable[rdfvalue.RDFValue],
      ) -> Iterator[None]:
        raise NotImplementedError()

    class BarParser(parsers.MultiResponseParser[None]):

      supported_artifacts = ["Bar"]

      def ParseResponses(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          responses: Iterable[rdfvalue.RDFValue],
      ) -> Iterator[None]:
        raise NotImplementedError()

    parsers.MULTI_RESPONSE_PARSER_FACTORY.Register("Foo", FooParser)
    parsers.MULTI_RESPONSE_PARSER_FACTORY.Register("Bar", BarParser)

    foo_factory = parsers.ArtifactParserFactory("Foo")
    foo_parsers = foo_factory.MultiResponseParsers()
    self.assertCountEqual(map(type, foo_parsers), [FooParser])

    bar_factory = parsers.ArtifactParserFactory("Bar")
    bar_parsers = bar_factory.MultiResponseParsers()
    self.assertCountEqual(map(type, bar_parsers), [BarParser])

  @mock.patch.object(parsers, "SINGLE_FILE_PARSER_FACTORY",
                     factory.Factory(parsers.SingleFileParser))
  def testSingleFileParsers(self):

    class FooParser(parsers.SingleFileParser):

      supported_artifacts = ["Bar"]

      def ParseFile(self, knowledge_base, pathspec, filedesc):
        raise NotImplementedError()

    parsers.SINGLE_FILE_PARSER_FACTORY.Register("Foo", FooParser)

    bar_factory = parsers.ArtifactParserFactory("Bar")
    bar_parsers = bar_factory.SingleFileParsers()
    self.assertCountEqual(map(type, bar_parsers), [FooParser])

    baz_factory = parsers.ArtifactParserFactory("Baz")
    baz_parsers = baz_factory.SingleFileParsers()
    self.assertCountEqual(map(type, baz_parsers), [])

  @mock.patch.object(parsers, "MULTI_FILE_PARSER_FACTORY",
                     factory.Factory(parsers.MultiFileParser))
  def testMultiFileParsers(self):

    class FooParser(parsers.MultiFileParser[None]):

      supported_artifacts = ["Quux", "Norf"]

      def ParseFiles(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          pathspecs: Iterable[rdf_paths.PathSpec],
          filedescs: Iterable[IO[bytes]],
      ) -> Iterator[None]:
        raise NotImplementedError()

    class BarParser(parsers.MultiFileParser[None]):

      supported_artifacts = ["Quux", "Thud"]

      def ParseFiles(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          pathspecs: Iterable[rdf_paths.PathSpec],
          filedescs: Iterable[IO[bytes]],
      ) -> Iterator[None]:
        raise NotImplementedError()

    parsers.MULTI_FILE_PARSER_FACTORY.Register("Foo", FooParser)
    parsers.MULTI_FILE_PARSER_FACTORY.Register("Bar", BarParser)

    quux_factory = parsers.ArtifactParserFactory("Quux")
    quux_parsers = quux_factory.MultiFileParsers()
    self.assertCountEqual(map(type, quux_parsers), [FooParser, BarParser])

    norf_factory = parsers.ArtifactParserFactory("Norf")
    norf_parsers = norf_factory.MultiFileParsers()
    self.assertCountEqual(map(type, norf_parsers), [FooParser])

    thud_factory = parsers.ArtifactParserFactory("Thud")
    thud_parsers = thud_factory.MultiFileParsers()
    self.assertCountEqual(map(type, thud_parsers), [BarParser])

  @mock.patch.object(parsers, "SINGLE_FILE_PARSER_FACTORY",
                     factory.Factory(parsers.SingleFileParser))
  @mock.patch.object(parsers, "MULTI_RESPONSE_PARSER_FACTORY",
                     factory.Factory(parsers.MultiResponseParser))
  def testAllParsers(self):

    class FooParser(parsers.SingleFileParser[None]):

      supported_artifacts = ["Quux"]

      def ParseFile(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          pathspec: rdf_paths.PathSpec,
          filedesc: IO[bytes],
      ):
        raise NotImplementedError()

    class BarParser(parsers.MultiResponseParser[None]):

      supported_artifacts = ["Quux"]

      def ParseResponses(
          self,
          knowledge_base: rdf_client.KnowledgeBase,
          responses: Iterable[rdfvalue.RDFValue],
      ) -> Iterator[None]:
        raise NotImplementedError()

    parsers.SINGLE_FILE_PARSER_FACTORY.Register("Foo", FooParser)
    parsers.MULTI_RESPONSE_PARSER_FACTORY.Register("Bar", BarParser)

    quux_factory = parsers.ArtifactParserFactory("Quux")
    quux_parsers = quux_factory.AllParserTypes()
    self.assertCountEqual(quux_parsers, [FooParser, BarParser])
Esempio n. 15
0
  def testCreateUnregisteredThrows(self):
    int_factory = factory.Factory(int)

    with self.assertRaisesRegex(ValueError, "foo"):
      int_factory.Create("foo")
Esempio n. 16
0
  def testGetAllTypesWithoutResults(self):
    obj_factory = factory.Factory(object)

    self.assertCountEqual(list(obj_factory.GetTypes()), [])
Esempio n. 17
0
 def testUsesClassConstructor(self):
   str_factory = factory.Factory(str)
   str_factory.Register("foo", str)
   self.assertEqual(str_factory.Create("foo"), "")
Esempio n. 18
0
    def testCreateWrongTypeThrows(self):
        int_factory = factory.Factory(int)
        int_factory.Register("foo", lambda: "Foo")

        with self.assertRaises(TypeError):
            int_factory.Create("foo")
Esempio n. 19
0
  def testUnregisterThrowsForUnknown(self):
    obj_factory = factory.Factory(object)

    with self.assertRaisesRegex(ValueError, "foo"):
      obj_factory.Unregister("foo")
Esempio n. 20
0
    def testCreateAllEmpty(self):
        obj_factory = factory.Factory(object)

        self.assertCountEqual(list(obj_factory.CreateAll()), [])
Esempio n. 21
0
#!/usr/bin/env python
"""A root module with decoder definitions."""
from __future__ import absolute_import
from __future__ import unicode_literals

from grr_response_core.lib import factory
from grr_response_server.decoders import _abstract

AbstractDecoder = _abstract.AbstractDecoder  # pylint: disable=invalid-name

FACTORY = factory.Factory(AbstractDecoder)