Add check for trailing semicolon to osquery query value (#916)

* Updates

* Add trailing colon...

* Fixes

* Update logic to fail

* Fix

dftimewolf/lib/collectors/grr_hosts.py

@@ -1240,8 +1240,12 @@ def _ProcessQuery(
       client: the GRR Client.
       osquery_container: the OSQuery.
     """
+    query = osquery_container.query
+    if not query.strip().endswith(';'):
+      query += ';'
+
     flow_args = osquery_flows.OsqueryFlowArgs()
-    flow_args.query = osquery_container.query
+    flow_args.query = query
     flow_args.timeout_millis = self.timeout_millis
     flow_args.ignore_stderr_errors = self.ignore_stderr_errors
     flow_args.configuration_content = osquery_container.configuration_content
@@ -1260,7 +1264,6 @@ def _ProcessQuery(
 
     name = osquery_container.name
     description = osquery_container.description
-    query = osquery_container.query
     flow_identifier = flow_id
     client_identifier = client.client_id
 

dftimewolf/lib/collectors/grr_hunt.py

@@ -397,8 +397,12 @@ def Process(self) -> None:
     osquery_containers = self.GetContainers(containers.OsqueryQuery)
 
     for osquery_container in osquery_containers:
+      query = osquery_container.query
+      if not query.strip().endswith(';'):
+        query += ';'
+
       hunt_args = osquery_flows.OsqueryFlowArgs()
-      hunt_args.query = osquery_container.query
+      hunt_args.query = query
       hunt_args.timeout_millis = self.timeout_millis
       hunt_args.ignore_stderr_errors = self.ignore_stderr_errors
       hunt_args.configuration_content = osquery_container.configuration_content

dftimewolf/lib/collectors/osquery.py

@@ -54,7 +54,9 @@ def _ValidateOsquery(self, query: str) -> bool:
       True if the query appears to be valid, False otherwise
     """
     # TODO(sydp): add more checks.
-    return query.upper().startswith('SELECT ')
+    return (
+        query.strip().upper().startswith('SELECT ')
+        and query.strip().endswith(';'))
 
   def _ParsePlatforms(self, platforms: str) -> List[str]:
     """Parse and normalise the platforms value from an osquery pack.
@@ -79,7 +81,8 @@ def _ParsePlatforms(self, platforms: str) -> List[str]:
       elif platform in _ALL_PLATFORMS:
         unique_platforms.add(platform)
       else:
-        self.logger.warning(f'Unexpected value {platform} in platform value.')
+        self.ModuleError(
+            f'Unexpected value {platform} in platform value.', critical=True)
 
     return list(unique_platforms)
 
@@ -101,7 +104,7 @@ def _LoadOsqueryPackToState(self, path: str) -> None:
     for num, (name, entry) in enumerate(query_pack.get('queries', {}).items()):
       query = entry['query']
       if not self._ValidateOsquery(query):
-        self.logger.warning(
+        self.ModuleError(
             f'Entry {num} in query pack {path} does not appear to be valid.')
         continue
 
@@ -139,8 +142,9 @@ def _LoadTextFileToState(self, path: str) -> None:
                   configuration_path=self.configuration_path,
                   file_collection_columns=self.file_collection_columns))
         else:
-          self.logger.warning(f'Osquery on line {line_number} of {path} '
-                              'does not appear to be valid.')
+          self.ModuleError(
+              f'Osquery on line {line_number} of {path} '
+              'does not appear to be valid.', critical=True)
 
   # pylint: disable=arguments-differ
   def SetUp(
@@ -166,9 +170,9 @@ def SetUp(
     either:
     * as an existing file on the GRR client using remote_configuration_path
     * as a temporary file on the GRR client where the content can come from
-    a file, using local_cofiguration_path, on the user's local machine or a 
+    a file, using local_cofiguration_path, on the user's local machine or a
     string value, using configuration_content.
-    
+
     GRR can also collect files based on the results of an Osquery flow using the
     file_collection_columns argument.
 
@@ -211,15 +215,17 @@ def SetUp(
       self.file_collection_columns = [
           col.strip() for col in file_collection_columns.split(',')]
 
-    if query and self._ValidateOsquery(query):
-      self.osqueries.append(containers.OsqueryQuery(
-          query=query,
-          configuration_content=self.configuration_content,
-          configuration_path=self.configuration_path,
-          file_collection_columns=self.file_collection_columns))
-    else:
-      self.logger.warning(
-          'Osquery parameter not set or does not appear to be valid.')
+    if query:
+      if self._ValidateOsquery(query):
+        self.osqueries.append(containers.OsqueryQuery(
+            query=query,
+            configuration_content=self.configuration_content,
+            configuration_path=self.configuration_path,
+            file_collection_columns=self.file_collection_columns))
+      else:
+        self.ModuleError(
+            'Osquery parameter not set or does not appear to be valid.',
+            critical=True)
 
     if paths:
       split_paths = [path.strip() for path in paths.split(',')]
@@ -235,7 +241,7 @@ def SetUp(
 
     if not self.osqueries:
       self.ModuleError(
-        message='No valid osquery collected.', critical=True)
+          message='No valid osquery collected.', critical=True)
 
   def Process(self) -> None:
     """Collects osquery from the command line and local file system."""

tests/lib/collectors/grr_hosts.py

@@ -980,7 +980,7 @@ def testProcess(self, mock_LaunchFlow, mock_DownloadResults, _):
         mock_grr_hosts.MOCK_CLIENT_RECENT,
         'OsqueryFlow',
         osquery_pb2.OsqueryFlowArgs(
-            query='SELECT * FROM processes',
+            query='SELECT * FROM processes;',
             timeout_millis=300000,
             ignore_stderr_errors=False,
             configuration_content='',
@@ -991,7 +991,7 @@ def testProcess(self, mock_LaunchFlow, mock_DownloadResults, _):
 
     results = self.grr_osquery_collector.GetContainers(containers.OsqueryResult)
     self.assertEqual(len(results), 1)
-    self.assertEqual(results[0].query, 'SELECT * FROM processes')
+    self.assertEqual(results[0].query, 'SELECT * FROM processes;')
     self.assertEqual(results[0].name, None)
     self.assertEqual(results[0].description, None)
     self.assertEqual(results[0].hostname, 'C.0000000000000001')

tests/lib/collectors/grr_hunt.py

@@ -156,7 +156,7 @@ def testProcess(self):
     # extract call kwargs
     call_kwargs = self.mock_grr_api.CreateHunt.call_args[1]
     self.assertEqual(call_kwargs['flow_args'].query,
-                     'SELECT * FROM processes')
+                     'SELECT * FROM processes;')
     self.assertEqual(call_kwargs['flow_args'].timeout_millis,
                      300000)
     self.assertEqual(call_kwargs['flow_args'].ignore_stderr_errors, False)

tests/lib/collectors/osquery.py

@@ -42,7 +42,16 @@ def testSetupQueryError(self) -> None:
     with self.assertRaises(DFTimewolfError) as context:
       self.osquery_collector.SetUp(query='not a query', paths='')
 
-    self.assertEqual(context.exception.message, 'No valid osquery collected.')
+    self.assertEqual(
+        context.exception.message,
+        'Osquery parameter not set or does not appear to be valid.')
+
+    with self.assertRaises(DFTimewolfError) as context:
+      self.osquery_collector.SetUp(query='SELECT * FROM processes', paths='')
+
+    self.assertEqual(
+        context.exception.message,
+        'Osquery parameter not set or does not appear to be valid.')
 
   def testSetupPathsError(self) -> None:
     """Tests the collector's Setup() method with invalid paths parameter."""
@@ -69,31 +78,31 @@ def testSetUpConfigurationError(self) -> None:
     """Tests the collector's SetUp() function with invalid configuration."""
     with self.assertRaises(DFTimewolfError) as context:
       self.osquery_collector.SetUp(
-          query='SELECT * FROM processes', paths='',
+          query='SELECT * FROM processes;', paths='',
           configuration_content='test', remote_configuration_path='test')
     self.assertEqual(
         context.exception.message,
         'Only one configuration argument can be set.')
 
     with self.assertRaises(DFTimewolfError) as context:
       self.osquery_collector.SetUp(
-          query='SELECT * FROM processes', paths='',
+          query='SELECT * FROM processes;', paths='',
           local_configuration_path ='test', remote_configuration_path='test')
     self.assertEqual(
         context.exception.message,
         'Only one configuration argument can be set.')
 
     with self.assertRaises(DFTimewolfError) as context:
       self.osquery_collector.SetUp(
-          query='SELECT * FROM processes', paths='',
+          query='SELECT * FROM processes;', paths='',
           local_configuration_path ='test', configuration_content='test')
     self.assertEqual(
         context.exception.message,
         'Only one configuration argument can be set.')
 
     with self.assertRaises(DFTimewolfError) as context:
       self.osquery_collector.SetUp(
-          query='SELECT * from processes', paths='',
+          query='SELECT * from processes;', paths='',
           configuration_content='invalid content')
     self.assertEqual(
         context.exception.message,
@@ -102,7 +111,7 @@ def testSetUpConfigurationError(self) -> None:
   def testSetUpRemoteConfigurationPath(self) -> None:
     """Tests the collector's SetUp() function with the remote config path."""
     self.osquery_collector.SetUp(
-        query='SELECT * from test',
+        query='SELECT * from test;',
         paths='ok',
         remote_configuration_path='/test/path')
     self.assertEqual(self.osquery_collector.configuration_path, '/test/path')
@@ -113,7 +122,7 @@ def testSetUpLocalConfigurationPath(self) -> None:
         'builtins.open',
         new=mock.mock_open(read_data='{"test": "test"}')) as _:
       self.osquery_collector.SetUp(
-          query='SELECT * from test',
+          query='SELECT * from test;',
           paths='ok',
           local_configuration_path='test')
     self.assertEqual(
@@ -122,7 +131,7 @@ def testSetUpLocalConfigurationPath(self) -> None:
   def testSetUpConfigurationContent(self) -> None:
     """Tests the collector's SetUp() function with configuration content."""
     self.osquery_collector.SetUp(
-        query='SELECT * from test',
+        query='SELECT * from test;',
         paths='ok',
         configuration_content='{"test": "test"}')
     self.assertEqual(
@@ -131,7 +140,7 @@ def testSetUpConfigurationContent(self) -> None:
   def testSetUpFileCollectionColumns(self) -> None:
     """Tests the collector's SetUp() function with file collection columns."""
     self.osquery_collector.SetUp(
-        query='SELECT * from test',
+        query='SELECT * from test;',
         paths='ok',
         file_collection_columns='a,b')
     self.assertEqual(
@@ -142,7 +151,7 @@ def testProcessTextFile(self, mock_exists) -> None:
     """Tests the collector's Process() function with a text file."""
     mock_exists.return_value = True
 
-    test_ok_data = "SELECT * FROM processes"
+    test_ok_data = "SELECT * FROM processes;"
 
     with mock.patch(
         'builtins.open',
@@ -153,7 +162,7 @@ def testProcessTextFile(self, mock_exists) -> None:
 
     containers = self.osquery_collector.GetContainers(OsqueryQuery)
     self.assertEqual(len(containers), 1)
-    self.assertEqual(containers[0].query, "SELECT * FROM processes")
+    self.assertEqual(containers[0].query, "SELECT * FROM processes;")
     self.assertEqual(containers[0].configuration_content, '')
     self.assertEqual(containers[0].configuration_path, '')