Added new APIv2 endpoints that allow advanced searches and included p…

…agination functionality (#47) * refactor: updated HttpSession class to be much more flexible * feat: allow org-level access and sub-account to be modified dynamically * refactor: created three base class types that endpoints should inherit * refactor: reorganized existing code and refactored to use new base classes * feat: implemented all new APIv2 endpoints including pagination support * chore: included flake8-quotes in dev dependencies * docs: small update to README.md for latest API endpoints * refactor: corrected docs and simplified `build_dict_from_items` method * fix: sanitizing access to 'query_data' * fix: narrowing down exception handling for pagination iteration * refactor: moved resource attribute to base of SeachEndpoint * refactor: removed unnecessary override method * docs: added docstrings for 'query_data' where needed * docs: added docstrings for methods which are passed * refactor: changed LaceworkException name to match conventions * tests: massively de-duped code for new search API tests * chore: importing Retry directly from urllib Ref: https://github.com/psf/requests/blob/v2.22.0/requests/packages.py * fix: passthrough Lacework response when maximum retries reached * fix: supply json to the query execute function * docs: improved function documentation for the BaseEndpoint class * chore: removed references to `query_data` in favor of `json` * chore: removed redundant `search()` method overrides * fix: improved consistency of variable naming in AgentAccessTokensAPI * chore: improved error message for JSONDecodeError * docs: added additional class docstrings * fix: fixed bugs and consistency issues found in testing * refactor: simplified and modernized all APIv2 tests * tests: fixed dependency issue with tests * fix: modified error logging for 'nextPage' parsing * refactor: changed LaceworksdkException to LaceworkSDKException
lacework · Jan 31, 2022 · ee3f6f8 · ee3f6f8
1 parent 6f0b47e
commit ee3f6f8
Show file tree

Hide file tree

Showing 107 changed files with 4,208 additions and 3,753 deletions.
diff --git a/.github/workflows/python-test-flaky.yml b/.github/workflows/python-test-flaky.yml
@@ -26,6 +26,7 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install flake8 flake8-quotes pytest pytest-rerunfailures
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+          if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
           if [ -f jupyter/requirements.txt ]; then pip install -r jupyter/requirements.txt; fi
 
       - name: Run setup.py

diff --git a/.github/workflows/python-test.yml b/.github/workflows/python-test.yml
@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.6", "3.7", "3.8", "3.9"]
+        python-version: ["3.6", "3.7", "3.8", "3.9", "3.10"]
 
     steps:
       - uses: actions/checkout@v2
@@ -25,8 +25,9 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install flake8 flake8-quotes pytest pytest-rerunfailures
+          python -m pip install flake8 flake8-quotes pytest pytest-lazy-fixture pytest-rerunfailures
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+          if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
           if [ -f jupyter/requirements.txt ]; then pip install -r jupyter/requirements.txt; fi
 
       - name: Lint with flake8

diff --git a/README.md b/README.md
@@ -6,22 +6,42 @@
 **laceworksdk** is a community developed Python library for interacting with the Lacework APIs.
 
 The purpose of this library is to simplify the common tasks required for interacting with the Lacework API, and allow
-users write simple code to automate tasks related to their Lacework instance. From data retrieval to configuration,
+users write simple code to automate tasks related to their Lacework instance(s). From data retrieval to configuration,
 this library aims to expose all publicly available APIs. For example, the following code would authenticate,
-fetch events, fetch host vulnerabilities, and fetch container vulnerabilities - in 5 lines of code.
+fetch events, fetch host vulnerabilities, and fetch container vulnerabilities. The latest version of the SDK supports
+expressive searches as enabled by v2 of the Lacework APIs.
 
-```
+```python
 from laceworksdk import LaceworkClient
 
+lw = LaceworkClient() # This would leverage your default Lacework CLI profile.
 lw = LaceworkClient(account="ACCOUNT",
+                    subaccount="SUBACCOUNT",
                     api_key="API KEY",
                     api_secret="API SECRET")
 
 events = lw.events.get_for_date_range(start_time=start_time, end_time=end_time)
 
-host_vulns = lw.vulnerabilities.get_host_vulnerabilities()
-
-container_vulns = lw.vulnerabilities.get_container_vulnerabilities(image_digest="sha256:123")
+host_vulns = lw.vulnerabilities.hosts.search(json={
+    "timeFilters": {
+        "startTime": start_time,
+        "endTime": end_time
+    }
+})
+
+container_vulns = lw.vulnerabilities.containers.search(json={
+    "timeFilters": {
+        "startTime": start_time,
+        "endTime": end_time
+    },
+    "filters": [
+        {
+            "field": "imageId",
+            "expression": "eq",
+            "value": "sha256:657922eb2d64b0a34fe7339f8b48afb9f2f44635d7d6eaa92af69591d29b3330"
+        }
+    ]
+})
 ```
 
 ## Requirements
@@ -34,30 +54,32 @@ container_vulns = lw.vulnerabilities.get_container_vulnerabilities(image_digest=
 
 ## How-To
 
-The following data points are required to instantiate a LaceworkClient instance:
+The following information is required to instantiate a LaceworkClient instance:
 
 - `account`: The Lacework account/organization domain. (`xxxxx`.lacework.net)
-- `subaccount`: (Optional) The Lacework sub-account domain. (`xxxxx`.lacework.net)
-  - This is only used if leveraging the Manage@Scale organization feature of Lacework
 - `api_key`: The API Key that was generated from the Lacework UI/API.
 - `api_secret`: The API Secret that was generated from the Lacework UI/API.
 
-To generate an API Key and Secret, do the following:
+Optionally, you can also set a Lacework Sub-Account using the `subaccount` parameter.
+
+To generate API credentials, you'll need to do the following in Lacework:
 
 1.  In the Lacework web interface, go to Settings -> API Keys
-2.  Create a new API Key, or download information for an existing one.
+2.  Create a new API Key and download information the credentials.
 
-### Environment Variables
+## Environment Variables
 
-The `account`, `subaccount`, `api_key`, and `api_secret` can also be set using environment variables or
-saved in ~/.lacework.toml configuration file (same file as the Lacework CLI uses).
+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
+can all be configured as specified below.
 
-| Environment Variable | Description                                                      | Required |
-| -------------------- | ---------------------------------------------------------------- | :------: |
-| `LW_ACCOUNT`         | Lacework account/organization domain (i.e. `xxxxx`.lacework.net) |    Y     |
-| `LW_SUBACCOUNT`      | Lacework sub-account domain (i.e. `xxxxx`.lacework.net)          |    N     |
-| `LW_API_KEY`         | Lacework API Access Key                                          |    Y     |
-| `LW_API_SECRET`      | Lacework API Access Secret                                       |    Y     |
+| 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     |
 
 ## Installation
 
@@ -75,41 +97,6 @@ Installing and upgrading `laceworksdk` is easy:
 
 Are you looking for some sample scripts? Check out the [examples](examples/) folder!
 
-## Implemented APIs
-
-### API v1
-
-- [x] Account API
-- [x] Compliance API
-- [x] Custom Compliance Config API
-- [x] Download File API
-- [x] Events API
-- [x] Integrations API
-- [x] Recommendations API
-- [x] Run Reports API
-- [x] Suppressions API
-- [x] Token API
-- [x] Vulnerability API
-
-### API v2
-
-- [x] Access Tokens
-- [x] Agent Access Tokens
-- [x] Alert Channels
-- [x] Alert Rules
-- [x] Audit Logs
-- [x] Cloud Accounts
-- [x] Cloud Activities
-- [x] Container Registries
-- [x] Contract Info
-- [x] Policies
-- [x] Queries
-- [x] Report Rules
-- [x] Resource Groups
-- [x] Schemas
-- [x] Team Members
-- [x] User Profile
-
 ### Contributing
 
 To install/configure the necessary requirements for contributing to this project, simply create a virtual environment, install `requirements.txt` and `requirements-dev.txt`, and set up a version file using the commands below:

diff --git a/examples/example_alert_channels.py b/examples/example_alert_channels.py
@@ -26,7 +26,7 @@
     lacework_client.alert_channels.get()
 
     # Search Alert Channels
-    lacework_client.alert_channels.search(query_data={
+    lacework_client.alert_channels.search(json={
         "filters": [
             {
                 "expression": "eq",

diff --git a/examples/example_audit_logs.py b/examples/example_audit_logs.py
@@ -36,7 +36,7 @@
     lacework_client.audit_logs.get(start_time=start_time, end_time=end_time)
 
     # Search Audit Logs
-    lacework_client.audit_logs.search(query_data={
+    lacework_client.audit_logs.search(json={
         "timeFilter": {
             "startTime": start_time,
             "endTime": end_time

diff --git a/examples/example_cloudtrail.py b/examples/example_cloudtrail.py
@@ -36,7 +36,7 @@
     lacework_client.cloudtrail.get(start_time=start_time, end_time=end_time)
 
     # Search CloudTrail
-    lacework_client.cloudtrail.search(query_data={
+    lacework_client.cloudtrail.search(json={
         "timeFilter": {
             "startTime": start_time,
             "endTime": end_time

diff --git a/laceworksdk/__init__.py b/laceworksdk/__init__.py
@@ -8,7 +8,7 @@
 import logging
 
 from .api import LaceworkClient  # noqa: F401
-from .exceptions import ApiError, laceworksdkException  # noqa: F401
+from .exceptions import ApiError, LaceworkSDKException  # noqa: F401
 
 # Initialize Package Logging
 logger = logging.getLogger(__name__)

diff --git a/laceworksdk/api/__init__.py b/laceworksdk/api/__init__.py
@@ -9,33 +9,41 @@
 import configparser
 
 from laceworksdk.http_session import HttpSession
-from .account import AccountAPI
-from .agent_access_tokens import AgentAccessTokensAPI
-from .alert_channels import AlertChannelsAPI
-from .alert_rules import AlertRulesAPI
-from .audit_logs import AuditLogsAPI
-from .cloud_accounts import CloudAccountsAPI
-from .cloud_activities import CloudActivitiesAPI
-from .compliance import ComplianceAPI
-from .container_registries import ContainerRegistriesAPI
-from .contract_info import ContractInfoAPI
-from .custom_compliance_config import CustomComplianceConfigAPI
-from .datasources import DatasourcesAPI
-from .download_file import DownloadFileAPI
-from .events import EventsAPI
-from .integrations import IntegrationsAPI
-from .policies import PoliciesAPI
-from .queries import QueriesAPI
-from .recommendations import RecommendationsAPI
-from .report_rules import ReportRulesAPI
-from .resource_groups import ResourceGroupsAPI
-from .run_reports import RunReportsAPI
-from .schemas import SchemasAPI
-from .suppressions import SuppressionsAPI
-from .team_members import TeamMembersAPI
-from .token import TokenAPI
-from .user_profile import UserProfileAPI
-from .vulnerability import VulnerabilityAPI
+
+from .v1.account import AccountAPI
+from .v1.compliance import ComplianceAPI
+from .v1.custom_compliance_config import CustomComplianceConfigAPI
+from .v1.download_file import DownloadFileAPI
+from .v1.events import EventsAPI
+from .v1.integrations import IntegrationsAPI
+from .v1.recommendations import RecommendationsAPI
+from .v1.run_reports import RunReportsAPI
+from .v1.suppressions import SuppressionsAPI
+from .v1.token import TokenAPI
+
+from .v2.activities import ActivitiesAPI
+from .v2.agent_access_tokens import AgentAccessTokensAPI
+from .v2.alert_channels import AlertChannelsAPI
+from .v2.alert_profiles import AlertProfilesAPI
+from .v2.alert_rules import AlertRulesAPI
+from .v2.alerts import AlertsAPI
+from .v2.audit_logs import AuditLogsAPI
+from .v2.cloud_accounts import CloudAccountsAPI
+from .v2.cloud_activities import CloudActivitiesAPI
+from .v2.configs import ConfigsAPI
+from .v2.container_registries import ContainerRegistriesAPI
+from .v2.contract_info import ContractInfoAPI
+from .v2.datasources import DatasourcesAPI
+from .v2.entities import EntitiesAPI
+from .v2.organization_info import OrganizationInfoAPI
+from .v2.policies import PoliciesAPI
+from .v2.queries import QueriesAPI
+from .v2.report_rules import ReportRulesAPI
+from .v2.resource_groups import ResourceGroupsAPI
+from .v2.schemas import SchemasAPI
+from .v2.team_members import TeamMembersAPI
+from .v2.user_profile import UserProfileAPI
+from .v2.vulnerabilities import VulnerabilitiesAPI
 
 from laceworksdk.config import (
     LACEWORK_ACCOUNT_ENVIRONMENT_VARIABLE,
@@ -50,7 +58,7 @@
 load_dotenv()
 
 
-class LaceworkClient(object):
+class LaceworkClient:
     """
     Lacework API wrapper for Python.
     """
@@ -87,28 +95,28 @@ def __init__(self,
             LACEWORK_API_BASE_DOMAIN_ENVIRONMENT_VARIABLE)
 
         config_file_path = os.path.join(
-            os.path.expanduser('~'), LACEWORK_CLI_CONFIG_RELATIVE_PATH)
+            os.path.expanduser("~"), LACEWORK_CLI_CONFIG_RELATIVE_PATH)
 
         if os.path.isfile(config_file_path):
             profile = profile or os.getenv(
-                LACEWORK_API_CONFIG_SECTION_ENVIRONMENT_VARIABLE, 'default')
+                LACEWORK_API_CONFIG_SECTION_ENVIRONMENT_VARIABLE, "default")
             config_obj = configparser.ConfigParser()
             config_obj.read([config_file_path])
             if config_obj.has_section(profile):
                 config_section = config_obj[profile]
-                api_key = config_section.get('api_key', '').strip('""')
+                api_key = config_section.get("api_key", "").strip('""')
                 if not self._api_key and api_key:
                     self._api_key = api_key
 
-                api_secret = config_section.get('api_secret', '').strip('""')
+                api_secret = config_section.get("api_secret", "").strip('""')
                 if not self._api_secret and api_secret:
                     self._api_secret = api_secret
 
-                account = config_section.get('account', '').strip('""')
+                account = config_section.get("account", "").strip('""')
                 if not self._account and account:
                     self._account = account
 
-                subaccount = config_section.get('subaccount', '').strip('""')
+                subaccount = config_section.get("subaccount", "").strip('""')
                 if not self._subaccount and subaccount:
                     self._subaccount = subaccount
 
@@ -123,20 +131,26 @@ def __init__(self,
 
         # API Wrappers
         self.account = AccountAPI(self._session)
+        self.activities = ActivitiesAPI(self._session)
         self.agent_access_tokens = AgentAccessTokensAPI(self._session)
         self.alert_channels = AlertChannelsAPI(self._session)
+        self.alert_profiles = AlertProfilesAPI(self._session)
         self.alert_rules = AlertRulesAPI(self._session)
+        self.alerts = AlertsAPI(self._session)
         self.audit_logs = AuditLogsAPI(self._session)
         self.cloud_accounts = CloudAccountsAPI(self._session)
         self.cloud_activities = CloudActivitiesAPI(self._session)
         self.compliance = ComplianceAPI(self._session)
         self.compliance.config = CustomComplianceConfigAPI(self._session)
+        self.configs = ConfigsAPI(self._session)
         self.container_registries = ContainerRegistriesAPI(self._session)
         self.contract_info = ContractInfoAPI(self._session)
         self.datasources = DatasourcesAPI(self._session)
+        self.entities = EntitiesAPI(self._session)
         self.events = EventsAPI(self._session)
         self.files = DownloadFileAPI(self._session)
         self.integrations = IntegrationsAPI(self._session)
+        self.organization_info = OrganizationInfoAPI(self._session)
         self.policies = PoliciesAPI(self._session)
         self.queries = QueriesAPI(self._session)
         self.recommendations = RecommendationsAPI(self._session)
@@ -148,4 +162,21 @@ def __init__(self,
         self.team_members = TeamMembersAPI(self._session)
         self.tokens = TokenAPI(self._session)
         self.user_profile = UserProfileAPI(self._session)
-        self.vulnerabilities = VulnerabilityAPI(self._session)
+        self.vulnerabilities = VulnerabilitiesAPI(self._session)
+
+    def set_org_level_access(self, org_level_access):
+        """
+        A method to set whether the client should use organization-level API calls.
+        """
+
+        if org_level_access is True:
+            self._session._org_level_access = True
+        else:
+            self._session._org_level_access = False
+
+    def set_subaccount(self, subaccount):
+        """
+        A method to update the subaccount the client should use for API calls.
+        """
+
+        self._session._subaccount = subaccount