qontract-reconcile 0.10.1rc701__py3-none-any.whl → 0.10.1rc702__py3-none-any.whl

reconcile/utils/aws_api_typed/api.py

@@ -11,10 +11,14 @@ from pydantic import BaseModel
 
 import reconcile.utils.aws_api_typed.iam
 import reconcile.utils.aws_api_typed.organization
+import reconcile.utils.aws_api_typed.service_quotas
 import reconcile.utils.aws_api_typed.sts
+import reconcile.utils.aws_api_typed.support
 from reconcile.utils.aws_api_typed.iam import AWSApiIam
 from reconcile.utils.aws_api_typed.organization import AWSApiOrganizations
+from reconcile.utils.aws_api_typed.service_quotas import AWSApiServiceQuotas
 from reconcile.utils.aws_api_typed.sts import AWSApiSts
+from reconcile.utils.aws_api_typed.support import AWSApiSupport
 
 SubApi = TypeVar("SubApi")
 
@@ -109,6 +113,29 @@ class AWSTemporaryCredentials(AWSStaticCredentials):
 
 
 class AWSApi:
+    """High-level API for AWS services.
+
+    This class provides a high-level API for AWS services like
+
+    * IAM
+    * Organizations
+    * Service Quotas
+    * STS
+    * Support
+
+    It also provides a way to assume roles and create temporary sessions.
+
+    Example:
+
+    with AWSApi(AWSStaticCredentials(...)) as api:
+        with api.assume_role(... role="MyRole") as role_api:
+            role_api.iam.create_user(...)
+
+
+    This new fully-tested and fully-typed API will replace the old one in the near future.
+    Feel free to implement missing methods and AWS servcies as needed.
+    """
+
     def __init__(self, aws_credentials: AWSCredentials) -> None:
         self.session = aws_credentials.build_session()
         self._session_clients: list[BaseClient] = []
@@ -134,9 +161,15 @@ class AWSApi:
             case reconcile.utils.aws_api_typed.organization.AWSApiOrganizations:
                 client = self.session.client("organizations")
                 api = api_cls(client)
+            case reconcile.utils.aws_api_typed.service_quotas.AWSApiServiceQuotas:
+                client = self.session.client("service-quotas")
+                api = api_cls(client)
             case reconcile.utils.aws_api_typed.sts.AWSApiSts:
                 client = self.session.client("sts")
                 api = api_cls(client)
+            case reconcile.utils.aws_api_typed.support.AWSApiSupport:
+                client = self.session.client("support")
+                api = api_cls(client)
             case _:
                 raise ValueError(f"Unknown API class: {api_cls}")
 
@@ -144,9 +177,9 @@ class AWSApi:
         return api
 
     @cached_property
-    def sts(self) -> AWSApiSts:
-        """Return an AWS STS Api client."""
-        return self._init_sub_api(AWSApiSts)
+    def iam(self) -> AWSApiIam:
+        """Return an AWS IAM Api client."""
+        return self._init_sub_api(AWSApiIam)
 
     @cached_property
     def organizations(self) -> AWSApiOrganizations:
@@ -154,9 +187,19 @@ class AWSApi:
         return self._init_sub_api(AWSApiOrganizations)
 
     @cached_property
-    def iam(self) -> AWSApiIam:
-        """Return an AWS IAM Api client."""
-        return self._init_sub_api(AWSApiIam)
+    def service_quotas(self) -> AWSApiServiceQuotas:
+        """Return an AWS Service Quotas Api client."""
+        return self._init_sub_api(AWSApiServiceQuotas)
+
+    @cached_property
+    def sts(self) -> AWSApiSts:
+        """Return an AWS STS Api client."""
+        return self._init_sub_api(AWSApiSts)
+
+    @cached_property
+    def support(self) -> AWSApiSupport:
+        """Return an AWS Support Api client."""
+        return self._init_sub_api(AWSApiSupport)
 
     def assume_role(self, account_id: str, role: str) -> AWSApi:
         """Return a new AWSApi with the assumed role."""

reconcile/utils/aws_api_typed/iam.py

