101 requesting automatic termination of search after x time or y results

products/cortex_xdr.py

@@ -55,14 +55,16 @@ class CortexXDR(Product):
     _session: requests.Session
     _queries: dict[Tag, list[Query]]
     _last_request: float
+    _limit: int = 1000 # Max is 1000 results otherwise have to get the results via stream
 
     def __init__(self, profile: str, creds_file: str, **kwargs):
         if not os.path.isfile(creds_file):
             raise ValueError(f'Credential file {creds_file} does not exist')
 
         self.creds_file = creds_file
         self._queries = dict()
-
+        if self._limit >= int(kwargs.get('limit',0)) > 0:
+            self._limit = int(kwargs['limit'])
         self._last_request = 0.0
 
         super().__init__(self.product, profile, **kwargs)
@@ -227,13 +229,12 @@ def nested_process_search(self, tag: Tag, criteria: dict, base_query: dict) -> N
         except KeyboardInterrupt:
             self._echo("Caught CTRL-C. Returning what we have...")
 
-    def _get_xql_results(self, query_id: str, limit: int = 1000) -> Tuple[dict, int]:
-        actual_limit = limit if limit < 1000 else 1000  # Max is 1000 results otherwise have to get the results via stream
+    def _get_xql_results(self, query_id: str) -> Tuple[dict, int]:
         params = {
             'request_data': {
                 'query_id': query_id,
                 'pending_flag': True,
-                'limit': actual_limit,
+                'limit': self._limit,
                 'format': 'json'
             }
         }

products/microsoft_defender_for_endpoints.py

@@ -42,13 +42,17 @@ class DefenderForEndpoints(Product):
     product: str = 'dfe'
     creds_file: str  # path to credential configuration file
     _token: str  # AAD access token
+    _limit: int = -1
 
     def __init__(self, profile: str, creds_file: str, **kwargs):
         if not os.path.isfile(creds_file):
             raise ValueError(f'Credential file {creds_file} does not exist')
 
         self.creds_file = creds_file
 
+        if 100000 >= int(kwargs.get('limit', -1)) > self._limit:
+            self._limit = int(kwargs['limit'])
+
         super().__init__(self.product, profile, **kwargs)
 
     def _authenticate(self) -> None:
@@ -139,6 +143,9 @@ def process_search(self, tag: Tag, base_query: dict, query: str) -> None:
 
         query += f" {self.build_query(base_query)}" if base_query != {} else ''
 
+        if self._limit > 0 and 'limit' not in query:
+            query += f"| limit {str(self._limit)}"
+
         self.log.debug(f'Query: {query}')
         full_query = {'Query': query}
 
@@ -154,10 +161,13 @@ def nested_process_search(self, tag: Tag, criteria: dict, base_query: dict) -> N
                     if isinstance(terms, list):
                         for query_entry in terms:
                             query_entry += f" {query_base}" if query_base != '' else ''
+                            if self._limit > 0: query_entry += f"| limit {str(self._limit)}"
                             self.process_search(tag, {}, query_entry)
                     else:
                         query_entry = terms
                         query_entry += f" {query_base}" if query_base != '' else ''
+                        if self._limit > 0: query_entry += f"| limit {str(self._limit)}"
+
                         self.process_search(tag, {}, query_entry)
                 else:
                     all_terms = ', '.join(f"'{term}'" for term in terms)
@@ -176,6 +186,8 @@ def nested_process_search(self, tag: Tag, criteria: dict, base_query: dict) -> N
 
                     query += f"| project Timestamp, {', '.join(PARAMETER_MAPPING[search_field]['projections'])}"
 
+                    if self._limit > 0: query += f"| limit {str(self._limit)}"
+
                     self.process_search(tag, {}, query)
         except KeyboardInterrupt:
             self._echo("Caught CTRL-C. Returning what we have...")

products/sentinel_one.py

@@ -76,6 +76,7 @@ class SentinelOne(Product):
     """
     product: str = 's1'
     creds_file: str  # path to credential configuration file
+    _limit: int = 1000
     _token: str  # AAD access token
     _url: str  # URL of SentinelOne console
     _site_id: Optional[str]  # Site ID for SentinelOne
@@ -98,6 +99,14 @@ def __init__(self, profile: str, creds_file: str, account_id: Optional[list[str]
         self._query_base = None
         self._pq = pq
 
+        if self._pq and self._limit >= int(kwargs.get('limit',0)) > 0:
+            self._limit = int(kwargs['limit'])
+
+        elif not self._pq and 20000 > int(kwargs.get('limit',0)) > 0:
+                self._limit = int(kwargs['limit'])
+        elif not self._pq: 
+            self._limit = 20000
+
         self._last_request = 0.0
 
         # Save these values to `self` for reference in _authenticate()
@@ -338,7 +347,7 @@ def build_query(self, filters: dict) -> Tuple[str, datetime, datetime]:
         return query_base, from_date, to_date
 
     def _get_all_paginated_data(self, url: str, params: Optional[dict] = None, headers: Optional[dict] = None,
-                                key: str = 'data', after_request: Optional[Callable] = None, limit: int = 1000,
+                                key: str = 'data', after_request: Optional[Callable] = None,
                                 no_progress: bool = True, progress_desc: str = 'Retrieving data',
                                 add_default_params: bool = True) -> list[dict]:
         """
