diff --git a/data/recipes/grr_hunt_osquery.json b/data/recipes/grr_hunt_osquery.json index 84c46407d..0d29f3e2f 100644 --- a/data/recipes/grr_hunt_osquery.json +++ b/data/recipes/grr_hunt_osquery.json @@ -7,7 +7,11 @@ "name": "OsqueryCollector", "args": { "query": "@osquery_query", - "paths": "@osquery_paths" + "paths": "@osquery_paths", + "remote_configuration_path": "@remote_configuration_path", + "local_configuration_path": "@local_configuration_path", + "configuration_content": "@configuration_content", + "file_collection_columns": "@file_collection_columns" } },{ "wants": ["OsqueryCollector"], @@ -30,6 +34,10 @@ ["reason", "Reason for collection.", null], ["--osquery_query", "Osquery query to hunt for.", null], ["--osquery_paths", "Path(s) to text file containing one osquery query per line.", null], + ["--remote_configuration_path", "Path to a remote osquery configuration file on the GRR client.", ""], + ["--local_configuration_path", "Path to a local osquery configuration file.", ""], + ["--configuration_content", "Osquery configuration as a JSON string.", ""], + ["--file_collection_columns", "The file collection columns.", ""], ["--timeout_millis", "Osquery timeout in milliseconds", 300000, {"format": "regex", "regex": "^\\d+$"}], ["--ignore_stderr_errors", "Ignore osquery stderr errors", false], ["--approvers", "Emails for GRR approval request.", null], diff --git a/data/recipes/grr_osquery_flow.json b/data/recipes/grr_osquery_flow.json index e07d9b340..58531a6c5 100644 --- a/data/recipes/grr_osquery_flow.json +++ b/data/recipes/grr_osquery_flow.json @@ -7,7 +7,11 @@ "name": "OsqueryCollector", "args": { "query": "@osquery_query", - "paths": "@osquery_paths" + "paths": "@osquery_paths", + "remote_configuration_path": "@remote_configuration_path", + "local_configuration_path": "@local_configuration_path", + "configuration_content": "@configuration_content", + "file_collection_columns": "@file_collection_columns" } },{ "wants": ["OsqueryCollector"], @@ -31,6 +35,10 @@ ["hostnames", "Hostname(s) to collect the osquery flow from.", null, {"format": "grr_host", "comma_separated": true}], ["--osquery_query", "Osquery query to hunt for.", null], ["--osquery_paths", "Path(s) to text file containing one osquery query per line.", null], + ["--remote_configuration_path", "Path to a remote osquery configuration file on the GRR client.", ""], + ["--local_configuration_path", "Path to a local osquery configuration file.", ""], + ["--configuration_content", "Osquery configuration as a JSON string.", ""], + ["--file_collection_columns", "The file collection columns.", ""], ["--timeout_millis", "Osquery timeout in milliseconds", 300000, {"format": "regex", "regex": "^\\d+$"}], ["--ignore_stderr_errors", "Ignore osquery stderr errors", false], ["--directory", "Directory in which to export results.", null], diff --git a/dftimewolf/lib/collectors/grr_hosts.py b/dftimewolf/lib/collectors/grr_hosts.py index 64e57680d..dd0622fb4 100644 --- a/dftimewolf/lib/collectors/grr_hosts.py +++ b/dftimewolf/lib/collectors/grr_hosts.py @@ -1154,13 +1154,17 @@ def _ProcessQuery( client: the GRR Client. osquery_container: the OSQuery. """ - hunt_args = osquery_flows.OsqueryFlowArgs( - query=osquery_container.query, - timeout_millis=self.timeout_millis, - ignore_stderr_errors=self.ignore_stderr_errors) + flow_args = osquery_flows.OsqueryFlowArgs() + flow_args.query = osquery_container.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 + flow_args.configuration_path = osquery_container.configuration_path + flow_args.file_collection_columns.extend( + osquery_container.file_collection_columns) try: - flow_id = self._LaunchFlow(client, 'OsqueryFlow', hunt_args) + flow_id = self._LaunchFlow(client, 'OsqueryFlow', flow_args) self._AwaitFlow(client, flow_id) except DFTimewolfError as error: self.ModuleError( diff --git a/dftimewolf/lib/collectors/grr_hunt.py b/dftimewolf/lib/collectors/grr_hunt.py index 1e4b61b2c..517a8dfda 100644 --- a/dftimewolf/lib/collectors/grr_hunt.py +++ b/dftimewolf/lib/collectors/grr_hunt.py @@ -401,6 +401,10 @@ def Process(self) -> None: hunt_args.query = osquery_container.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 + hunt_args.configuration_path = osquery_container.configuration_path + hunt_args.file_collection_columns.extend( + osquery_container.file_collection_columns) self._CreateAndStartHunt('OsqueryFlow', hunt_args) diff --git a/dftimewolf/lib/collectors/osquery.py b/dftimewolf/lib/collectors/osquery.py index 6b554bff1..9399b91f6 100644 --- a/dftimewolf/lib/collectors/osquery.py +++ b/dftimewolf/lib/collectors/osquery.py @@ -12,11 +12,17 @@ _ALL_PLATFORMS = ['darwin', 'freebsd', 'linux', 'windows'] + class OsqueryCollector(module.BaseModule): - """Osquey query collector. + """Osquery collector that creates OsqueryQuery containers. Attributes: - osqueries (List[containers.OsqueryQuery]): list of osquery containers. + osqueries (List[containers.OsqueryQuery]): list of osquery containers. + configuration_path (str): the path to a configuration file on the + client. + configuration_content (str): the JSON configuration content. + file_collection_columns (List[str]): The list of file collection + columns. """ def __init__(self, @@ -34,6 +40,9 @@ def __init__(self, super(OsqueryCollector, self).__init__( state, name=name, critical=critical) self.osqueries: List[containers.OsqueryQuery] = [] + self.configuration_path: str = '' + self.configuration_content: str = '' + self.file_collection_columns: List[str] = [] def _ValidateOsquery(self, query: str) -> bool: """Validate Osquery query. @@ -89,24 +98,27 @@ def _LoadOsqueryPackToState(self, path: str) -> None: if 'platform' in query_pack: global_platform = self._ParsePlatforms(query_pack.get('platform')) - for num, (name, entry) in enumerate( - query_pack.get('queries', {}).items()): + for num, (name, entry) in enumerate(query_pack.get('queries', {}).items()): query = entry['query'] if not self._ValidateOsquery(query): - self.logger.warning(f'Entry {num} in query pack' - f'{path} does not appear to be valid.') + self.logger.warning( + f'Entry {num} in query pack {path} does not appear to be valid.') continue if 'platform' in entry: platform = self._ParsePlatforms(entry.get('platform')) else: platform = global_platform + self.osqueries.append( containers.OsqueryQuery( query=query, name=name, description=entry.get('description', ''), - platforms=platform)) + platforms=platform, + configuration_content=self.configuration_content, + configuration_path=self.configuration_path, + file_collection_columns=self.file_collection_columns)) def _LoadTextFileToState(self, path: str) -> None: """Loads osquery from a text file and creates Osquery containers. @@ -122,32 +134,89 @@ def _LoadTextFileToState(self, path: str) -> None: query=line, name='', description='', - platforms=None)) + platforms=None, + configuration_content=self.configuration_content, + 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.') # pylint: disable=arguments-differ - def SetUp(self, - query: str, - paths: str) -> None: + def SetUp( + self, + query: str, + paths: str, + remote_configuration_path: str = '', + local_configuration_path: str = '', + configuration_content: str = '', + file_collection_columns: Optional[str] = None + ) -> None: """Sets up the osquery to collect. Supported files are: * text files that contain one Osquery * json files containing an osquery pack. See https://osquery.readthedocs.io - /en/stable/deployment/configuration/#query-packs for details and - https://github.com/osquery/osquery/tree/master/packs for examples. + /en/stable/deployment/configuration/#query-packs for details and + https://github.com/osquery/osquery/tree/master/packs for examples. + + The GRR osquery flow can also be set up to use a custom osquery + configuration on invocation (see + https://osquery.readthedocs.io/en/stable/deployment/configuration/) + 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 + string value, using configuration_content. + + GRR can also collect files based on the results of an Osquery flow using the + file_collection_columns argument. Args: - query (str): osquery query. - paths (str): osquery filepaths. + query: osquery query. + paths: osquery filepaths. + remote_configuration_path: the path to a remote osquery configuration file + on the GRR client. + configuration_content: the configuration content, in JSON format. + local_configuration_path: the path to a local osquery configuration file. + file_collection_columns: The comma-seaparated list of file collection + columns names. """ if not query and not paths: self.ModuleError('Both query and paths cannot be empty.', critical=True) + if (remote_configuration_path and ( + local_configuration_path or configuration_content) or ( + local_configuration_path and configuration_content + )): + self.ModuleError( + 'Only one configuration argument can be set.', critical=True) + + if remote_configuration_path: + self.configuration_path = remote_configuration_path + elif local_configuration_path: + with open(local_configuration_path, mode='r') as fd: + configuration_content = fd.read() + + if configuration_content: + try: + content = json.loads(configuration_content) + except json.JSONDecodeError: + self.ModuleError( + 'Osquery configuration does not contain valid JSON.', + critical=True) + self.configuration_content = json.dumps(content) + + if file_collection_columns: + 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)) + 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.') diff --git a/dftimewolf/lib/containers/containers.py b/dftimewolf/lib/containers/containers.py index 0b270f4d1..0f004eae6 100644 --- a/dftimewolf/lib/containers/containers.py +++ b/dftimewolf/lib/containers/containers.py @@ -599,10 +599,16 @@ class OsqueryQuery(interface.AttributeContainer): Attributes: query (str): The osquery query. + configuration_content (str): The JSON content of an osquery + configuration. + configuration_path (str): The path to an osquery configuration + file on the client. name (Optional[str]): A name for the osquery. platforms (Optional[List[str]]): A constraint on the platform(s) the query should be run. Valid values are 'darwin', 'linux', 'windows', description (Optional[str]): A description for the query. + file_collection_columns (Optional[List[str]]): The list of file collection + columns. """ CONTAINER_TYPE = "osquery_query" @@ -610,15 +616,21 @@ class OsqueryQuery(interface.AttributeContainer): def __init__( self, query: str, + configuration_content: str = '', + configuration_path: str = '', name: Optional[str] = None, platforms: Optional[List[str]] = None, description: Optional[str] = None, + file_collection_columns: Optional[List[str]] = None ) -> None: super(OsqueryQuery, self).__init__() self.description = description self.name = name self.platforms = platforms self.query = query + self.configuration_content = configuration_content + self.configuration_path = configuration_path + self.file_collection_columns = file_collection_columns or [] def __str__(self) -> str: """Override __str()__.""" diff --git a/poetry.lock b/poetry.lock index 659ebe94b..968791237 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.6.1 and should not be changed by hand. [[package]] name = "adal" @@ -1445,28 +1445,28 @@ protobuf = ">=3.12.0" [[package]] name = "grr-api-client" -version = "3.4.7.post1" +version = "3.4.7.post5" description = "GRR API client library" optional = false python-versions = "*" files = [ - {file = "grr-api-client-3.4.7.post1.zip", hash = "sha256:e62f965abf940de71b9eb1a2589fbfe240b4d083f53f4e6533be9426baf91579"}, + {file = "grr-api-client-3.4.7.post5.zip", hash = "sha256:a633c621fbb63d7e613932898787963aec00d56fa6e248cb0a8cf8b37e3090d9"}, ] [package.dependencies] cryptography = ">=3.3.2" -grr_response_proto = "3.4.7post1" +grr_response_proto = "3.4.7post5" requests = ">=2.25.1,<3" Werkzeug = ">=2.1.2,<3" [[package]] name = "grr-response-proto" -version = "3.4.7.post1" +version = "3.4.7.post5" description = "GRR API client library" optional = false python-versions = "*" files = [ - {file = "grr-response-proto-3.4.7.post1.zip", hash = "sha256:84d6ff957df4721e873e57766ef3b210323a22ce1638ce94159733358c522590"}, + {file = "grr-response-proto-3.4.7.post5.zip", hash = "sha256:921bfab2b566002c39240eaec67384d41c195a3570fed75ffb0b3c7da8dd7689"}, ] [package.dependencies] @@ -2883,13 +2883,13 @@ rpds-py = ">=0.7.0" [[package]] name = "requests" -version = "2.32.0" +version = "2.31.0" description = "Python HTTP for Humans." optional = false -python-versions = ">=3.8" +python-versions = ">=3.7" files = [ - {file = "requests-2.32.0-py3-none-any.whl", hash = "sha256:f2c3881dddb70d056c5bd7600a4fae312b2a300e39be6a118d30b90bd27262b5"}, - {file = "requests-2.32.0.tar.gz", hash = "sha256:fa5490319474c82ef1d2c9bc459d3652e3ae4ef4c4ebdd18a21145a47ca4b6b8"}, + {file = "requests-2.31.0-py3-none-any.whl", hash = "sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f"}, + {file = "requests-2.31.0.tar.gz", hash = "sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1"}, ] [package.dependencies] @@ -3570,4 +3570,4 @@ turbinia-legacy = [] [metadata] lock-version = "2.0" python-versions = ">=3.11,<=3.13" -content-hash = "34d543a9f22230185531e9605bba253bee6c827793d837b8c7fdb5415bf4eaa2" +content-hash = "999edc2fbc2119868352dcede55f7a4b823d1ac146ba4d4e20177c588e594c76" diff --git a/pyproject.toml b/pyproject.toml index 925966ce5..1df780bcc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ boto3 = "^1.24.31" python-magic = "^0.4.27" turbinia-api-lib = "^1.0.2" turbinia-client = "^1.0.3" -grr-api-client = "^3.4.6" +grr-api-client = "^3.4.7" libcloudforensics = {git = "https://github.com/google/cloud-forensics-utils.git"} docker = "^7.0.0" setuptools = "^69.5.1" # needed by docker diff --git a/tests/lib/collectors/grr_hosts.py b/tests/lib/collectors/grr_hosts.py index 9253ac926..0dd8c19fa 100644 --- a/tests/lib/collectors/grr_hosts.py +++ b/tests/lib/collectors/grr_hosts.py @@ -951,7 +951,10 @@ def setUp(self, mock_InitHttp): self.grr_osquery_collector = grr_hosts.GRROsqueryCollector( self.test_state) self.grr_osquery_collector.StoreContainer( - containers.OsqueryQuery('SELECT * FROM processes')) + containers.OsqueryQuery( + 'SELECT * FROM processes', + configuration_path='/test/path', + file_collection_columns=['path'])) self.grr_osquery_collector.SetUp( hostnames='C.0000000000000001', reason='Random reason', @@ -993,7 +996,10 @@ def testProcess(self, mock_LaunchFlow, mock_DownloadResults, _): osquery_pb2.OsqueryFlowArgs( query='SELECT * FROM processes', timeout_millis=300000, - ignore_stderr_errors=False + ignore_stderr_errors=False, + configuration_content='', + configuration_path='/test/path', + file_collection_columns=['path'] ) ) diff --git a/tests/lib/collectors/grr_hunt.py b/tests/lib/collectors/grr_hunt.py index 75c83c482..d16109ad0 100644 --- a/tests/lib/collectors/grr_hunt.py +++ b/tests/lib/collectors/grr_hunt.py @@ -133,7 +133,10 @@ def setUp(self, mock_InitHttp): self.grr_hunt_osquery_collector = grr_hunt.GRRHuntOsqueryCollector( self.test_state) self.grr_hunt_osquery_collector.StoreContainer( - containers.OsqueryQuery(query='SELECT * FROM processes')) + containers.OsqueryQuery( + query='SELECT * FROM processes', + configuration_path='/test/path', + file_collection_columns=['path'])) self.grr_hunt_osquery_collector.SetUp( reason='random reason', timeout_millis=300000, @@ -157,6 +160,8 @@ def testProcess(self): self.assertEqual(call_kwargs['flow_args'].timeout_millis, 300000) self.assertEqual(call_kwargs['flow_args'].ignore_stderr_errors, False) + self.assertEqual(call_kwargs['flow_args'].configuration_path, '/test/path') + self.assertEqual(call_kwargs['flow_args'].file_collection_columns, ['path']) self.assertEqual(call_kwargs['flow_name'], 'OsqueryFlow') self.assertEqual(call_kwargs['hunt_runner_args'].description, 'random reason') diff --git a/tests/lib/collectors/osquery.py b/tests/lib/collectors/osquery.py index 1b86884ce..dc7ab6223 100644 --- a/tests/lib/collectors/osquery.py +++ b/tests/lib/collectors/osquery.py @@ -17,38 +17,35 @@ class OsqueryCollectorTest(unittest.TestCase): """Tests for the GRR osquery collector.""" - def testInitialization(self) -> None: - """Tests that the collector can be initialized.""" + def setUp(self) -> None: + super().setUp() test_state = state.DFTimewolfState(config.Config) - osquery_collector = osquery.OsqueryCollector(test_state) - self.assertIsNotNone(osquery_collector) + self.osquery_collector = osquery.OsqueryCollector(test_state) - def testSetup(self) -> None: - """Tests the collector's Setup() function.""" - test_state = state.DFTimewolfState(config.Config) - osquery_collector = osquery.OsqueryCollector(test_state) + def testInitialization(self) -> None: + """Tests that the collector can be initialized.""" + self.assertEqual(self.osquery_collector.osqueries, []) + self.assertEqual(self.osquery_collector.configuration_content, '') + self.assertEqual(self.osquery_collector.configuration_path, '') + self.assertEqual(self.osquery_collector.file_collection_columns, []) + def testSetupError(self) -> None: + """Tests the collector's Setup() function with invalid query, path.""" with self.assertRaises(DFTimewolfError) as context: - osquery_collector.SetUp(query='', paths='') + self.osquery_collector.SetUp(query='', paths='') self.assertEqual( context.exception.message, 'Both query and paths cannot be empty.') - def testSetupQuery(self) -> None: + def testSetupQueryError(self) -> None: """Tests the collector's Setup() function with invalid query parameter.""" - test_state = state.DFTimewolfState(config.Config) - osquery_collector = osquery.OsqueryCollector(test_state) - with self.assertRaises(DFTimewolfError) as context: - osquery_collector.SetUp(query='not a query', paths='') + self.osquery_collector.SetUp(query='not a query', paths='') self.assertEqual(context.exception.message, 'No valid osquery collected.') - def testSetupPaths(self) -> None: + def testSetupPathsError(self) -> None: """Tests the collector's Setup() method with invalid paths parameter.""" - test_state = state.DFTimewolfState(config.Config) - osquery_collector = osquery.OsqueryCollector(test_state) - test_empty_data = "" test_bad_data = "bad" @@ -56,7 +53,7 @@ def testSetupPaths(self) -> None: 'builtins.open', new=mock.mock_open(read_data=test_empty_data)) as _: with self.assertRaises(DFTimewolfError) as context: - osquery_collector.SetUp(query='', paths='empty') + self.osquery_collector.SetUp(query='', paths='empty') self.assertEqual(context.exception.message, 'No valid osquery collected.') @@ -64,15 +61,85 @@ def testSetupPaths(self) -> None: 'builtins.open', new=mock.mock_open(read_data=test_bad_data)) as _: with self.assertRaises(DFTimewolfError) as context: - osquery_collector.SetUp(query='', paths='fbad') + self.osquery_collector.SetUp(query='', paths='fbad') self.assertEqual(context.exception.message, 'No valid osquery collected.') + 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='', + 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='', + 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='', + 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='', + configuration_content='invalid content') + self.assertEqual( + context.exception.message, + 'Osquery configuration does not contain valid JSON.') + + def testSetUpRemoteConfigurationPath(self) -> None: + """Tests the collector's SetUp() function with the remote config path.""" + self.osquery_collector.SetUp( + query='SELECT * from test', + paths='ok', + remote_configuration_path='/test/path') + self.assertEqual(self.osquery_collector.configuration_path, '/test/path') + + def testSetUpLocalConfigurationPath(self) -> None: + """Tests the collector's SetUp() function with the local config path.""" + with mock.patch( + 'builtins.open', + new=mock.mock_open(read_data='{"test": "test"}')) as _: + self.osquery_collector.SetUp( + query='SELECT * from test', + paths='ok', + local_configuration_path='test') + self.assertEqual( + self.osquery_collector.configuration_content, '{"test": "test"}') + + def testSetUpConfigurationContent(self) -> None: + """Tests the collector's SetUp() function with configuration content.""" + self.osquery_collector.SetUp( + query='SELECT * from test', + paths='ok', + configuration_content='{"test": "test"}') + self.assertEqual( + self.osquery_collector.configuration_content, '{"test": "test"}') + + def testSetUpFileCollectionColumns(self) -> None: + """Tests the collector's SetUp() function with file collection columns.""" + self.osquery_collector.SetUp( + query='SELECT * from test', + paths='ok', + file_collection_columns='a,b') + self.assertEqual( + self.osquery_collector.file_collection_columns, ['a', 'b']) + @mock.patch('os.path.exists') def testProcessTextFile(self, mock_exists) -> None: """Tests the collector's Process() function with a text file.""" - test_state = state.DFTimewolfState(config.Config) - osquery_collector = osquery.OsqueryCollector(test_state) mock_exists.return_value = True test_ok_data = "SELECT * FROM processes" @@ -80,19 +147,19 @@ def testProcessTextFile(self, mock_exists) -> None: with mock.patch( 'builtins.open', new=mock.mock_open(read_data=test_ok_data)) as _: - osquery_collector.SetUp(query='', paths='ok') + self.osquery_collector.SetUp(query='', paths='ok') - osquery_collector.Process() + self.osquery_collector.Process() - containers = osquery_collector.GetContainers(OsqueryQuery) + containers = self.osquery_collector.GetContainers(OsqueryQuery) self.assertEqual(len(containers), 1) self.assertEqual(containers[0].query, "SELECT * FROM processes") + self.assertEqual(containers[0].configuration_content, '') + self.assertEqual(containers[0].configuration_path, '') @mock.patch('os.path.exists') def testProcessQueryPack(self, mock_exists) -> None: """Tests the collector's Process() function with a Osquery Pack file.""" - test_state = state.DFTimewolfState(config.Config) - osquery_collector = osquery.OsqueryCollector(test_state) mock_exists.return_value = True test_ok_data = json.dumps({ @@ -120,11 +187,11 @@ def testProcessQueryPack(self, mock_exists) -> None: with mock.patch( 'builtins.open', new=mock.mock_open(read_data=test_ok_data)) as _: - osquery_collector.SetUp(query='', paths='ok.json') + self.osquery_collector.SetUp(query='', paths='ok.json') - osquery_collector.Process() + self.osquery_collector.Process() - containers = osquery_collector.GetContainers(OsqueryQuery) + containers = self.osquery_collector.GetContainers(OsqueryQuery) self.assertEqual(len(containers), 2) self.assertEqual(containers[0].name, 'query_1') @@ -139,5 +206,6 @@ def testProcessQueryPack(self, mock_exists) -> None: self.assertEqual(containers[1].query, 'select * from test where path like \'%user32.dll\';') + if __name__ == '__main__': unittest.main() diff --git a/tests/lib/containers/containers.py b/tests/lib/containers/containers.py index 1488c4e1b..cfa63461b 100644 --- a/tests/lib/containers/containers.py +++ b/tests/lib/containers/containers.py @@ -108,10 +108,14 @@ class OsqueryQueryDataTest(unittest.TestCase): def testGetAttributeNames(self): """Tests the GetAttributeNames function.""" attribute_container = containers.OsqueryQuery( - query='', name='', description='', platforms=[]) + query='', name='', description='', platforms=[], + configuration_content='', configuration_path='', + file_collection_columns=[]) expected_attribute_names = [ - 'description', 'metadata', 'name', 'platforms', 'query'] + 'configuration_content', 'configuration_path', 'description', + 'file_collection_columns', 'metadata', 'name', 'platforms', 'query', + ] attribute_names = sorted(attribute_container.GetAttributeNames())