@@ -20,6 +20,10 @@ class AWSUser(BaseModel):
     path: str = Field(..., alias="Path")
 
 
+class AWSEntityAlreadyExistsException(Exception):
+    """Raised when the user already exists in IAM."""
+
+
 class AWSApiIam:
     def __init__(self, client: IAMClient) -> None:
         self.client = client
@@ -33,10 +37,11 @@ class AWSApiIam:
 
     def create_user(self, user_name: str) -> AWSUser:
         """Create a new IAM user."""
-        user = self.client.create_user(
-            UserName=user_name,
-        )
-        return AWSUser(**user["User"])
+        try:
+            user = self.client.create_user(UserName=user_name)
+            return AWSUser(**user["User"])
+        except self.client.exceptions.EntityAlreadyExistsException:
+            raise AWSEntityAlreadyExistsException(f"User {user_name} already exists")
 
     def attach_user_policy(self, user_name: str, policy_arn: str) -> None:
         """Attach a policy to a user."""
@@ -45,6 +50,16 @@ class AWSApiIam:
             PolicyArn=policy_arn,
         )
 
-    def create_account_alias(self, account_alias: str) -> None:
-        """Create an account alias."""
-        self.client.create_account_alias(AccountAlias=account_alias)
+    def get_account_alias(self) -> str:
+        """Get the account alias."""
+        return self.client.list_account_aliases()["AccountAliases"][0]
+
+    def set_account_alias(self, account_alias: str) -> None:
+        """Set the account alias."""
+        try:
+            self.client.create_account_alias(AccountAlias=account_alias)
+        except self.client.exceptions.EntityAlreadyExistsException:
+            if self.get_account_alias() != account_alias:
+                raise ValueError(
+                    "Account alias already exists for another AWS account. Choose another one!"
+                )

reconcile/utils/aws_api_typed/organization.py

@@ -1,5 +1,6 @@
-from collections.abc import Mapping
-from typing import TYPE_CHECKING
+import functools
+from collections.abc import Iterable, Mapping
+from typing import TYPE_CHECKING, Optional
 
 from pydantic import BaseModel, Field
 
@@ -17,38 +18,65 @@ class AwsOrganizationOU(BaseModel):
     name: str = Field(..., alias="Name")
     children: list["AwsOrganizationOU"] = []
 
-    def find(self, path: str) -> "AwsOrganizationOU":
-        """Return an organizational unit by its path."""
-        name, *rest = path.strip("/").split("/")
-        subs = "/".join(rest)
-        if self.name == name:
-            if not rest:
-                return self
-            for child in self.children:
-                try:
-                    return child.find(subs)
-                except KeyError:
-                    pass
-        raise KeyError(f"OU not found: {path}")
+    def locate(
+        self, path: list[str], ignore_case: bool = True
+    ) -> Optional["AwsOrganizationOU"]:
+        name, *sub = path
+        match = self.name.lower() == name.lower() if ignore_case else self.name == name
+        if not match:
+            return None
+        if not sub:
+            return self
+        return next(
+            (
+                result
+                for child in self.children
+                if (result := child.locate(sub, ignore_case=ignore_case))
+            ),
+            None,
+        )
+
+    def find(self, path: str, ignore_case: bool = True) -> "AwsOrganizationOU":
+        node = self.locate(path.strip("/").split("/"), ignore_case=ignore_case)
+        if not node:
+            raise KeyError(f"OU not found: {path}")
+        return node
+
+    def __hash__(self) -> int:
+        return hash(self.id)
 
 
 class AWSAccountStatus(BaseModel):
     id: str = Field(..., alias="Id")
-    account_name: str = Field(..., alias="AccountName")
-    account_id: str = Field(..., alias="AccountId")
+    name: str = Field(..., alias="AccountName")
+    uid: str | None = Field(alias="AccountId")
     state: str = Field(..., alias="State")
     failure_reason: CreateAccountFailureReasonType | None = Field(alias="FailureReason")
 
 
