feat: add support for user supplied access token (#174)

Signed-off-by: Timothy MacDonald <[email protected]>
lacework · Jun 11, 2024 · 977808a · 977808a
1 parent 2acc932
commit 977808a
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -78,16 +78,20 @@ To generate API credentials, you'll need to do the following in Lacework:
 ## Environment Variables
 
 If you wish to configure the LaceworkClient instance using environment variables, this module honors the same
-variables used by the Lacework CLI. The `account`, `subaccount`, `api_key`, `api_secret`, and `profile` parameters
+variables used by the Lacework CLI. The `account`, `subaccount`, `api_key`, `api_secret`, `api_token`, and `profile` parameters
 can all be configured as specified below.
 
 | Environment Variable | Description                                                          | Required |
 | -------------------- | -------------------------------------------------------------------- | :------: |
 | `LW_PROFILE`         | Lacework CLI profile to use (configured at ~/.lacework.toml)         |    N     |
 | `LW_ACCOUNT`         | Lacework account/organization domain (i.e. `<account>`.lacework.net) |    Y     |
 | `LW_SUBACCOUNT`      | Lacework sub-account                                                 |    N     |
-| `LW_API_KEY`         | Lacework API Access Key                                              |    Y     |
-| `LW_API_SECRET`      | Lacework API Access Secret                                           |    Y     |
+| `LW_API_KEY`         | Lacework API Access Key                                              |    N     |
+| `LW_API_SECRET`      | Lacework API Access Secret                                           |    N     |
+| `LW_API_TOKEN`       | Lacework API Token (alternative to key and secret)                   |    N     |
+
+NOTE: To authenticate with the Lacework API you must specify either a key and secret OR a token. If you specify both the 
+token will be used. 
 
 ## Installation
 

diff --git a/laceworksdk/api/__init__.py b/laceworksdk/api/__init__.py
@@ -49,6 +49,7 @@
     LACEWORK_SUBACCOUNT_ENVIRONMENT_VARIABLE,
     LACEWORK_API_KEY_ENVIRONMENT_VARIABLE,
     LACEWORK_API_SECRET_ENVIRONMENT_VARIABLE,
+    LACEWORK_API_TOKEN_ENVIRONMENT_VARIABLE,
     LACEWORK_API_BASE_DOMAIN_ENVIRONMENT_VARIABLE,
     LACEWORK_API_CONFIG_SECTION_ENVIRONMENT_VARIABLE,
     LACEWORK_CLI_CONFIG_RELATIVE_PATH,
@@ -66,6 +67,7 @@ def __init__(
         subaccount=None,
         api_key=None,
         api_secret=None,
+        api_token=None,
         instance=None,
         base_domain=None,
         profile=None,
@@ -86,6 +88,7 @@ def __init__(
         self._subaccount = subaccount or os.getenv(
             LACEWORK_SUBACCOUNT_ENVIRONMENT_VARIABLE
         )
+        self._api_token = api_token or os.getenv(LACEWORK_API_TOKEN_ENVIRONMENT_VARIABLE)
         self._api_key = api_key or os.getenv(LACEWORK_API_KEY_ENVIRONMENT_VARIABLE)
         self._api_secret = api_secret or os.getenv(
             LACEWORK_API_SECRET_ENVIRONMENT_VARIABLE
@@ -135,6 +138,7 @@ def __init__(
             self._api_key,
             self._api_secret,
             self._base_domain,
+            api_token=self._api_token
         )
 
         # API Wrappers

diff --git a/laceworksdk/config.py b/laceworksdk/config.py
@@ -14,6 +14,7 @@
 LACEWORK_SUBACCOUNT_ENVIRONMENT_VARIABLE = "LW_SUBACCOUNT"
 LACEWORK_API_KEY_ENVIRONMENT_VARIABLE = "LW_API_KEY"
 LACEWORK_API_SECRET_ENVIRONMENT_VARIABLE = "LW_API_SECRET"
+LACEWORK_API_TOKEN_ENVIRONMENT_VARIABLE = "LW_API_TOKEN"
 LACEWORK_API_BASE_DOMAIN_ENVIRONMENT_VARIABLE = "LW_BASE_DOMAIN"
 LACEWORK_API_CONFIG_SECTION_ENVIRONMENT_VARIABLE = "LW_PROFILE"
 

diff --git a/laceworksdk/http_session.py b/laceworksdk/http_session.py
@@ -30,7 +30,7 @@ class HttpSession:
     _access_token = None
     _access_token_expiry = None
 
-    def __init__(self, account, subaccount, api_key, api_secret, base_domain):
+    def __init__(self, account, subaccount, api_key, api_secret, base_domain, api_token=None):
         """
             Initializes the HttpSession object.
 
@@ -40,6 +40,7 @@ def __init__(self, account, subaccount, api_key, api_secret, base_domain):
             api_key (str): a Lacework API Key
             api_secret (str): a Lacework API Secret
             base_domain (str): a Lacework Domain (defaults to "lacework.net")
+            api_token (str): a Lacework API token (instead of key and secret)
 
         Returns:
             HttpSession: An instance of this class
@@ -59,7 +60,7 @@ def __init__(self, account, subaccount, api_key, api_secret, base_domain):
         self._account = account
         self._subaccount = subaccount
         self._org_level_access = False
-
+        self._access_token = api_token
         # Get an access token
         self._check_access_token()
 
@@ -99,8 +100,11 @@ def _check_access_token(self):
         """
         A method to check the validity of the access token.
         """
-
-        if self._access_token is None or self._access_token_expiry < datetime.now(
+        if self._access_token and self._access_token_expiry is None:
+            # This catches the case that the user has provided an access token instead of 
+            # key and secret. We cannot know the expiry date so we simply return
+            return
+        elif self._access_token is None or self._access_token_expiry < datetime.now(
             timezone.utc
         ):
             response = self._get_access_token()