Esempio n. 1
0
    def testFlowReportsErrorWhenCollectingFileAboveSingleLimit(self):
        with temp.AutoTempFilePath() as temp_file_path:
            target_bytes = 2**20  # 1 MiB
            less_than_necessary_bytes = target_bytes // 2

            with io.open(temp_file_path, "wb") as fd:
                fd.write(b"1" * target_bytes)

            table = f"""
      [
        {{ "collect_column": "{temp_file_path}" }}
      ]
      """

            with mock.patch.object(osquery_flow,
                                   "FILE_COLLECTION_MAX_SINGLE_FILE_BYTES",
                                   less_than_necessary_bytes):
                with osquery_test_lib.FakeOsqueryiOutput(stdout=table,
                                                         stderr=""):
                    flow_id = self._InitializeFlow("Doesn't matter",
                                                   ["collect_column"],
                                                   check_flow_errors=False)
                    progress = flow_test_lib.GetFlowProgress(
                        self.client_id, flow_id)

            self.assertEqual(
                f"File with path {temp_file_path} is too big: "
                f"{target_bytes} bytes when the limit is "
                f"{less_than_necessary_bytes} bytes.", progress.error_message)
Esempio n. 2
0
    def testFlowReportsErrorWhenCollectingRowsAboveLimit(self):
        with temp.AutoTempFilePath() as temp_file_path:
            with io.open(temp_file_path, mode="w", encoding="utf-8") as fd:
                fd.write("Just sample text to put in the file.")

            table = f"""
      [
        {{ "collect_column": "{temp_file_path}"}},
        {{ "collect_column": "{temp_file_path}"}}
      ]
      """

            with mock.patch.object(osquery_flow, "FILE_COLLECTION_MAX_ROWS",
                                   1):
                with osquery_test_lib.FakeOsqueryiOutput(stdout=table,
                                                         stderr=""):
                    flow_id = self._InitializeFlow(
                        file_collection_columns=["collect_column"],
                        check_flow_errors=False)
                    progress = flow_test_lib.GetFlowProgress(
                        self.client_id, flow_id)

        self.assertEqual(
            "Requested file collection on a table with 2 rows, "
            "but the limit is 1 rows.", progress.error_message)
Esempio n. 3
0
    def testFlowReportsErrorWhenCollectingSingleFileAboveTotalLimit(self):
        with temp.AutoTempFilePath() as temp_file_path:
            target_bytes = 2**20  # 1 MiB
            less_than_necessary_bytes = target_bytes // 2

            with io.open(temp_file_path, "wb") as fd:
                fd.write(b"1" * target_bytes)

            table = f"""
      [
        {{ "collect_column": "{temp_file_path}" }}
      ]
      """

            with mock.patch.object(osquery_flow,
                                   "FILE_COLLECTION_MAX_TOTAL_BYTES",
                                   less_than_necessary_bytes):
                with osquery_test_lib.FakeOsqueryiOutput(stdout=table,
                                                         stderr=""):
                    flow_id = self._InitializeFlow(
                        file_collection_columns=["collect_column"],
                        check_flow_errors=False)
                    progress = flow_test_lib.GetFlowProgress(
                        self.client_id, flow_id)

        self.assertEqual(
            "Files for collection exceed the total size limit of "
            f"{less_than_necessary_bytes} bytes.", progress.error_message)
Esempio n. 4
0
  def testFlowDoesntFailWhenCollectingFilesFromEmptyResult(self):
    empty_table = """
    [

    ]
    """

    with osquery_test_lib.FakeOsqueryiOutput(stdout=empty_table, stderr=""):
      results = self._RunFlow(file_collection_columns=["collect_column"])

    self.assertEmpty(results)
Esempio n. 5
0
    def testFlowDoesntFailWhenCollectingFilesFromEmptyResult(self):
        empty_table = """
    [

    ]
    """

        with osquery_test_lib.FakeOsqueryiOutput(stdout=empty_table,
                                                 stderr=""):
            results = self._RunFlow("Doesn't matter", ["collect_column"])

        self.assertLen(results, 1)
        self.assertIsInstance(results[0], rdf_osquery.OsqueryResult)
Esempio n. 6
0
    def testFlowFailsWhenCollectingUnexistingColumn(self):
        with temp.AutoTempFilePath() as temp_file_path:
            with io.open(temp_file_path, mode="w", encoding="utf-8") as fd:
                fd.write("Just sample text to put in the file.")

            table = f"""
      [
        {{ "collect_column": "{temp_file_path}" }}
      ]
      """

            with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""):
                with self.assertRaises(RuntimeError):
                    self._RunFlow("Doesn't matter", ["non_existent_column"])