+class AWSAccount(BaseModel):
+    name: str = Field(..., alias="Name")
+    uid: str = Field(..., alias="Id")
+    email: str = Field(..., alias="Email")
+    state: str = Field(..., alias="Status")
+
+
 class AWSAccountCreationException(Exception):
-    pass
+    """Exception raised when account creation failed."""
+
+
+class AWSAccountNotFoundException(Exception):
+    """Exception raised when the account cannot be found in the specified OU."""
 
 
 class AWSApiOrganizations:
     def __init__(self, client: OrganizationsClient) -> None:
         self.client = client
+        self.get_organizational_units_tree = functools.lru_cache(maxsize=None)(
+            self._get_organizational_units_tree
+        )
 
-    def get_organizational_units_tree(
+    def _get_organizational_units_tree(
         self, root: AwsOrganizationOU | None = None
     ) -> AwsOrganizationOU:
         """List all organizational units for a given root recursively."""
@@ -64,18 +92,13 @@ class AWSApiOrganizations:
         return root
 
     def create_account(
-        self,
-        email: str,
-        account_name: str,
-        tags: Mapping[str, str],
-        access_to_billing: bool = True,
+        self, email: str, name: str, access_to_billing: bool = True
     ) -> AWSAccountStatus:
         """Create a new account in the organization."""
         resp = self.client.create_account(
             Email=email,
-            AccountName=account_name,
+            AccountName=name,
             IamUserAccessToBilling="ALLOW" if access_to_billing else "DENY",
-            Tags=[{"Key": k, "Value": v} for k, v in tags.items()],
         )
         status = AWSAccountStatus(**resp["CreateAccountStatus"])
         if status.state == "FAILED":
@@ -93,12 +116,37 @@ class AWSApiOrganizations:
         )
         return AWSAccountStatus(**resp["CreateAccountStatus"])
 
-    def move_account(
-        self, account_id: str, source_parent_id: str, destination_parent_id: str
-    ) -> None:
+    def get_ou(self, uid: str) -> str:
+        """Return the organizational unit ID of an account."""
+        resp = self.client.list_parents(ChildId=uid)
+        for p in resp.get("Parents", []):
+            if p["Type"] in {"ORGANIZATIONAL_UNIT", "ROOT"}:
+                return p["Id"]
+        raise AWSAccountNotFoundException(f"Account {uid} not found!")
+
+    def move_account(self, uid: str, destination_parent_id: str) -> None:
         """Move an account to a different organizational unit."""
+        source_parent_id = self.get_ou(uid=uid)
+        if source_parent_id == destination_parent_id:
+            return
         self.client.move_account(
-            AccountId=account_id,
+            AccountId=uid,
             SourceParentId=source_parent_id,
             DestinationParentId=destination_parent_id,
         )
+
+    def describe_account(self, uid: str) -> AWSAccount:
+        """Return the status of an account."""
+        resp = self.client.describe_account(AccountId=uid)
+        return AWSAccount(**resp["Account"])
+
+    def tag_resource(self, resource_id: str, tags: Mapping[str, str]) -> None:
+        """Tag a resource."""
+        self.client.tag_resource(
+            ResourceId=resource_id,
+            Tags=[{"Key": k, "Value": v} for k, v in tags.items()],
+        )
+
+    def untag_resource(self, resource_id: str, tag_keys: Iterable[str]) -> None:
+        """Untag a resource."""
+        self.client.untag_resource(ResourceId=resource_id, TagKeys=list(tag_keys))

reconcile/utils/aws_api_typed/service_quotas.py

