Use version 2 STS tokens to query opt-in regions

botocove/cove_decorator.py

@@ -1,5 +1,6 @@
 import functools
 import logging
+from itertools import chain
 from typing import Any, Callable, Dict, List, Optional
 from warnings import warn
 
@@ -64,26 +65,35 @@ def wrapper(*args: Any, **kwargs: Any) -> CoveOutput:
                 thread_workers=thread_workers,
             )
 
-            output = runner.run_cove_function()
+            function_output = runner.run_cove_function()
 
-            # Rewrite dataclasses into untyped dicts to retain current functionality
-            return CoveOutput(
+            cove_output = CoveOutput(
                 Results=[
                     {k: v for k, v in r.items() if v is not None}
-                    for r in output["Results"]
+                    for r in function_output["Results"]
                 ],
                 Exceptions=[
                     {k: v for k, v in e.items() if v is not None}
-                    for e in output["Exceptions"]
+                    for e in function_output["Exceptions"]
                     if e["AssumeRoleSuccess"] is True
                 ],
                 FailedAssumeRole=[
                     {k: v for k, v in f.items() if v is not None}
-                    for f in output["Exceptions"]
+                    for f in function_output["Exceptions"]
                     if f["AssumeRoleSuccess"] is False
                 ],
             )
 
+            if regions is None:
+                for result in chain(
+                    cove_output["Results"],
+                    cove_output["Exceptions"],
+                    cove_output["FailedAssumeRole"],
+                ):
+                    del result["Region"]
+
+            return cove_output
+
         return wrapper
 
     # Handle both bare decorator and with argument

botocove/cove_host_account.py

@@ -1,26 +1,12 @@
 import logging
 import re
 from functools import lru_cache
-from typing import (
-    Any,
-    Dict,
-    Iterable,
-    List,
-    Literal,
-    Optional,
-    Sequence,
-    Set,
-    Tuple,
-    Union,
-)
-
-import boto3
+from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple
+
 from boto3.session import Session
 from botocore.config import Config
 from botocore.exceptions import ClientError
-from mypy_boto3_organizations.client import OrganizationsClient
 from mypy_boto3_organizations.type_defs import AccountTypeDef
-from mypy_boto3_sts.client import STSClient
 from mypy_boto3_sts.type_defs import PolicyDescriptorTypeTypeDef
 
 from botocove.cove_types import CoveSessionInformation
@@ -52,15 +38,33 @@ def __init__(
 
         self.thread_workers = thread_workers
 
-        self.sts_client = self._get_boto3_sts_client(assuming_session)
-        self.org_client = self._get_boto3_org_client(assuming_session)
+        if assuming_session:
+            logger.info(f"Using provided Boto3 session {assuming_session}")
+            self._session = assuming_session
+        else:
+            logger.info("No Boto3 session argument: using credential chain")
+            self._session = Session()
+
+        # Obtain version 2 session tokens valid in all AWS regions.
+        # https://repost.aws/knowledge-center/iam-validate-access-credentials
+        self.sts_client = self._session.client(
+            service_name="sts",
+            region_name=self._session.region_name,
+            endpoint_url=f"https://sts.{self._session.region_name}.amazonaws.com",
+            config=Config(max_pool_connections=self.thread_workers),
+        )
+
+        self.org_client = self._session.client(
+            service_name="organizations",
+            config=Config(max_pool_connections=self.thread_workers),
+        )
 
         caller_id = self.sts_client.get_caller_identity()
         self.host_account_id = caller_id["Account"]
         self.host_account_partition = caller_id["Arn"].split(":")[1]
 
         if regions is None:
-            self.target_regions = [None]
+            self.target_regions = [self._session.region_name]
         else:
             self.target_regions = regions
 
@@ -147,35 +151,6 @@ def _generate_account_sessions(self) -> Iterable[CoveSessionInformation]:
                         Result=None,
                     )
 
-    def _get_boto3_client(
-        self,
-        clientname: Union[Literal["organizations"], Literal["sts"]],
-        assuming_session: Optional[Session],
-    ) -> Any:
-        if assuming_session:
-            logger.info(f"Using provided Boto3 session {assuming_session}")
-            return assuming_session.client(
-                service_name=clientname,
-                config=Config(max_pool_connections=self.thread_workers),
-            )
-        logger.info("No Boto3 session argument: using credential chain")
-        return boto3.client(
-            service_name=clientname,
-            config=Config(max_pool_connections=self.thread_workers),
-        )
-
-    def _get_boto3_org_client(
-        self, assuming_session: Optional[Session]
-    ) -> OrganizationsClient:
-        client: OrganizationsClient = self._get_boto3_client(
-            "organizations", assuming_session
-        )
-        return client
-
-    def _get_boto3_sts_client(self, assuming_session: Optional[Session]) -> STSClient:
-        client: STSClient = self._get_boto3_client("sts", assuming_session)
-        return client
-
     def _resolve_target_accounts(self, target_ids: Optional[List[str]]) -> Set[str]:
         accounts_to_ignore = self._gather_ignored_accounts()
         logger.info(f"Ignoring account IDs: {accounts_to_ignore=}")

pyproject.toml

@@ -32,7 +32,8 @@ flake8-print = ">=5.0.0"
 mypy = ">=1.5.1"
 pre-commit = ">=3.4.0"
 flakeheaven = ">=3.3.0"
-moto = {extras = ["organizations", "sts"], version = ">=4.2.5"}
+moto = {extras = ["organizations", "sts", "ec2"], version = ">=4.2.5"}
+boto3-stubs = {extras = ["ec2"], version = "*"}
 pytest-randomly = ">=3.15.0"
 
 # These are the last versions compatible with flake8 4. flakeheaven 3.3.0 is

tests/conftest.py

@@ -5,7 +5,7 @@
 from _pytest.logging import LogCaptureFixture
 from _pytest.monkeypatch import MonkeyPatch
 from boto3 import Session
-from moto import mock_organizations, mock_sts
+from moto import mock_ec2, mock_organizations, mock_sts
 
 from tests.moto_mock_org.moto_models import LargeOrg, SmallOrg
 
@@ -27,7 +27,7 @@ def _clean_env(monkeypatch: MonkeyPatch) -> None:
 @pytest.fixture()
 def mock_session() -> Iterator[Session]:
     """Returns a session with mock AWS services."""
-    with mock_sts(), mock_organizations():
+    with mock_sts(), mock_organizations(), mock_ec2():
         yield Session()