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

cartography/util.py

@@ -1,9 +1,11 @@
 import asyncio
 import logging
 import re
-import sys
 from functools import partial
 from functools import wraps
+from importlib.resources import open_binary
+from importlib.resources import read_text
+from itertools import islice
 from string import Template
 from typing import Any
 from typing import Awaitable
@@ -15,6 +17,7 @@ from typing import Iterable
 from typing import List
 from typing import Optional
 from typing import Set
+from typing import Type
 from typing import TypeVar
 from typing import Union
 
@@ -22,32 +25,45 @@ import backoff
 import boto3
 import botocore
 import neo4j
+from botocore.exceptions import EndpointConnectionError
+from botocore.parsers import ResponseParserError
 
 from cartography.graph.job import GraphJob
 from cartography.graph.statement import get_job_shortname
 from cartography.stats import get_stats_client
 from cartography.stats import ScopedStatsClient
 
+logger = logging.getLogger(__name__)
 
-if sys.version_info >= (3, 7):
-    from importlib.resources import open_binary, read_text
-else:
-    from importlib_resources import open_binary, read_text
 
-logger = logging.getLogger(__name__)
+def is_service_control_policy_explicit_deny(
+    error: botocore.exceptions.ClientError,
+) -> bool:
+    """Return True if the ClientError was caused by an explicit service control policy deny."""
+    error_code = error.response.get("Error", {}).get("Code")
+    if error_code not in {"AccessDenied", "AccessDeniedException"}:
+        return False
+
+    message = error.response.get("Error", {}).get("Message")
+    if not message:
+        return False
+
+    lowered = message.lower()
+    return "explicit deny" in lowered and "service control policy" in lowered
 
 
 STATUS_SUCCESS = 0
 STATUS_FAILURE = 1
 STATUS_KEYBOARD_INTERRUPT = 130
 DEFAULT_BATCH_SIZE = 1000
+DEFAULT_MAX_PAGES = 10000
 
 
 def run_analysis_job(
     filename: str,
     neo4j_session: neo4j.Session,
     common_job_parameters: Dict,
-    package: str = 'cartography.data.jobs.analysis',
+    package: str = "cartography.data.jobs.analysis",
 ) -> None:
     """
     Enriches existing graph data with analysis jobs. This is designed for use with the sync stage
@@ -69,11 +85,11 @@ def run_analysis_job(
 
 
 def run_analysis_and_ensure_deps(
-        analysis_job_name: str,
-        resource_dependencies: Set[str],
-        requested_syncs: Set[str],
-        common_job_parameters: Dict[str, Any],
-        neo4j_session: neo4j.Session,
+    analysis_job_name: str,
+    resource_dependencies: Set[str],
+    requested_syncs: Set[str],
+    common_job_parameters: Dict[str, Any],
+    neo4j_session: neo4j.Session,
 ) -> None:
     """
     Runs analysis job only if the given set of resource dependencies was included in the requested_syncs.
@@ -103,7 +119,7 @@ def run_scoped_analysis_job(
     filename: str,
     neo4j_session: neo4j.Session,
     common_job_parameters: Dict,
-    package: str = 'cartography.data.jobs.scoped_analysis',
+    package: str = "cartography.data.jobs.scoped_analysis",
 ) -> None:
     """
     Enriches existing graph data scoped to a given sub resource - e.g. the current AWS account.
@@ -119,8 +135,10 @@ def run_scoped_analysis_job(
 
 
 def run_cleanup_job(
-    filename: str, neo4j_session: neo4j.Session, common_job_parameters: Dict,
-    package: str = 'cartography.data.jobs.cleanup',
+    filename: str,
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+    package: str = "cartography.data.jobs.cleanup",
 ) -> None:
     GraphJob.run_from_json(
         neo4j_session,
@@ -141,7 +159,7 @@ def merge_module_sync_metadata(
     update_tag: int,
     stat_handler: ScopedStatsClient,
 ) -> None:
-    '''
+    """
     This creates `ModuleSyncMetadata` nodes when called from each of the individual modules or sub-modules.
     The 'types' used here should be actual node labels. For example, if we did sync a particular AWSAccount's S3Buckets,
     the `grouptype` is 'AWSAccount', the `groupid` is the particular account's `id`, and the `syncedtype` is 'S3Bucket'.