@@ -0,0 +1,79 @@
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from mypy_boto3_service_quotas import ServiceQuotasClient
+    from mypy_boto3_service_quotas.literals import RequestStatusType
+else:
+    ServiceQuotasClient = RequestStatusType = object
+
+
+class AWSRequestedServiceQuotaChange(BaseModel):
+    id: str = Field(..., alias="Id")
+    status: RequestStatusType = Field(..., alias="Status")
+    service_code: str = Field(..., alias="ServiceCode")
+    quota_code: str = Field(..., alias="QuotaCode")
+    desired_value: float = Field(..., alias="DesiredValue")
+
+
+class AWSQuota(BaseModel):
+    service_code: str = Field(..., alias="ServiceCode")
+    service_name: str = Field(..., alias="ServiceName")
+    quota_code: str = Field(..., alias="QuotaCode")
+    quota_name: str = Field(..., alias="QuotaName")
+    value: float = Field(..., alias="Value")
+
+    def __str__(self) -> str:
+        return f"{self.service_name=} {self.service_code=} {self.quota_name=} {self.quota_code=}: {self.value=}"
+
+    def __repr__(self) -> str:
+        return str(self)
+
+
+class AWSNoSuchResourceException(Exception):
+    """Raised when a resource is not found in a service quotas API call."""
+
+
+class AWSResourceAlreadyExistsException(Exception):
+    """Raised when quota increase request already exists."""
+
+
+class AWSApiServiceQuotas:
+    def __init__(self, client: ServiceQuotasClient) -> None:
+        self.client = client
+
+    def get_requested_service_quota_change(
+        self, request_id: str
+    ) -> AWSRequestedServiceQuotaChange:
+        """Return the requested service quota change."""
+        req = self.client.get_requested_service_quota_change(RequestId=request_id)
+        return AWSRequestedServiceQuotaChange(**req["RequestedQuota"])
+
+    def request_service_quota_change(
+        self, service_code: str, quota_code: str, desired_value: float
+    ) -> AWSRequestedServiceQuotaChange:
+        """Request a service quota change."""
+        try:
+            req = self.client.request_service_quota_increase(
+                ServiceCode=service_code,
+                QuotaCode=quota_code,
+                DesiredValue=desired_value,
+            )
+            return AWSRequestedServiceQuotaChange(**req["RequestedQuota"])
+        except self.client.exceptions.ResourceAlreadyExistsException:
+            raise AWSResourceAlreadyExistsException(
+                f"Service quota increase request {service_code=}, {quota_code=} already exists."
+            )
+
+    def get_service_quota(self, service_code: str, quota_code: str) -> AWSQuota:
+        """Return the current value of the service quota."""
+        try:
+            quota = self.client.get_service_quota(
+                ServiceCode=service_code, QuotaCode=quota_code
+            )
+            return AWSQuota(**quota["Quota"])
+        except self.client.exceptions.NoSuchResourceException:
+            raise AWSNoSuchResourceException(
+                f"Service quota {service_code=}, {quota_code=} not found."
+            )

reconcile/utils/aws_api_typed/support.py

@@ -0,0 +1,79 @@
+from enum import Enum
+from typing import TYPE_CHECKING, Literal
+
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from mypy_boto3_support import SupportClient
+else:
+    SupportClient = object
+
+
+class AWSCase(BaseModel):
+    case_id: str = Field(..., alias="caseId")
+    subject: str
+    status: str
+
+
+class SUPPORT_PLAN(Enum):
+    BASIC = "basic"
+    DEVELOPER = "developer"
+    BUSINESS = "business"
+    ENTERPRISE = "enterprise"
+
+
+class AWSApiSupport:
+    def __init__(self, client: SupportClient) -> None:
+        self.client = client
+
+    def create_case(
+        self,
+        subject: str,
+        message: str,
+        category: str = "other-account-issues",
+        service: str = "customer-account",
+        issue_type: Literal["customer-service", "technical"] = "customer-service",
+        language: Literal["en", "zh", "ja", "ko"] = "en",
+        severity: str = "high",
+    ) -> str:
+        """Create a support case and return the case id."""
+        case = self.client.create_case(
+            subject=subject,
+            communicationBody=message,
+            categoryCode=category,
+            serviceCode=service,
+            issueType=issue_type,
+            language=language,
+            severityCode=severity,
+        )
+        return case["caseId"]
+
+    def describe_case(self, case_id: str) -> AWSCase:
+        """Return the status of a support case."""
+        case = self.client.describe_cases(caseIdList=[case_id])["cases"][0]
+        return AWSCase(**case)
+
+    def get_support_level(self) -> SUPPORT_PLAN:
+        """Return the support level of the account."""
+
+        try:
+            response = self.client.describe_severity_levels(language="en")
+        except self.client.exceptions.ClientError as err:
+            if err.response["Error"]["Code"] == "SubscriptionRequiredException":
+                return SUPPORT_PLAN.BASIC
+            raise err
+
+        severity_levels = {
+            level["code"].lower() for level in response["severityLevels"]
+        }
+        if "critical" in severity_levels:
+            return SUPPORT_PLAN.ENTERPRISE
+        if "urgent" in severity_levels:
+            return SUPPORT_PLAN.BUSINESS
+        if "high" in severity_levels:
+            return SUPPORT_PLAN.BUSINESS
+        if "normal" in severity_levels:
+            return SUPPORT_PLAN.DEVELOPER
+        if "low" in severity_levels:
+            return SUPPORT_PLAN.DEVELOPER
+        return SUPPORT_PLAN.BASIC

