[DPDV-5997] feat: allow query all or selected accounts (#134)

* [DPDV-5997] feat: allow query all or selected accounts (#127) * [DPDV-5997] adding account and tenant to queries * [DPDV-5997] lint fix * [DPDV-5997] lint fix * [DPDV-5997] updated config * [DPDV-5997] updated documentation * [DPDV-5997] updated auth token encryption * [DPDV-5997] pre-commit fix * [DPDV-5997] pre-commit fix * [DPDV-5997] review comments fix * [DPDV-5997] pre commit fix * [DPDV-5997] review comment fix * [DPDV-5997] pre-commit fix * [DPDV-5997] review comment fix * [DPDV-5997] pre-commit fix * [DPDV-5997] review comment fixes * [DPDV-5997] pre-commit fixes * [DPDV-5997] conflict fixes * [DPDV-5997] fix: fixing the playwright errors * [DPDV-5997] fix: rollback the changes of tests * [DPDV-5997] fix: new review comments fixes * [DPDV-5997] fix: lint fix * [DPDV-5997] fix: review comment fix
scalyr · May 13, 2024 · d0e72d7 · d0e72d7
1 parent 22fba6f
commit d0e72d7
Show file tree

Hide file tree

Showing 14 changed files with 877 additions and 651 deletions.
diff --git a/README.md b/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
 

diff --git a/README_images/acc_details_new.png b/README_images/acc_details_new.png
diff --git a/README_images/generate_token.png b/README_images/generate_token.png
diff --git a/README_images/my_user.png b/README_images/my_user.png
diff --git a/TA_dataset/README/ta_dataset_settings.conf.spec b/TA_dataset/README/ta_dataset_settings.conf.spec
@@ -1,5 +1,9 @@
 [account]
 url = <string>
+tenant = <string>
+account_ids = <string>
+authn_token_part_one = <string>
+authn_token_part_two = <string>
 dataset_log_read_access_key = <string>
 dataset_log_write_access_key = <string>
 

diff --git a/TA_dataset/bin/dataset_alerts.py b/TA_dataset/bin/dataset_alerts.py
@@ -12,6 +12,7 @@
     get_acct_info,
     get_log_level,
     get_proxy,
+    get_tenant_related_payload,
     get_url,
     relative_to_epoch,
 )
@@ -110,7 +111,13 @@ 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

diff --git a/TA_dataset/bin/dataset_api.py b/TA_dataset/bin/dataset_api.py
@@ -53,6 +53,8 @@ def ds_lrq_log_query(
     filter_expr,
     limit,
     proxy,
+    logger,
+    tenant_related_payload,
 ):
     client = AuthenticatedClient(
         base_url=base_url, token=api_key, proxy=convert_proxy(proxy)
@@ -63,11 +65,21 @@ 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)
+    body.additional_properties = tenant_related_payload
+    return ds_lrq_run_loop(logger, client=client, body=body)
 
 
 # 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,
+    tenant_related_payload,
+):
     client = AuthenticatedClient(
         base_url=base_url, token=api_key, proxy=convert_proxy(proxy)
     )
@@ -77,7 +89,8 @@ 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)
+    body.additional_properties = tenant_related_payload
+    return ds_lrq_run_loop(logger, client=client, body=body)
 
 
 # Executes Dataset LongRunningQuery to fetch facet values
@@ -90,6 +103,8 @@ def ds_lrq_facet_values(
     name,
     max_values,
     proxy,
+    logger,
+    tenant_related_payload,
 ):
     client = AuthenticatedClient(
         base_url=base_url, token=api_key, proxy=convert_proxy(proxy)
@@ -102,17 +117,24 @@ def ds_lrq_facet_values(
             filter_=filter, name=name, max_values=max_values
         ),
     )
-    return ds_lrq_run_loop(client=client, body=body)
+    body.additional_properties = tenant_related_payload
+    return ds_lrq_run_loop(logger, client=client, body=body)
 
 
 # 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,
 ):
+    """
+    The log parameter assists in generating logs with the chosen level
+    in the index="_internal". This helps in debugging.
+    """
     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, logger=log)
     logger().debug(response)
     result = response.parsed
     if result:
@@ -242,7 +264,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]

diff --git a/TA_dataset/bin/dataset_common.py b/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,63 @@ 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.debug("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.debug(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, "authn_token_part_one"):
+        logger.debug("The AuthN api token first part was available")
+        first_half = get_token(self, name, "authn", logger, "authn_token_part_one")
+        authn_token += first_half
+    if hasattr(conf, "authn_token_part_two"):
+        logger.debug("The AuthN api token second part was available")
+        second_part = conf.authn_token_part_two
+        authn_token += second_part
+    if not hasattr(conf, "authn_token_part_one") and not hasattr(
+        conf, "authn_token_part_two"
+    ):
+        logger.debug("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"],
+        }
+    return {}
+
+
+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 +301,21 @@ 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:
+                        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,

diff --git a/TA_dataset/bin/dataset_powerquery.py b/TA_dataset/bin/dataset_powerquery.py
@@ -8,7 +8,13 @@
 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_tenant_related_payload,
+    get_url,
+)
 from solnlib import log
 from solnlib.modular_input import checkpointer
 from splunklib import modularinput as smi
@@ -117,7 +123,13 @@ 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

diff --git a/TA_dataset/bin/dataset_query.py b/TA_dataset/bin/dataset_query.py
@@ -16,6 +16,7 @@
     get_acct_info,
     get_logger,
     get_proxy,
+    get_tenant_related_payload,
     get_url,
     relative_to_epoch,
 )
@@ -144,7 +145,12 @@ 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 = {