@@ -151,8 +169,12 @@ def merge_module_sync_metadata(
     :param group_id: The parent module's id
     :param synced_type: The sub-module's type
     :param update_tag: Timestamp used to determine data freshness
-    '''
-    template = Template("""
+    """
+    # Import here to avoid circular import with cartography.client.core.tx
+    from cartography.client.core.tx import run_write_query
+
+    template = Template(
+        """
         MERGE (n:ModuleSyncMetadata{id:'${group_type}_${group_id}_${synced_type}'})
         ON CREATE SET
             n:SyncMetadata, n.firstseen=timestamp()
@@ -160,20 +182,26 @@ def merge_module_sync_metadata(
             n.grouptype='${group_type}',
             n.groupid='${group_id}',
             n.lastupdated=$UPDATE_TAG
-    """)
-    neo4j_session.run(
-        template.safe_substitute(group_type=group_type, group_id=group_id, synced_type=synced_type),
+    """,
+    )
+    run_write_query(
+        neo4j_session,
+        template.safe_substitute(
+            group_type=group_type,
+            group_id=group_id,
+            synced_type=synced_type,
+        ),
         UPDATE_TAG=update_tag,
     )
-    stat_handler.incr(f'{group_type}_{group_id}_{synced_type}_lastupdated', update_tag)
+    stat_handler.incr(f"{group_type}_{group_id}_{synced_type}_lastupdated", update_tag)
 
 
 def load_resource_binary(package: str, resource_name: str) -> BinaryIO:
     return open_binary(package, resource_name)
 
 
-R = TypeVar('R')
-F = TypeVar('F', bound=Callable[..., Any])
+R = TypeVar("R")
+F = TypeVar("F", bound=Callable[..., Any])
 
 
 def timeit(method: F) -> F:
@@ -182,6 +210,7 @@ def timeit(method: F) -> F:
     This is only active if config.statsd_enabled is True.
     :param method: The function to measure execution
     """
+
     # Allow access via `inspect` to the wrapped function. This is used in integration tests to standardize param names.
     @wraps(method)
     def timed(*args, **kwargs):  # type: ignore
@@ -203,42 +232,79 @@ def aws_paginate(
     client: boto3.client,
     method_name: str,
     object_name: str,
+    max_pages: int | None = DEFAULT_MAX_PAGES,
     **kwargs: Any,
-) -> List[Dict]:
-    '''
+) -> Iterable[Dict]:
+    """
     Helper method for boilerplate boto3 pagination
     The **kwargs will be forwarded to the paginator
-    '''
+    """
     paginator = client.get_paginator(method_name)
-    items = []
-    i = 0
     for i, page in enumerate(paginator.paginate(**kwargs), start=1):
         if i % 100 == 0:
-            logger.info(f'fetching page number {i}')
+            logger.info(f"fetching page number {i}")
         if object_name in page:
-            items.extend(page[object_name])
+            items = page[object_name]
+            yield from items
         else:
             logger.warning(
-                f'''aws_paginate: Key "{object_name}" is not present, check if this is a typo.
-If not, then the AWS datatype somehow does not have this key.''',
+                f"""aws_paginate: Key "{object_name}" is not present, check if this is a typo.
+If not, then the AWS datatype somehow does not have this key.""",
             )
-    return items
+        if max_pages is not None and i >= max_pages:
+            logger.warning(f"Reached max batch size of {max_pages} pages")
+            break
 
 
-AWSGetFunc = TypeVar('AWSGetFunc', bound=Callable[..., List])
+AWSGetFunc = TypeVar("AWSGetFunc", bound=Callable[..., Iterable])
 
 # fix for AWS TooManyRequestsException
-# https://github.com/lyft/cartography/issues/297
-# https://github.com/lyft/cartography/issues/243
-# https://github.com/lyft/cartography/issues/65
-# https://github.com/lyft/cartography/issues/25
+# https://github.com/cartography-cncf/cartography/issues/297
+# https://github.com/cartography-cncf/cartography/issues/243
+# https://github.com/cartography-cncf/cartography/issues/65
+# https://github.com/cartography-cncf/cartography/issues/25
 
 
 def backoff_handler(details: Dict) -> None:
     """
-    Handler that will be executed on exception by backoff mechanism
+    Handler that will be executed on exception by backoff mechanism.
+
+    The backoff library may provide partial details (e.g. ``wait`` can be ``None`` when a
+    retry is triggered immediately). Format the message defensively so logging never raises.
     """