reconcile/utils/state.py

@@ -1,9 +1,10 @@
+import contextlib
 import json
 import logging
 import os
 from abc import abstractmethod
-from collections.abc import Callable, Mapping
-from types import TracebackType
+from collections.abc import Callable, Generator, Mapping
+from dataclasses import dataclass, field
 from typing import (
     Any,
     Optional,
@@ -212,6 +213,10 @@ def acquire_state_settings(
     )
 
 
+class AbortStateTransaction(Exception):
+    """Raise to abort a state transaction."""
+
+
 class State:
     """
     A state object to be used by stateful integrations.
@@ -405,52 +410,51 @@ class State:
     def __setitem__(self, key: str, value: Any) -> None:
         self._set(key, value)
 
-    def transaction(self, key: str, value: Any) -> "_TransactionContext":
+    @contextlib.contextmanager
+    def transaction(
+        self, key: str, value: Any = None
+    ) -> Generator["TransactionStateObj", None, None]:
         """Get a context manager to set the key in the state if no exception occurs.
 
+        You can set the value either via the value parameter or by setting the value attribute of the returned object.
+        If both are provided, the value attribute of the state object will take precedence.
+
         Attention!
 
         This is not a locking mechanism. It is a way to ensure that a key is set in the state if no exception occurs.
         This method is not thread-safe nor multi-process-safe! There is no locking mechanism in place.
         """
-        return _TransactionContext(self, key, value)
+        try:
+            _current_value = self[key]
+        except KeyError:
+            _current_value = None
+        state_obj = TransactionStateObj(key, value=_current_value)
+        try:
+            yield state_obj
+        except AbortStateTransaction:
+            return
+        else:
+            if state_obj.changed and state_obj.value != _current_value:
+                self[state_obj.key] = state_obj.value
+            elif value is not None and state_obj.value != value:
+                self[state_obj.key] = value
 
 
-class _TransactionContext:
-    """A context manager to set a key in the state if no exception occurs."""
+@dataclass
+class TransactionStateObj:
+    """Represents a transistion state object with a key and a value."""
 
-    def __init__(
-        self,
-        state: State,
-        key: str,
-        value: Any,
-    ):
-        self.state = state
-        self.key = key
-        self.value = value
+    key: str
+    value: Any = None
+    _init_value: Any = field(init=False)
 
-    def __enter__(self) -> bool:
-        """Return True if the key exists in the state, False otherwise.
+    def __post_init__(self) -> None:
+        self._init_value = self.value
 
-        Cache the previous value to avoid unnecessary updates.
-        """
-        self._previous_value = None
-        try:
-            self._previous_value = self.state[self.key]
-            return True
-        except KeyError:
-            return False
+    @property
+    def changed(self) -> bool:
+        return self.value != self._init_value
 
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> None:
-        if exc_type:
-            # if an exception occurred, we don't want to write to the state
-            return
-        if self._previous_value == self.value:
-            # if the value didn't change, we don't want to write to the state
-            return
-        self.state[self.key] = self.value
+    @property
+    def exists(self) -> bool:
+        return self._init_value is not None