cartography 0.93.0rc1__py3-none-any.whl → 0.123.0__py3-none-any.whl

cartography/intel/airbyte/users.py

@@ -0,0 +1,106 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Tuple
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.models.airbyte.user import AirbyteUserSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: AirbyteClient,
+    org_id: str,
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    users = get(api_session, org_id=org_id)
+    for user in users:
+        permissions = get_permissions(api_session, user["id"], org_id=org_id)
+        org_admin, workspace_admin, workspace_member = transform_permissions(
+            permissions
+        )
+        user["adminOfOrganization"] = org_admin
+        user["adminOfWorkspace"] = workspace_admin
+        user["memberOfWorkspace"] = workspace_member
+    load_users(neo4j_session, users, org_id, common_job_parameters["UPDATE_TAG"])
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: AirbyteClient,
+    org_id: str,
+) -> List[Dict[str, Any]]:
+    return api_session.get("/users", params={"organizationId": org_id})
+
+
+@timeit
+def get_permissions(
+    api_session: AirbyteClient,
+    user_id: str,
+    org_id: str,
+) -> List[Dict[str, Any]]:
+    return api_session.get(
+        "/permissions",
+        params={"organizationId": org_id, "userId": user_id},
+    )
+
+
+@timeit
+def transform_permissions(
+    permissions: List[Dict[str, Any]],
+) -> Tuple[
+    List[str],  # org_admin
+    List[str],  # workspace_admin
+    List[str],  # workspace_member
+]:
+    org_admin: list[str] = []
+    workspace_admin: list[str] = []
+    workspace_member: list[str] = []
+
+    for permission in permissions:
+        # workspace
+        if permission["scope"] == "workspace":
+            if permission["permissionType"] in ("workspace_owner", "workspace_admin"):
+                workspace_admin.append(permission["scopeId"])
+            workspace_member.append(permission["scopeId"])
+        # organization
+        elif permission["scope"] == "organization":
+            if permission["permissionType"] in ("organization_admin",):
+                org_admin.append(permission["scopeId"])
+    return org_admin, workspace_admin, workspace_member
+
+
+@timeit
+def load_users(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    org_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Users into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteUserSchema(),
+        data,
+        ORG_ID=org_id,
+        lastupdated=update_tag,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AirbyteUserSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/airbyte/util.py

@@ -0,0 +1,122 @@
+import logging
+import time
+from typing import Any
+
+import requests
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+class AirbyteClient:
+    """A client for interacting with the Airbyte API.
+    This client handles authentication and provides methods to make GET requests to the Airbyte API.
+    It automatically handles pagination for GET requests that return multiple pages of data.
+    """
+
+    def __init__(self, base_url: str, client_id: str, client_secret: str) -> None:
+        self._client_id = client_id
+        self._client_secret = client_secret
+        self.base_url = base_url
+        self._access_token_expiry: int | None = None
+        self._session = requests.Session()
+
+    def get(self, uri: str, params: dict | None = None, offset: int = 0) -> list[dict]:
+        """Make a GET request to the Airbyte API.
+        This method handles authentication and pagination.
+        Args:
+            uri (str): The URI to make the GET request to.
+            params (dict | None): Optional parameters to include in the request.
+            offset (int): The offset for pagination, defaults to 0.
+        Returns:
+            list[dict]: A list of dictionaries containing the data from the response.
+        """
+        self.authenticate()
+        if params is None:
+            params_with_pagination = {}
+        else:
+            params_with_pagination = params.copy()
+        params_with_pagination["offset"] = offset
+        response = self._session.get(
+            f"{self.base_url}{uri}", params=params_with_pagination, timeout=_TIMEOUT
+        )
+        response.raise_for_status()
+        data = response.json().get("data")
+        if response.json().get("next", "") != "":
+            data.extend(
+                self.get(
+                    uri,
+                    params=params,
+                    offset=offset + len(data),
+                )
+            )
+        return data
+
+    def authenticate(self) -> None:
+        """Authenticate with the Airbyte API using client credentials.
+        This method checks if the access token is still valid and renews it if necessary.
+        If the access token is expired or not set, it will make a request to obtain a new access token.
+        """
+        if self._access_token_expiry and self._access_token_expiry >= time.time():
+            return
+        self._session.headers.pop("Authorization", None)
+        payload = {
+            "grant-type": "client_credentials",
+            "client_id": self._client_id,
+            "client_secret": self._client_secret,
+        }
+        response = self._session.post(
+            f"{self.base_url}/applications/token", json=payload, timeout=_TIMEOUT
+        )
+        response.raise_for_status()
+        data = response.json()
+        self._session.headers["Authorization"] = (
+            f"Bearer {data.get('access_token', '')}"
+        )
+        token_expiry = data.get("expires_in", 0)
+        self._access_token_expiry = time.time() + token_expiry
+        logger.debug("Access token renewed, expires in %s seconds.", token_expiry)
+
+
+def normalize_airbyte_config(config: dict[str, Any]) -> dict[str, Any]:
+    """Normalize the Airbyte configuration dictionary.
+    This function takes a configuration dictionary and normalizes it by mapping keys to a standard set of keys.
+    This is useful for ensuring consistency across different configurations, and will allow to connect to existing nodes.
+    Args:
+        config (dict[str, Any]): The configuration dictionary to normalize.
+    Returns:
+        dict[str, Any]: A normalized configuration dictionary with standardized keys.
+    """
+    normalized_config = {}
+    for key in config:
+        if key in ("host", "port", "name", "region", "endpoint", "account", ""):
+            normalized_config[key] = config[key]
+        elif key in ("aws_region_name", "region_name", "s3_bucket_region"):
+            normalized_config["region"] = config[key]
+        elif key in ("queue_url", "url", "s3_endpoint"):
+            normalized_config["endpoint"] = config[key]
+        elif key in ("azure_blob_storage_account_name", "storage_account_name"):
+            normalized_config["account"] = config[key]
+        elif key in (
+            "azure_blob_storage_container_name",
+            "bucket",
+            "database",
+            "s3_bucket_name",
+        ):
+            normalized_config["name"] = config[key]
+    return normalized_config
+
+
+def list_to_string(lst: list[str]) -> str | None:
+    """Convert a list of strings to a comma-separated string."""
+    if len(lst) == 0:
+        return None
+    # Sublist
+    formated_list: list[str] = []
+    for item in lst:
+        if isinstance(item, list):
+            formated_list.append("|".join(item))
+        else:
+            formated_list.append(str(item))
+    return ",".join(formated_list)

cartography/intel/airbyte/workspaces.py

@@ -0,0 +1,63 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.airbyte.util import AirbyteClient
+from cartography.models.airbyte.workspace import AirbyteWorkspaceSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: AirbyteClient,
+    org_id: str,
+    common_job_parameters: Dict[str, Any],
+) -> List[Dict[str, Any]]:
+    workspaces = get(api_session, org_id)
+    load_workspaces(
+        neo4j_session, workspaces, org_id, common_job_parameters["UPDATE_TAG"]
+    )
+    cleanup(neo4j_session, common_job_parameters)
+    return workspaces
+
+
+@timeit
+def get(
+    api_session: AirbyteClient,
+    org_id: str,
+) -> List[Dict[str, Any]]:
+    return api_session.get("/workspaces", params={"organizationId": org_id})
+
+
+@timeit
+def load_workspaces(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    org_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Airbyte Workspaces into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AirbyteWorkspaceSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=org_id,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AirbyteWorkspaceSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/analysis.py

@@ -39,4 +39,7 @@ def run(neo4j_session: neo4j.Session, config: Config) -> None:
         except (KeyboardInterrupt, SystemExit):
             raise
         except Exception:
-            logger.exception("An exception occurred while executing discovered analysis job: %s", path)
+            logger.exception(
+                "An exception occurred while executing discovered analysis job: %s",
+                path,
+            )

cartography/intel/anthropic/__init__.py

@@ -0,0 +1,62 @@
+import logging
+
+import neo4j
+import requests
+
+import cartography.intel.anthropic.apikeys
+import cartography.intel.anthropic.users
+import cartography.intel.anthropic.workspaces
+from cartography.config import Config
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def start_anthropic_ingestion(neo4j_session: neo4j.Session, config: Config) -> None:
+    """
+    If this module is configured, perform ingestion of Anthropic data. Otherwise warn and exit
+    :param neo4j_session: Neo4J session for database interface
+    :param config: A cartography.config object
+    :return: None
+    """
+
+    if not config.anthropic_apikey:
+        logger.info(
+            "Anthropic import is not configured - skipping this module. "
+            "See docs to configure.",
+        )
+        return
+
+    # Create requests sessions
+    api_session = requests.session()
+    api_session.headers.update(
+        {
+            "X-Api-Key": config.anthropic_apikey,
+            "anthropic-version": "2023-06-01",
+        }
+    )
+
+    common_job_parameters = {
+        "UPDATE_TAG": config.update_tag,
+        "BASE_URL": "https://api.anthropic.com/v1",
+    }
+
+    # Organization node is created during the users sync
+    cartography.intel.anthropic.users.sync(
+        neo4j_session,
+        api_session,
+        common_job_parameters,
+    )
+
+    cartography.intel.anthropic.workspaces.sync(
+        neo4j_session,
+        api_session,
+        common_job_parameters,
+    )
+
+    cartography.intel.anthropic.apikeys.sync(
+        neo4j_session,
+        api_session,
+        common_job_parameters,
+    )

cartography/intel/anthropic/apikeys.py

@@ -0,0 +1,72 @@
+import logging
+from typing import Any
+from typing import Tuple
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.anthropic.util import paginated_get
+from cartography.models.anthropic.apikey import AnthropicApiKeySchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: requests.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    org_id, apikeys = get(
+        api_session,
+        common_job_parameters["BASE_URL"],
+    )
+    common_job_parameters["ORG_ID"] = org_id
+    load_apikeys(
+        neo4j_session,
+        apikeys,
+        org_id,
+        common_job_parameters["UPDATE_TAG"],
+    )
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: requests.Session,
+    base_url: str,
+) -> Tuple[str, list[dict[str, Any]]]:
+    return paginated_get(
+        api_session, f"{base_url}/organizations/api_keys", timeout=_TIMEOUT
+    )
+
+
+@timeit
+def load_apikeys(
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    ORG_ID: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Anthropic ApiKey into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AnthropicApiKeySchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=ORG_ID,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AnthropicApiKeySchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/anthropic/users.py

@@ -0,0 +1,75 @@
+import logging
+from typing import Any
+from typing import Tuple
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.anthropic.util import paginated_get
+from cartography.models.anthropic.organization import AnthropicOrganizationSchema
+from cartography.models.anthropic.user import AnthropicUserSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: requests.Session,
+    common_job_parameters: dict[str, Any],
+) -> str:
+    org_id, users = get(
+        api_session,
+        common_job_parameters["BASE_URL"],
+    )
+    common_job_parameters["ORG_ID"] = org_id
+    load_users(neo4j_session, users, org_id, common_job_parameters["UPDATE_TAG"])
+    cleanup(neo4j_session, common_job_parameters)
+    return org_id
+
+
+@timeit
+def get(
+    api_session: requests.Session,
+    base_url: str,
+) -> Tuple[str, list[dict[str, Any]]]:
+    return paginated_get(
+        api_session, f"{base_url}/organizations/users", timeout=_TIMEOUT
+    )
+
+
+@timeit
+def load_users(
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    ORG_ID: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        AnthropicOrganizationSchema(),
+        [{"id": ORG_ID}],
+        lastupdated=update_tag,
+    )
+    logger.info("Loading %d Anthropic User into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AnthropicUserSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=ORG_ID,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AnthropicUserSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/anthropic/util.py

@@ -0,0 +1,51 @@
+from typing import Any
+
+import requests
+
+
+def paginated_get(
+    api_session: requests.Session,
+    url: str,
+    timeout: tuple[int, int],
+    after: str | None = None,
+) -> tuple[str, list[dict[str, Any]]]:
+    """Helper function to get paginated data from the Anthropic API.
+
+    This function handles the pagination of the API requests and returns
+    the results in a list. It also retrieves the organization ID from the
+    response headers. The function will continue to make requests until
+    all pages of data have been retrieved. The results are returned as a
+    list of dictionaries, where each dictionary represents a single
+    entity.
+
+    Args:
+        api_session (requests.Session): The requests session to use for making API calls.
+        url (str): The URL to make the API call to.
+        timeout (tuple[int, int]): The timeout for the API call.
+        after (str | None): The ID of the last item retrieved in the previous request.
+            If None, the first page of results will be retrieved.
+    Returns:
+        tuple[str, list[dict[str, Any]]]: A tuple containing the organization ID and a list of
+            dictionaries representing the results.
+    """
+    results: list[dict[str, Any]] = []
+    params = {"after_id": after} if after else {}
+    req = api_session.get(
+        url,
+        params=params,
+        timeout=timeout,
+    )
+    req.raise_for_status()
+    # Get organization_id from the headers
+    organization_id = req.headers.get("anthropic-organization-id", "")
+    result = req.json()
+    results.extend(result.get("data", []))
+    if result.get("has_more"):
+        _, next_results = paginated_get(
+            api_session,
+            url,
+            timeout=timeout,
+            after=result.get("last_id"),
+        )
+        results.extend(next_results)
+    return organization_id, results

cartography/intel/anthropic/workspaces.py

@@ -0,0 +1,95 @@
+import logging
+from typing import Any
+from typing import Tuple
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.anthropic.util import paginated_get
+from cartography.models.anthropic.workspace import AnthropicWorkspaceSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: requests.Session,
+    common_job_parameters: dict[str, Any],
+) -> list[dict]:
+    org_id, workspaces = get(
+        api_session,
+        common_job_parameters["BASE_URL"],
+    )
+    common_job_parameters["ORG_ID"] = org_id
+    for workspace in workspaces:
+        workspace["users"] = []
+        workspace["admins"] = []
+        for user in get_workspace_users(
+            api_session,
+            common_job_parameters["BASE_URL"],
+            workspace["id"],
+        ):
+            workspace["users"].append(user["user_id"])
+            if user["workspace_role"] == "workspace_admin":
+                workspace["admins"].append(user["user_id"])
+    load_workspaces(
+        neo4j_session, workspaces, org_id, common_job_parameters["UPDATE_TAG"]
+    )
+    cleanup(neo4j_session, common_job_parameters)
+    return workspaces
+
+
+@timeit
+def get(
+    api_session: requests.Session,
+    base_url: str,
+) -> Tuple[str, list[dict[str, Any]]]:
+    return paginated_get(
+        api_session, f"{base_url}/organizations/workspaces", timeout=_TIMEOUT
+    )
+
+
+@timeit
+def get_workspace_users(
+    api_session: requests.Session,
+    base_url: str,
+    workspace_id: str,
+) -> list[dict[str, Any]]:
+    _, result = paginated_get(
+        api_session,
+        f"{base_url}/organizations/workspaces/{workspace_id}/members",
+        timeout=_TIMEOUT,
+    )
+    return result
+
+
+@timeit
+def load_workspaces(
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    ORG_ID: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d Anthropic workspaces into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        AnthropicWorkspaceSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=ORG_ID,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(AnthropicWorkspaceSchema(), common_job_parameters).run(
+        neo4j_session
+    )