Esempio n. 7
0
  def testSmallerTruncationLimit(self):
    two_row_table = """
    [
      { "col1": "cell-1-1", "col2": "cell-1-2", "col3": "cell-1-3" },
      { "col1": "cell-2-1", "col2": "cell-2-2", "col3": "cell-2-3" }
    ]
    """
    max_rows = 1

    with osquery_test_lib.FakeOsqueryiOutput(stdout=two_row_table, stderr=""):
      with mock.patch.object(osquery_flow, "TRUNCATED_ROW_COUNT", max_rows):
        flow_id = self._InitializeFlow()
        progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id)

    self.assertEqual(len(progress.partial_table.rows), max_rows)
Esempio n. 8
0
    def _RunOsqueryExportResults(self,
                                 stdout: str) -> utils.BinaryChunkIterator:
        client_id = self.SetupClient(0)

        with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=""):
            flow_id = flow_test_lib.TestFlowHelper(
                osquery_flow.OsqueryFlow.__name__,
                action_mocks.OsqueryClientMock(),
                client_id=client_id,
                token=self.token,
                query="doesn't matter")
            result_flow = self.api.Client(client_id=client_id).Flow(flow_id)
            result_flow.WaitUntilDone()

        format_csv = api_osquery_pb2.ApiGetOsqueryResultsArgs.Format.CSV
        return result_flow.GetOsqueryResults(format_csv)
Esempio n. 9
0
  def testFlowDoesntCollectWhenRowsAboveLimit(self):
    with temp.AutoTempFilePath() as temp_file_path:
      with io.open(temp_file_path, mode="w", encoding="utf-8") as fd:
        fd.write("Just sample text to put in the file.")

      table = f"""
      [
        {{ "collect_column": "{temp_file_path}"}},
        {{ "collect_column": "{temp_file_path}"}}
      ]
      """

      with mock.patch.object(osquery_flow, "FILE_COLLECTION_MAX_ROWS", 1):
        with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""):
          with self.assertRaises(RuntimeError):
            self._RunFlow(file_collection_columns=["collect_column"])
Esempio n. 10
0
  def testTotalRowCountIncludesAllChunks(self):
    row_count = 100
    split_pieces = 10

    cell_value = "fixed"
    table = [{"column1": cell_value}] * row_count
    table_json = json.dumps(table)

    table_bytes = row_count * len(cell_value.encode("utf-8"))
    chunk_bytes = table_bytes // split_pieces

    with test_lib.ConfigOverrider({"Osquery.max_chunk_size": chunk_bytes}):
      with osquery_test_lib.FakeOsqueryiOutput(stdout=table_json, stderr=""):
        flow_id = self._InitializeFlow()
        progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id)

    self.assertEqual(progress.total_row_count, row_count)
Esempio n. 11
0
    def testOutput(self):
        stdout = """
    [
      { "foo": "bar", "quux": "norf" },
      { "foo": "baz", "quux": "thud" }
    ]
    """
        with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=""):
            results = _Query("SELECT foo, quux FROM blargh;")

        self.assertLen(results, 1)

        table = results[0].table
        self.assertLen(table.header.columns, 2)
        self.assertEqual(table.header.columns[0].name, "foo")
        self.assertEqual(table.header.columns[1].name, "quux")
        self.assertEqual(list(table.Column("foo")), ["bar", "baz"])
        self.assertEqual(list(table.Column("quux")), ["norf", "thud"])
Esempio n. 12
0
  def testChunksSmallerThanTruncation(self):
    row_count = 100
    max_rows = 70
    split_pieces = 10

    cell_value = "fixed"
    table = [{"column1": cell_value}] * row_count
    table_json = json.dumps(table)

    table_bytes = row_count * len(cell_value.encode("utf-8"))
    chunk_bytes = table_bytes // split_pieces

    with test_lib.ConfigOverrider({"Osquery.max_chunk_size": chunk_bytes}):
      with osquery_test_lib.FakeOsqueryiOutput(stdout=table_json, stderr=""):
        with mock.patch.object(osquery_flow, "TRUNCATED_ROW_COUNT", max_rows):
          flow_id = self._InitializeFlow()
          progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id)

    self.assertEqual(len(progress.partial_table.rows), max_rows)
Esempio n. 13
0
  def testFlowDoesntCollectSingleFileAboveTotalLimit(self):
    with temp.AutoTempFilePath() as temp_file_path:
      target_bytes = 2**20  # 1 MiB
      less_than_necessary_bytes = target_bytes // 2

      with io.open(temp_file_path, "wb") as fd:
        fd.write(b"1" * target_bytes)

      table = f"""
      [
        {{ "collect_column": "{temp_file_path}" }}
      ]
      """

      with mock.patch.object(osquery_flow, "FILE_COLLECTION_MAX_TOTAL_BYTES",
                             less_than_necessary_bytes):
        with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""):
          with self.assertRaises(RuntimeError):
            self._RunFlow(file_collection_columns=["collect_column"])
