Update TimesketchSearchEventCollector to use the new TimesketchEvents…

… container (#955) * Add TimesketchEvents to containers.py * Update containers.py * Update containers.py Fix pandas -> pd * Update containers.py fix import order * Update TimesketchSearchEventCollector to use the new TimesketchEvents container
log2timeline · Jan 16, 2025 · a7d3bf0 · a7d3bf0
1 parent defb3b9
commit a7d3bf0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/dftimewolf/lib/collectors/timesketch.py b/dftimewolf/lib/collectors/timesketch.py
@@ -233,10 +233,13 @@ def _OutputSearchResults(self, data_frame: pd.DataFrame) -> None:
 
     if self.output_format == 'pandas':
       self.StoreContainer(
-          containers.DataFrame(
+          containers.TimesketchEvents(
               name=self.search_name,
               description=self.search_description,
-              data_frame=data_frame))
+              data_frame=data_frame,
+              query=self.query_string,
+              sketch_id=self.sketch_id
+          ))
     else:
       with tempfile.NamedTemporaryFile(
           mode='w',

diff --git a/tests/lib/collectors/timesketch.py b/tests/lib/collectors/timesketch.py
@@ -120,7 +120,7 @@ def testProcessPandas(self, mock_get_search_results, _mock_get_api_client):
         token_password='test_token')
     self._ProcessModule()
 
-    state_containers = self._module.GetContainers(containers.DataFrame)
+    state_containers = self._module.GetContainers(containers.TimesketchEvents)
     self.assertEqual(len(state_containers), 1)
     pd.testing.assert_frame_equal(
         state_containers[0].data_frame, pd.DataFrame([1, 2]))