[DPDV-5997] feat: allow query all or selected accounts

README.md

@@ -52,16 +52,32 @@ The add-on uses Splunk encrypted secrets storage, so admins require `admin_all_o
 4. Under Log Access Keys, click Add Key > Add Write Key (required for alert action).
 5. Optionally, click the pencil icon to rename the keys.
 
+### SentinelOne Platform with Singularity Data Lake
+To get the AuthN API token follow the below mentioned details:
+1. Login into the SentinelOne console, click on your User Name > My User
+![My User Page](README_images/my_user.png)
+2. Click on Actions > API Token Operations > Generate / Regenrate API token.
+    - If you are generating the API Token for the first time then you will have the Generate API Token option. Otherwise you will find the generate API Token.
+![Token generation](README_images/generate_token.png)
+3. Copy the API Token and save it for configuration.
+
 ### Splunk
 1. In Splunk, open the Add-on
 
-![Configuring DataSet Account](README_images/setup_account.png)
+![Configuring DataSet Account](README_images/acc_details_new.png)
 
 2. On the configuration > account tab:
 - Click Add
 - Enter a user-friendly account name. For multiple accounts, the account name can be used in queries (more details below).
 - Enter the full URL noted above (e.g.: `https://app.scalyr.com`, `https://xdr.us1.sentinelone.net` or `https://xdr.eu1.sentinelone.net`).
-- Enter the DataSet read key from above (required for searching)
+- Enter Tenant value, it can be True/False/Blank. If set to True, the queries will run for the entire Tenant or if set to False, provide Account IDs as a comma separated values to run searches in those specific accounts. Leave it blank if you are not trying use the Tenant level searches.
+- Provide the comma seperated Account Ids, if Tenant is False. eg: 1234567890,9876543210.
+- Enter the AuthN API Token First part which includes first 220 characters.
+- Enter the AuthN API Token Second part which includes remaining characters.
+    - Use this command to prepare both parts of AuthN API token:
+    `read -p "Enter Token: " input_string && echo "Part1: $(echo $input_string | cut -c 1-220)"; echo "Part2: $(echo $input_string | cut -c 221-)"`
+    - Reason for creating 2 parts of AuthN Token: Splunk Storage Manager has a limitation of storing only 256 characters of encrypted data from inputs. And the AuthN Token can have length <256, hence its split into 2 parts, the first one is encrypted (first 220 chars) and the second one is not. As we are encrypting most of the Token, its use is safe.
+- Enter the DataSet read key from above (required for searching), please ignore this if AuthN token value is provided.
 - Enter the DataSet write key from above (only required for alert actions).
 - Click Save
 

TA_dataset/README/ta_dataset_settings.conf.spec

@@ -1,5 +1,9 @@
 [account]
 url = <string>
+tenant = <bool>
+account_ids = <string>
+an_fir_part = <string>
+an_sec_part = <string>
 dataset_log_read_access_key = <string>
 dataset_log_write_access_key = <string>
 

TA_dataset/bin/dataset_alerts.py

@@ -14,6 +14,7 @@
     get_proxy,
     get_url,
     relative_to_epoch,
+    get_tenant_related_payload,
 )
 from solnlib import log
 from solnlib.modular_input import checkpointer
@@ -110,7 +111,11 @@ def stream_events(self, inputs, ew):
                 ds_headers = {
                     "Authorization": "Bearer " + acct_dict[ds_acct]["ds_api_key"]
                 }