Esempio n. 14
0
  def testFlowReportsErrorWhenCollectingUnexistingColumn(self):
    with temp.AutoTempFilePath() as temp_file_path:
      with io.open(temp_file_path, mode="w", encoding="utf-8") as fd:
        fd.write("Just sample text to put in the file.")

      table = f"""
      [
        {{ "collect_column": "{temp_file_path}" }}
      ]
      """

      with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""):
        flow_id = self._InitializeFlow(
            file_collection_columns=["non_existent_column"],
            check_flow_errors=False)
        progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id)

    self.assertEqual(
        "No such column 'non_existent_column' to collect files from.",
        progress.error_message)
Esempio n. 15
0
    def testOsquery(self):
        data_store.REL_DB.WriteClientMetadata(
            client_id=ClientTest.FAKE_CLIENT_ID, fleetspeak_enabled=False)

        client = grr_colab.Client.with_id(ClientTest.FAKE_CLIENT_ID)

        stdout = """
    [
      { "foo": "test1", "bar": "test2" },
      { "foo": "test3", "bar": "test4" }
    ]
    """
        with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=''):
            table = client.osquery('SELECT foo, bar FROM table;')

        self.assertLen(table.header.columns, 2)
        self.assertEqual(table.header.columns[0].name, 'foo')
        self.assertEqual(table.header.columns[1].name, 'bar')
        self.assertEqual(list(list(table.rows)[0].values), ['test1', 'test2'])
        self.assertEqual(list(list(table.rows)[1].values), ['test3', 'test4'])
Esempio n. 16
0
    def testFailure(self):
        stderr = "Error: query syntax error"

        with osquery_test_lib.FakeOsqueryiOutput(stdout="", stderr=stderr):
            flow_id = flow_test_lib.StartFlow(
                flow_cls=osquery_flow.OsqueryFlow,
                client_id=self.client_id,
                creator=self.test_username,
                query="<<<FAKE OSQUERY FLOW QUERY PLACEHOLDER>>>")

            with self.assertRaises(RuntimeError):
                flow_test_lib.RunFlow(
                    client_id=self.client_id,
                    flow_id=flow_id,
                    client_mock=action_mocks.OsqueryClientMock())

        flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
        self.assertEqual(flow.flow_state,
                         rdf_flow_objects.Flow.FlowState.ERROR)
        self.assertIn(stderr, flow.error_message)
Esempio n. 17
0
    def testSuccess(self):
        stdout = """
    [
      { "foo": "quux", "bar": "norf", "baz": "thud" },
      { "foo": "blargh", "bar": "plugh", "baz": "ztesch" }
    ]
    """
        with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=""):
            results = self._RunFlow()

        self.assertLen(results, 1)

        table = results[0].table
        self.assertLen(table.header.columns, 3)
        self.assertEqual(table.header.columns[0].name, "foo")
        self.assertEqual(table.header.columns[1].name, "bar")
        self.assertEqual(table.header.columns[2].name, "baz")
        self.assertEqual(list(table.Column("foo")), ["quux", "blargh"])
        self.assertEqual(list(table.Column("bar")), ["norf", "plugh"])
        self.assertEqual(list(table.Column("baz")), ["thud", "ztesch"])
Esempio n. 18
0
    def testArchiveMappingsForDuplicateFilesInResult(self):
        with temp.AutoTempFilePath() as temp_file_path:
            with io.open(temp_file_path, mode="w", encoding="utf-8") as fd:
                fd.write("Just sample text to put in the file.")

            table = f"""
      [
        {{ "collect_column": "{temp_file_path}" }}
      ]
      """

            with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""):
                flow_id = self._InitializeFlow(
                    file_collection_columns=["collect_column"])

        flow = flow_base.FlowBase.CreateFlowInstance(
            flow_test_lib.GetFlowObj(self.client_id, flow_id))
        results = list(flow_test_lib.GetRawFlowResults(self.client_id,
                                                       flow_id))

        # This is how we emulate duplicate filenames in the results
        duplicated_results = results + results + results

        mappings = list(flow.GetFilesArchiveMappings(iter(duplicated_results)))
        self.assertCountEqual(mappings, [
            flow_base.ClientPathArchiveMapping(
                db.ClientPath.OS(self.client_id,
                                 temp_file_path.split("/")[1:]),
                f"osquery_collected_files{temp_file_path}",
            ),
            flow_base.ClientPathArchiveMapping(
                db.ClientPath.OS(self.client_id,
                                 temp_file_path.split("/")[1:]),
                f"osquery_collected_files{temp_file_path}-1",
            ),
            flow_base.ClientPathArchiveMapping(
                db.ClientPath.OS(self.client_id,
                                 temp_file_path.split("/")[1:]),
                f"osquery_collected_files{temp_file_path}-2",
            ),
        ])