-    logger.warning("Backing off {wait:0.1f} seconds after {tries} tries. Calling function {target}".format(**details))
+    wait = details.get("wait")
+    if isinstance(wait, (int, float)):
+        wait_display = f"{wait:0.1f}"
+    elif wait is None:
+        wait_display = "unknown"
+    else:
+        wait_display = str(wait)
+
+    tries = details.get("tries")
+    tries_display = str(tries) if tries is not None else "unknown"
+
+    target = details.get("target", "<unknown>")
+
+    logger.warning(
+        "Backing off %s seconds after %s tries. Calling function %s",
+        wait_display,
+        tries_display,
+        target,
+    )
+
+
+# Error codes that indicate a service is unavailable in a region or blocked by policies
+AWS_REGION_ACCESS_DENIED_ERROR_CODES = [
+    "AccessDenied",
+    "AccessDeniedException",
+    "AuthFailure",
+    "AuthorizationError",
+    "AuthorizationErrorException",
+    "InvalidClientTokenId",
+    "UnauthorizedOperation",
+    "UnrecognizedClientException",
+    "InternalServerErrorException",
+]
 
 
 # TODO Move this to cartography.intel.aws.util.common
@@ -252,25 +318,16 @@ def aws_handle_regions(func: AWSGetFunc) -> AWSGetFunc:
 
     This should be used on `get_` functions that normally return a list of items.
     """
-    ERROR_CODES = [
-        'AccessDenied',
-        'AccessDeniedException',
-        'AuthFailure',
-        'InvalidClientTokenId',
-        'UnauthorizedOperation',
-        'UnrecognizedClientException',
-        'InternalServerErrorException',
-    ]
 
     @wraps(func)
     # fix for AWS TooManyRequestsException
-    # https://github.com/lyft/cartography/issues/297
-    # https://github.com/lyft/cartography/issues/243
-    # https://github.com/lyft/cartography/issues/65
-    # https://github.com/lyft/cartography/issues/25
+    # https://github.com/cartography-cncf/cartography/issues/297
+    # https://github.com/cartography-cncf/cartography/issues/243
+    # https://github.com/cartography-cncf/cartography/issues/65
+    # https://github.com/cartography-cncf/cartography/issues/25
     @backoff.on_exception(
         backoff.expo,
-        botocore.exceptions.ClientError,
+        (botocore.exceptions.ClientError, ResponseParserError),
         max_time=600,
         on_backoff=backoff_handler,
     )
@@ -278,16 +335,65 @@ def aws_handle_regions(func: AWSGetFunc) -> AWSGetFunc:
         try:
             return func(*args, **kwargs)
         except botocore.exceptions.ClientError as e:
+            error_code = e.response.get("Error", {}).get("Code")
+            if error_code == "InvalidToken":
+                raise RuntimeError(
+                    "AWS returned an InvalidToken error. Configure regional STS endpoints by "
+                    "setting environment variable AWS_STS_REGIONAL_ENDPOINTS=regional or adding "
+                    "'sts_regional_endpoints = regional' to your AWS config file."
+                ) from e
             # The account is not authorized to use this service in this region
             # so we can continue without raising an exception
-            if e.response['Error']['Code'] in ERROR_CODES:
-                logger.warning("{} in this region. Skipping...".format(e.response['Error']['Message']))
+            if error_code in AWS_REGION_ACCESS_DENIED_ERROR_CODES:
+                error_message = e.response.get("Error", {}).get("Message")
+                if is_service_control_policy_explicit_deny(e):
+                    logger.warning(
+                        "Service control policy denied access while calling %s: %s",
+                        func.__name__,
+                        error_message,
+                    )
+                else:
+                    logger.warning(
+                        "{} in this region. Skipping...".format(
+                            error_message,
+                        ),
+                    )
                 return []
             else:
                 raise
+        except EndpointConnectionError:
+            logger.warning(
+                "Encountered an EndpointConnectionError. This means that the AWS "
+                "resource is not available in this region. Skipping.",
+            )
+            return []
+
     return cast(AWSGetFunc, inner_function)
 
 
+def retries_with_backoff(
+    func: Callable,
+    exception_type: Type[Exception],
+    max_tries: int,
+    on_backoff: Callable,
+) -> Callable:
+    """
+    Adds retry with backoff to the given function.  (Could expand the possible input parameters as needed.)
+    """
+
+    @wraps(func)
+    @backoff.on_exception(
+        backoff.expo,
+        exception_type,
+        max_tries=max_tries,
+        on_backoff=on_backoff,
+    )
+    def inner_function(*args, **kwargs):  # type: ignore
+        return func(*args, **kwargs)
+
+    return cast(Callable, inner_function)
+
+
 def dict_value_to_str(obj: Dict, key: str) -> Optional[str]:
     """
     Convert the value referenced by the key in the dict to a string, if it exists, and return it. If it doesn't exist,