-
+                tenant_related_payload = get_tenant_related_payload(acct_dict.get(ds_acct))
+                ds_payload.update(tenant_related_payload)
+                logger.debug(
+                    "ds payload in power query stream events = {}".format(ds_payload)
+                )
                 # Create checkpointer
                 checkpoint = checkpointer.KVStoreCheckpointer(
                     input_name, session_key, APP_NAME

TA_dataset/bin/dataset_api.py

@@ -44,15 +44,17 @@ def convert_proxy(proxy):
     return new_proxy
 
 
+def get_tenant_details(acc_conf):
+    if acc_conf.get("tenant") is not None:
+        tenant_value = acc_conf.get("tenant")
+        if tenant_value:
+            return {"tenant": True}
+        return {"tenant": False, "accountIds": acc_conf["account_ids"]}
+
+
 # Executes Dataset LongRunningQuery for log events
 def ds_lrq_log_query(
-    base_url,
-    api_key,
-    start_time,
-    end_time,
-    filter_expr,
-    limit,
-    proxy,
+    base_url, api_key, start_time, end_time, filter_expr, limit, proxy, logger, acc_conf
 ):
     client = AuthenticatedClient(
         base_url=base_url, token=api_key, proxy=convert_proxy(proxy)
@@ -63,11 +65,16 @@ def ds_lrq_log_query(
         end_time=end_time,
         log=LogAttributes(filter_=filter_expr, limit=limit),
     )
-    return ds_lrq_run_loop(client=client, body=body)
+    tenant_details = get_tenant_details(acc_conf=acc_conf)
+    return ds_lrq_run_loop(
+        logger, client=client, body=body, tenant_details=tenant_details
+    )
 
 
 # Executes Dataset LongRunningQuery using PowerQuery language
-def ds_lrq_power_query(base_url, api_key, start_time, end_time, query, proxy):
+def ds_lrq_power_query(
+    base_url, api_key, start_time, end_time, query, proxy, logger, acc_conf
+):
     client = AuthenticatedClient(
         base_url=base_url, token=api_key, proxy=convert_proxy(proxy)
     )
@@ -77,7 +84,10 @@ def ds_lrq_power_query(base_url, api_key, start_time, end_time, query, proxy):
         end_time=end_time,
         pq=PQAttributes(query=query),
     )
-    return ds_lrq_run_loop(client=client, body=body)
+    tenant_details = get_tenant_details(acc_conf=acc_conf)
+    return ds_lrq_run_loop(
+        logger, client=client, body=body, tenant_details=tenant_details
+    )
 
 
 # Executes Dataset LongRunningQuery to fetch facet values
@@ -90,6 +100,8 @@ def ds_lrq_facet_values(
     name,
     max_values,
     proxy,
+    logger,
+    acc_conf,
 ):
     client = AuthenticatedClient(
         base_url=base_url, token=api_key, proxy=convert_proxy(proxy)
@@ -102,17 +114,25 @@ def ds_lrq_facet_values(
             filter_=filter, name=name, max_values=max_values
         ),
     )
-    return ds_lrq_run_loop(client=client, body=body)
+    tenant_details = get_tenant_details(acc_conf=acc_conf)
+    return ds_lrq_run_loop(
+        logger, client=client, body=body, tenant_details=tenant_details
+    )
 
 
 # Executes LRQ run loop of launch-ping-remove API requests until the query completes
 # with a result
 # Returns tuple - value, error message
 def ds_lrq_run_loop(
-    client: AuthenticatedClient, body: PostQueriesLaunchQueryRequestBody
+    log,
+    client: AuthenticatedClient,
+    body: PostQueriesLaunchQueryRequestBody,
+    tenant_details=None,
 ):
     body.query_priority = PostQueriesLaunchQueryRequestBodyQueryPriority.HIGH
-    response = post_queries.sync_detailed(client=client, json_body=body)
+    response = post_queries.sync_detailed(
+        client=client, json_body=body, tenant_details=tenant_details, logger=log
+    )
     logger().debug(response)
     result = response.parsed
     if result:
@@ -242,7 +262,7 @@ def parse_query(ds_columns, match_list, sessions):
     if ds_columns is None:
         session_key = match_list["session"]
 
-        for session_entry, session_dict in sessions.items():
+        for session_entry, session_dict in list(sessions.items()):
             if session_entry == session_key:
                 for key in session_dict:
                     ds_event_dict[key] = session_dict[key]

TA_dataset/bin/dataset_common.py

@@ -167,9 +167,16 @@ def get_acct_info(self, logger, account=None):
                 for conf in confs:
                     acct_dict[conf.name] = {}
                     acct_dict[conf.name]["base_url"] = conf.url
-                    acct_dict[conf.name]["ds_api_key"] = get_token(
-                        self, conf.name, "read", logger
+                    acct_dict[conf.name]["ds_api_key"] = get_token_from_config(
+                        self, conf, conf.name, logger
                     )
+                    if hasattr(conf, "tenant"):
+                        tenant_value = get_tenant_value(conf, logger)
+                        acct_dict[conf.name]["tenant"] = tenant_value
+                        if not tenant_value:
+                            acct_dict[conf.name]["account_ids"] = get_account_ids(
+                                conf, logger
+                            )
             except Exception as e:
                 msg = "Error retrieving add-on settings, error = {}".format(e)
                 logger.error(msg + " - %s", e, exc_info=True)
@@ -182,9 +189,16 @@ def get_acct_info(self, logger, account=None):
                     conf = self.service.confs[conf_name][entry]
                     acct_dict[entry] = {}
                     acct_dict[entry]["base_url"] = conf.url
-                    acct_dict[entry]["ds_api_key"] = get_token(
-                        self, entry, "read", logger
+                    acct_dict[entry]["ds_api_key"] = get_token_from_config(
+                        self, conf, entry, logger
                     )
+                    if hasattr(conf, "tenant"):
+                        tenant_value = get_tenant_value(conf, logger)
+                        acct_dict[entry]["tenant"] = tenant_value
+                        if not tenant_value:
+                            acct_dict[entry]["account_ids"] = get_account_ids(
+                                conf, logger
+                            )
             except Exception as e:
                 msg = "Error retrieving account settings, error = {}".format(e)
                 logger.error(msg + " - %s", e, exc_info=True)
@@ -197,9 +211,16 @@ def get_acct_info(self, logger, account=None):
             for conf in confs:
                 acct_dict[conf.name] = {}
                 acct_dict[conf.name]["base_url"] = conf.url
-                acct_dict[conf.name]["ds_api_key"] = get_token(
-                    self, conf.name, "read", logger
+                acct_dict[conf.name]["ds_api_key"] = get_token_from_config(
+                    self, conf, conf.name, logger
                 )
+                if hasattr(conf, "tenant"):
+                    tenant_value = get_tenant_value(conf, logger)
+                    acct_dict[conf.name]["tenant"] = tenant_value
+                    if not tenant_value:
+                        acct_dict[conf.name]["account_ids"] = get_account_ids(
+                            conf, logger
+                        )
                 break
         except Exception as e:
             msg = (
@@ -212,7 +233,57 @@ def get_acct_info(self, logger, account=None):
     return acct_dict
 
 
-def get_token(self, account, rw, logger):
+def get_tenant_value(conf, logger):
+    tenant_value = conf.tenant
+    tenant_value = tenant_value.strip()
+    logger.info("The provided tenant value in config is {}".format(tenant_value))
+    if tenant_value.lower() == "false" or tenant_value.lower() == "0":
+        return False
+    return True
+
+
+def get_account_ids(conf, logger):
+    account_ids_array = []
+    if hasattr(conf, "account_ids"):
+        account_ids_conf = conf.account_ids
+        account_ids_conf = account_ids_conf.strip()
+        if account_ids_conf:
+            account_ids_array = account_ids_conf.split(",")
+        logger.info(f"the provided account ids in config: {account_ids_array}")
+    if not account_ids_array:
+        raise Exception(
+            "Tenant is false, so please provide the valid comma-separated account IDs"
+            " in the account configuration page."
+        )
+    return account_ids_array
+
+
+def get_token_from_config(self, conf, name, logger):
+    authn_token = ""
+    if hasattr(conf, "an_fir_part"):
+        logger.info("The AuthN api token first part was available")
+        first_half = get_token(self, name, "authn", logger, "an_fir_part")
+        authn_token += first_half
+    if hasattr(conf, "an_sec_part"):
+        logger.info("The AuthN api token second part was available")
+        second_part = conf.an_sec_part
+        authn_token += second_part
+    if not hasattr(conf, "an_fir_part") and not hasattr(conf, "an_sec_part"):
+        logger.info("The AuthN api token was not available")
+        return get_token(self, name, "read", logger)
+
+    return authn_token
+
+def get_tenant_related_payload(ds_acct):
+    if ds_acct.get("tenant") is not None:
+        tenant_value = ds_acct.get("tenant")
+        if tenant_value:
+            return {"tenant": True}
+        return {"tenant": False, "accountIds": ds_acct["account_ids"],}
+
+
+
+def get_token(self, account, token_type, logger, config_key=None):
     try:
         # use Python SDK secrets retrieval
         for credential in self.service.storage_passwords:
@@ -224,15 +295,22 @@ def get_token(self, account, rw, logger):
                 and credential.username.startswith(account)
             ):
                 cred = credential.content.get("clear_password")
-                if rw == "read":
+                if token_type == "authn":
+                    if config_key in cred:
+                        logger.info("the yes on authn token")
+                        cred_json = json.loads(cred)
+                        token = cred_json[config_key]
+                if token_type == "read":
                     if "dataset_log_read_access_key" in cred:
                         cred_json = json.loads(cred)
                         token = cred_json["dataset_log_read_access_key"]
-                elif rw == "write":
+                elif token_type == "write":
                     if "dataset_log_write_access_key" in cred:
                         cred_json = json.loads(cred)
                         token = cred_json["dataset_log_write_access_key"]
                 return token
+            else:
+                logger.debug("the credentials were not retireived")
     except Exception as e:
         logger.error(
             self,

TA_dataset/bin/dataset_powerquery.py

@@ -8,7 +8,7 @@
 import import_declare_test  # noqa: F401
 import requests
 from dataset_api import build_payload, parse_powerquery
-from dataset_common import get_acct_info, get_log_level, get_proxy, get_url
+from dataset_common import get_acct_info, get_log_level, get_proxy, get_url, get_tenant_related_payload
 from solnlib import log
 from solnlib.modular_input import checkpointer
 from splunklib import modularinput as smi
@@ -117,7 +117,11 @@ def stream_events(self, inputs, ew):
                 ds_headers = {
                     "Authorization": "Bearer " + acct_dict[ds_acct]["ds_api_key"]
                 }
-
+                tenant_related_payload = get_tenant_related_payload(acct_dict.get(ds_acct))
+                ds_payload.update(tenant_related_payload)
+                logger.info(
+                    "ds payload in power query stream events = {}".format(ds_payload)
+                )
                 # Create checkpointer
                 checkpoint = checkpointer.KVStoreCheckpointer(
                     input_name, session_key, APP_NAME

TA_dataset/bin/dataset_query.py

@@ -18,6 +18,7 @@
     get_proxy,
     get_url,
     relative_to_epoch,
+    get_tenant_related_payload,
 )
 from dataset_query_api_client.client import get_user_agent
 from solnlib.modular_input import checkpointer
@@ -144,7 +145,10 @@ def stream_events(self, inputs, ew):
             proxy = get_proxy(session_key, logger)
             acct_dict = get_acct_info(self, logger, ds_account)
             for ds_acct in acct_dict.keys():
+                tenant_related_payload = get_tenant_related_payload(acct_dict.get(ds_acct))
+                ds_payload.update(tenant_related_payload)
                 curr_payload = copy.deepcopy(ds_payload)
+                logger.info("query api account curr payload {}".format(curr_payload))
                 curr_maxcount = copy.copy(ds_maxcount)
                 ds_url = get_url(acct_dict[ds_acct]["base_url"], "query")
                 ds_headers = {

TA_dataset/bin/dataset_search_command.py

@@ -29,6 +29,7 @@
     get_url,
     logger,
     relative_to_epoch,
+    get_tenant_related_payload,
 )
 
 # Dataset V2 API client (generated)
@@ -326,6 +327,8 @@ def generate(self):
                         filter_expr=ds_search,
                         limit=ds_maxcount,
                         proxy=proxy,
+                        logger=logger,
+                        acc_conf=acct_dict[ds_acct],
                     )
 
                     logger.debug("QUERY RESULT, result={}".format(result))
@@ -368,6 +371,8 @@ def generate(self):
                         end_time=ds_end,
                         query=pq,
                         proxy=proxy,
+                        logger=logger,
+                        acc_conf=acct_dict[ds_acct],
                     )
                     logger.debug("QUERY RESULT, result={}".format(result))
                     data = result.data  # TableResultData
@@ -404,6 +409,8 @@ def generate(self):
                         name=f_field,
                         max_values=ds_maxcount,
                         proxy=proxy,
+                        logger=logger,
+                        acc_conf=acct_dict[ds_acct],
                     )
                     logger.debug("QUERY RESULT, result={}".format(result))
                     facet = result.data.facet  # FacetValuesResultData.data -> FacetData
@@ -436,6 +443,11 @@ def generate(self):
                         "DataSetFunction=makeRequest, destination={}, startTime={}"
                         .format(ds_url, time.time())
                     )
+                    tenant_related_payload = get_tenant_related_payload(acct_dict.get(ds_acct))
+                    ds_payload.update(tenant_related_payload)
+                    logger.info(
+                        "The paylaod for the timeseries api {}".format(ds_payload)
+                    )
                     r = requests.post(
                         url=ds_url, headers=ds_headers, json=ds_payload, proxies=proxy
                     )