Esempio n. 19
0
    def testIgnoreStderrErrors(self):
        query = "SELECT foo, bar, baz FROM blargh;"
        stdout = """
    [
      { "foo": "quux", "bar": "norf", "baz": "thud" }
    ]
    """
        stderr = "Warning: near 'FROM': table 'blargh' might be very large"
        with osquery_test_lib.FakeOsqueryiOutput(stdout=stdout, stderr=stderr):
            results = _Query(query, ignore_stderr_errors=True)

        self.assertLen(results, 1)
        self.assertEqual(results[0].stderr, stderr)

        table = results[0].table
        self.assertLen(table.header.columns, 3)
        self.assertEqual(table.header.columns[0].name, "foo")
        self.assertEqual(table.header.columns[1].name, "bar")
        self.assertEqual(table.header.columns[2].name, "baz")
        self.assertEqual(list(table.Column("foo")), ["quux"])
        self.assertEqual(list(table.Column("bar")), ["norf"])
        self.assertEqual(list(table.Column("baz")), ["thud"])
Esempio n. 20
0
    def testArchiveMappingsForMultipleFiles(self):
        with temp.AutoTempDirPath(remove_non_empty=True) as temp_dir_path:
            temp_file_path1 = os.path.join(temp_dir_path, "foo")
            temp_file_path2 = os.path.join(temp_dir_path, "bar")

            with io.open(temp_file_path1, mode="w", encoding="utf-8") as fd:
                fd.write("Just sample text to put in the file 1.")
            with io.open(temp_file_path2, mode="w", encoding="utf-8") as fd:
                fd.write("Just sample text to put in the file 2.")

            table = f"""
      [
        {{ "collect_column": "{temp_file_path1}" }},
        {{ "collect_column": "{temp_file_path2}" }}
      ]
      """

            with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""):
                flow_id = self._InitializeFlow(
                    file_collection_columns=["collect_column"])

        flow = flow_base.FlowBase.CreateFlowInstance(
            flow_test_lib.GetFlowObj(self.client_id, flow_id))
        results = flow_test_lib.GetRawFlowResults(self.client_id, flow_id)

        mappings = list(flow.GetFilesArchiveMappings(iter(results)))
        self.assertCountEqual(mappings, [
            flow_base.ClientPathArchiveMapping(
                db.ClientPath.OS(self.client_id,
                                 temp_file_path1.split("/")[1:]),
                f"osquery_collected_files{temp_file_path1}",
            ),
            flow_base.ClientPathArchiveMapping(
                db.ClientPath.OS(self.client_id,
                                 temp_file_path2.split("/")[1:]),
                f"osquery_collected_files{temp_file_path2}",
            ),
        ])
Esempio n. 21
0
    def testFlowCollectFile(self):
        with temp.AutoTempFilePath() as temp_file_path:
            with io.open(temp_file_path, mode="w", encoding="utf-8") as fd:
                fd.write("Just sample text to put in the file.")

            table = f"""
      [
        {{ "collect_column": "{temp_file_path}" }}
      ]
      """

            with osquery_test_lib.FakeOsqueryiOutput(stdout=table, stderr=""):
                results = self._RunFlow("Doesn't matter", ["collect_column"])

        self.assertLen(results, 2)
        self.assertIsInstance(results[0], rdf_osquery.OsqueryResult)
        self.assertIsInstance(results[1], rdf_client_fs.StatEntry)

        pathspec = results[1].pathspec
        client_path = db.ClientPath.FromPathSpec(self.client_id, pathspec)
        fd_rel_db = file_store.OpenFile(client_path)
        file_text = fd_rel_db.read().decode("utf-8")
        self.assertEqual(file_text, "Just sample text to put in the file.")
Esempio n. 22
0
 def testFailure(self):
   stderr = "Error: near '*': syntax error"
   with osquery_test_lib.FakeOsqueryiOutput(stdout="", stderr=stderr):
     with self.assertRaises(RuntimeError):
       self._RunQuery("SELECT * FROM *;")
Esempio n. 23
0
    def testSkipEmptyTable(self):
        with osquery_test_lib.FakeOsqueryiOutput(stdout="[]", stderr=""):
            results = self._RunFlow()

        self.assertEmpty(results)
Esempio n. 24
0
 def testError(self):
     stderr = "Error: near 'FROM': syntax error"
     with osquery_test_lib.FakeOsqueryiOutput(stdout="", stderr=stderr):
         with self.assertRaises(osquery.QueryError):
             _Query("FROM bar SELECT foo;")