@@ -313,32 +419,30 @@ def dict_date_to_epoch(obj: Dict, key: str) -> Optional[int]:
 
 
 def camel_to_snake(name: str) -> str:
-    return re.sub(r'(?<!^)(?=[A-Z])', '_', name).lower()
+    return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
 
 
-def batch(items: Iterable, size: int = DEFAULT_BATCH_SIZE) -> List[List]:
-    '''
-    Takes an Iterable of items and returns a list of lists of the same items,
+def batch(items: Iterable, size: int = DEFAULT_BATCH_SIZE) -> Iterable[List[Any]]:
+    """
+    Takes an Iterable of items and returns a Generator of lists of the same items,
      batched into chunks of the provided `size`.
 
     Use:
     x = [1,2,3,4,5,6,7,8]
-    batch(x, size=3) -> [[1, 2, 3], [4, 5, 6], [7, 8]]
-    '''
-    items = list(items)
-    return [
-        items[i: i + size]
-        for i in range(0, len(items), size)
-    ]
+    batch(x, size=3) -> Iterator yielding [1, 2, 3], [4, 5, 6], [7, 8]
+    """
+    it = iter(items)
+    while chunk := list(islice(it, size)):
+        yield chunk
 
 
 def is_throttling_exception(exc: Exception) -> bool:
-    '''
+    """
     Returns True if the exception is caused by a client libraries throttling mechanism
-    '''
+    """
     # https://boto3.amazonaws.com/v1/documentation/api/1.19.9/guide/error-handling.html
     if isinstance(exc, botocore.exceptions.ClientError):
-        if exc.response['Error']['Code'] in ['LimitExceededException', 'Throttling']:
+        if exc.response["Error"]["Code"] in ["LimitExceededException", "Throttling"]:
             return True
     # add other exceptions here, if needed, like:
     # https://cloud.google.com/python/docs/reference/storage/1.39.0/retry_timeout#configuring-retries
@@ -348,7 +452,7 @@ def is_throttling_exception(exc: Exception) -> bool:
 
 
 def to_asynchronous(func: Callable[..., R], *args: Any, **kwargs: Any) -> Awaitable[R]:
-    '''
+    """
     Returns a Future that will run a function and its arguments in the default threadpool.
     Helper until we start using python 3.9's asyncio.to_thread
 
@@ -378,8 +482,12 @@ def to_asynchronous(func: Callable[..., R], *args: Any, **kwargs: Any) -> Awaita
     NOTE: to use this in a Jupyter notebook, you need to do:
     # import nest_asyncio
     # nest_asyncio.apply()
-    '''
-    CartographyThrottlingException = type('CartographyThrottlingException', (Exception,), {})
+    """
+    CartographyThrottlingException = type(
+        "CartographyThrottlingException",
+        (Exception,),
+        {},
+    )
 
     @wraps(func)
     def wrapper(*args: Any, **kwargs: Any) -> R:
@@ -391,13 +499,15 @@ def to_asynchronous(func: Callable[..., R], *args: Any, **kwargs: Any) -> Awaita
             raise
 
     # don't use @backoff as decorator, to preserve typing
-    wrapped = backoff.on_exception(backoff.expo, CartographyThrottlingException)(wrapper)
+    wrapped = backoff.on_exception(backoff.expo, CartographyThrottlingException)(
+        wrapper,
+    )
     call = partial(wrapped, *args, **kwargs)
     return asyncio.get_event_loop().run_in_executor(None, call)
 
 
 def to_synchronous(*awaitables: Awaitable[Any]) -> List[Any]:
-    '''
+    """
     Synchronously waits for the Awaitable(s) to complete and returns their result(s).
     See https://docs.python.org/3.8/library/asyncio-task.html#asyncio-awaitables
 
@@ -419,5 +529,5 @@ def to_synchronous(*awaitables: Awaitable[Any]) -> List[Any]:
     future_2 = another_async_func(2)
 
     results = to_synchronous(future_1, future_2)
-    '''
+    """
     return asyncio.get_event_loop().run_until_complete(asyncio.gather(*awaitables))