qontract-reconcile 0.10.1rc696__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/merge_request_manager/merge_request_manager.py

@@ -0,0 +1,102 @@
+import logging
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+from gitlab.v4.objects import ProjectMergeRequest
+from pydantic import BaseModel
+
+from reconcile.utils.merge_request_manager.parser import (
+    Parser,
+    ParserError,
+    ParserVersionError,
+)
+from reconcile.utils.vcs import VCS
+
+T = TypeVar("T", bound=BaseModel)
+
+
+@dataclass
+class OpenMergeRequest(Generic[T]):
+    raw: ProjectMergeRequest
+    mr_info: T
+
+
+class MergeRequestManagerBase(Generic[T]):
+    """ """
+
+    def __init__(self, vcs: VCS, parser: Parser, mr_label: str):
+        self._vcs = vcs
+        self._parser = parser
+        self._mr_label = mr_label
+        self._open_mrs: list[OpenMergeRequest] = []
+        self._open_mrs_with_problems: list[OpenMergeRequest] = []
+        self._housekeeping_ran = False
+
+    @abstractmethod
+    def create_merge_request(self, data: Any) -> None:
+        pass
+
+    def _merge_request_already_exists(
+        self,
+        expected_data: dict[str, Any],
+    ) -> OpenMergeRequest | None:
+        for mr in self._open_mrs:
+            mr_info_dict = mr.mr_info.dict()
+            if all(
+                mr_info_dict.get(k) == expected_data.get(k)
+                for k in expected_data.keys()
+            ):
+                return mr
+
+        return None
+
+    def _fetch_managed_open_merge_requests(self) -> list[ProjectMergeRequest]:
+        all_open_mrs = self._vcs.get_open_app_interface_merge_requests()
+        return [mr for mr in all_open_mrs if self._mr_label in mr.labels]
+
+    def housekeeping(self) -> None:
+        """
+        Close bad MRs:
+        - bad description format
+        - wrong version
+        - merge conflict
+
+        --> if we update the template output, we automatically close
+        old open MRs and replace them with new ones.
+        """
+        for mr in self._fetch_managed_open_merge_requests():
+            attrs = mr.attributes
+            desc = attrs.get("description")
+            has_conflicts = attrs.get("has_conflicts", False)
+            if has_conflicts:
+                logging.info(
+                    "Merge-conflict detected. Closing %s",
+                    mr.attributes.get("web_url", "NO_WEBURL"),
+                )
+                self._vcs.close_app_interface_mr(
+                    mr, "Closing this MR because of a merge-conflict."
+                )
+                continue
+            try:
+                mr_info = self._parser.parse(description=desc)
+            except ParserVersionError:
+                logging.info(
+                    "Old MR version detected! Closing %s",
+                    mr.attributes.get("web_url", "NO_WEBURL"),
+                )
+                self._vcs.close_app_interface_mr(
+                    mr, "Closing this MR because it has an outdated integration version"
+                )
+                continue
+            except ParserError:
+                logging.info(
+                    "Bad MR description format. Closing %s",
+                    mr.attributes.get("web_url", "NO_WEBURL"),
+                )
+                self._vcs.close_app_interface_mr(
+                    mr, "Closing this MR because of bad description format."
+                )
+                continue
+            self._open_mrs.append(OpenMergeRequest(raw=mr, mr_info=mr_info))
+        self._housekeeping_ran = True

reconcile/utils/oauth2_backend_application_session.py

@@ -0,0 +1,102 @@
+import threading
+from collections.abc import Mapping
+from typing import Any, Self
+
+from oauthlib.oauth2 import BackendApplicationClient, TokenExpiredError
+from requests import Response
+from requests_oauthlib import OAuth2Session
+
+FETCH_TOKEN_HEADERS = {
+    "Accept": "application/json",
+    "Content-Type": "application/x-www-form-urlencoded;charset=UTF-8",
+    "Connection": "close",  # Close connection to avoid RemoteDisconnected ConnectionError
+}
+
+
+class OAuth2BackendApplicationSession:
+    """
+    OAuth2 session using Backend Application flow with auto and thread-safe token fetch.
+    """
+
+    def __init__(
+        self,
+        client_id: str,
+        client_secret: str,
+        token_url: str,
+        scope: list[str] | None = None,
+    ) -> None:
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.token_url = token_url
+        self.scope = scope
+        client = BackendApplicationClient(client_id=client_id)
+        self.session = OAuth2Session(client=client, scope=scope)
+        self.token_lock = threading.Lock()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self.session.close()
+
+    def fetch_token(self) -> dict:
+        """
+        Fetch token from token_url and store it in the session.
+        Thread-safe method to avoid multiple threads fetching the token at the same time.
+        """
+        token = self.session.token
+        with self.token_lock:
+            # Check if token is already fetched by another thread
+            if token is not self.session.token:
+                return self.session.token
+            return self.session.fetch_token(
+                token_url=self.token_url,
+                client_id=self.client_id,
+                client_secret=self.client_secret,
+                scope=self.scope,
+                headers=FETCH_TOKEN_HEADERS,
+            )
+
+    def request(
+        self,
+        method: str,
+        url: str,
+        data: Any = None,
+        headers: Mapping[str, str] | None = None,
+        withhold_token: bool = False,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        **kwargs: Any,
+    ) -> Response:
+        """
+        Make a request using OAuth2 session, compatible with the OAuth2Session.request method.
+        Auto fetch token if never fetched before or if the token is expired.
+        """
+        if not self.session.authorized:
+            self.fetch_token()
+        try:
+            return self.session.request(
+                method=method,
+                url=url,
+                data=data,
+                headers=headers,
+                withhold_token=withhold_token,
+                client_id=client_id,
+                client_secret=client_secret,
+                **kwargs,
+            )
+        except TokenExpiredError:
+            self.fetch_token()
+            return self.session.request(
+                method=method,
+                url=url,
+                data=data,
+                headers=headers,
+                withhold_token=withhold_token,
+                client_id=client_id,
+                client_secret=client_secret,
+                **kwargs,
+            )