@@ -371,7 +380,7 @@ def _get_all_paginated_data(self, url: str, params: Optional[dict] = None, heade
         if add_default_params:
             params.update(self._get_default_body())
 
-        params['limit'] = limit
+        params['limit'] = self._limit
 
         if headers is None:
             headers = dict()
@@ -603,7 +612,7 @@ def _run_query(self, merged_query: str, start_date: datetime, end_date: datetime
             params.update({
                 "fromDate": datetime_to_epoch_millis(start_date),
                 "toDate": datetime_to_epoch_millis(end_date),
-                "limit": 20000,
+                "limit": self._limit,
                 "query": merged_query
             })
 

products/vmware_cb_enterprise_edr.py

@@ -38,10 +38,12 @@ def _convert_relative_time(relative_time) -> str:
 class CbEnterpriseEdr(Product):
     product: str = 'cbc'
     _conn: CBCloudAPI  # CB Cloud API
+    _limit: int = -1
 
     def __init__(self, profile: str, **kwargs):
         self._device_group = kwargs['device_group'] if 'device_group' in kwargs else None
         self._device_policy = kwargs['device_policy'] if 'device_group' in kwargs else None
+        self._limit = int(kwargs['limit']) if 'limit' in kwargs else self._limit
 
         super().__init__(self.product, profile, **kwargs)
 
@@ -118,6 +120,9 @@ def perform_query(self, tag: Tag, base_query: dict, query: str) -> set[Result]:
                 result = Result(hostname, user, proc_name, cmdline, (ts, proc_guid,))
 
                 results.add(result)
+                if self._limit > 0 and len(results)+1 > self._limit:
+                    break
+
         except cbc_sdk.errors.ApiError as e:
             self._echo(f'CbC SDK Error (see log for details): {e}', logging.ERROR)
             self.log.exception(e)

products/vmware_cb_response.py

@@ -9,9 +9,11 @@
 class CbResponse(Product):
     product: str = 'cbr'
     _conn: CbEnterpriseResponseAPI  # CB Response API
+    _limit: int = -1
 
     def __init__(self, profile: str, **kwargs):
         self._sensor_group = kwargs['sensor_group'] if 'sensor_group' in kwargs else None
+        self._limit = int(kwargs['limit']) if 'limit' in kwargs else self._limit
 
         super().__init__(self.product, profile, **kwargs)
 
@@ -58,6 +60,10 @@ def process_search(self, tag: Tag, base_query: dict, query: str) -> None:
                 result = Result(proc.hostname.lower(), proc.username.lower(), proc.path, proc.cmdline,
                                 (proc.start, proc.id))
                 results.add(result)
+
+                if self._limit > 0 and len(results)+1 > self._limit:
+                        break
+
         except KeyboardInterrupt:
             self._echo("Caught CTRL-C. Returning what we have . . .")
 
@@ -89,6 +95,9 @@ def nested_process_search(self, tag: Tag, criteria: dict, base_query: dict) -> N
                     result = Result(proc.hostname.lower(), proc.username.lower(), proc.path, proc.cmdline,
                                     (proc.start,))
                     results.add(result)
+                    if self._limit > 0 and len(results)+1 > self._limit:
+                        break
+
         except Exception as e:
             self._echo(f'Error (see log for details): {e}', logging.ERROR)
             self.log.exception(e)

surveyor.py

@@ -85,6 +85,7 @@ class ExecutionOptions:
     days: Optional[int]
     minutes: Optional[int]
     username: Optional[str]
+    limit: Optional[int]
     ioc_file: Optional[str]
     ioc_type: Optional[str]
     query: Optional[str]
@@ -106,6 +107,17 @@ class ExecutionOptions:
 @click.option("--profile", help="The credentials profile to use.", type=click.STRING)
 @click.option("--days", help="Number of days to search.", type=click.INT)
 @click.option("--minutes", help="Number of minutes to search.", type=click.INT)
+@click.option("--limit",help="""
+              Number of results to return. Cortex XDR: Default: 1000, Max: Default
+              Microsoft Defender for Endpoint: Default/Max: 100000
+              SentinelOne (PowerQuery): Default/Max: 1000
+              SentinelOne (Deep Visibility): Default/Max: 20000
+              VMware Carbon Black EDR: Default/Max: None
+              VMware Carbon Black Cloud Enterprise EDR: Default/Max: None
+
+              Note: Exceeding the maximum limits will automatically set the limit to its maximum value, where applicable.
+              """
+              , type=click.INT)
 @click.option("--hostname", help="Target specific host by name.", type=click.STRING)
 @click.option("--username", help="Target specific username.")
 # different ways you can survey the EDR
@@ -127,14 +139,14 @@ class ExecutionOptions:
 @click.option("--log-dir", 'log_dir', help="Specify the logging directory.", type=click.STRING, default='logs')
 @click.pass_context
 def cli(ctx, prefix: Optional[str], hostname: Optional[str], profile: str, days: Optional[int], minutes: Optional[int],
-        username: Optional[str],
+        username: Optional[str], limit: Optional[int],
         ioc_file: Optional[str], ioc_type: Optional[str], query: Optional[str], output: Optional[str],
         def_dir: Optional[str], def_file: Optional[str], no_file: bool, no_progress: bool,
         sigma_rule: Optional[str], sigma_dir: Optional[str],
         log_dir: str) -> None:
 
     ctx.ensure_object(dict)
-    ctx.obj = ExecutionOptions(prefix, hostname, profile, days, minutes, username, ioc_file, ioc_type, query, output,
+    ctx.obj = ExecutionOptions(prefix, hostname, profile, days, minutes, username, limit, ioc_file, ioc_type, query, output,
                                def_dir, def_file, sigma_rule, sigma_dir, no_file, no_progress, log_dir, dict())
 
     if ctx.invoked_subcommand is None:
@@ -265,6 +277,10 @@ def survey(ctx, product_str: str = 'cbr') -> None:
     if len(opt.product_args) > 0:
         kwargs.update(opt.product_args)
 
+    if opt.limit:
+        kwargs['limit'] = str(opt.limit)
+
+
     kwargs['tqdm_echo'] = str(not opt.no_progress)
 
     # instantiate a product class instance based on the product string