qontract-reconcile 0.10.0__py3-none-any.whl → 0.10.1.dev1203__py3-none-any.whl

reconcile/utils/aws_api_typed/api.py

@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+import textwrap
+from abc import ABC, abstractmethod
+from functools import cached_property
+from typing import Any, TypeVar
+
+from boto3 import Session
+from botocore.client import BaseClient
+from pydantic import BaseModel
+
+import reconcile.utils.aws_api_typed.account
+import reconcile.utils.aws_api_typed.dynamodb
+import reconcile.utils.aws_api_typed.iam
+import reconcile.utils.aws_api_typed.organization
+import reconcile.utils.aws_api_typed.s3
+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.account import AWSApiAccount
+from reconcile.utils.aws_api_typed.dynamodb import AWSApiDynamoDB
+from reconcile.utils.aws_api_typed.iam import AWSApiIam
+from reconcile.utils.aws_api_typed.organization import AWSApiOrganizations
+from reconcile.utils.aws_api_typed.s3 import AWSApiS3
+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",
+    AWSApiAccount,
+    AWSApiDynamoDB,
+    AWSApiIam,
+    AWSApiOrganizations,
+    AWSApiS3,
+    AWSApiServiceQuotas,
+    AWSApiSts,
+    AWSApiSupport,
+)
+
+
+class AWSCredentials(ABC):
+    @abstractmethod
+    def as_env_vars(self) -> dict[str, str]:
+        """
+        Returns a dictionary of environment variables that can be used to authenticate with AWS.
+        """
+        ...
+
+    @abstractmethod
+    def as_credentials_file(self, profile_name: str = "default") -> str:
+        """
+        Returns a string that can be used to write an AWS credentials file.
+        """
+        ...
+
+    @abstractmethod
+    def build_session(self) -> Session:
+        """
+        Builds an AWS session using these credentials.
+        """
+        ...
+
+
+class AWSStaticCredentials(BaseModel, AWSCredentials):
+    """
+    A model representing AWS credentials.
+    """
+
+    access_key_id: str
+    secret_access_key: str
+    region: str
+
+    def as_env_vars(self) -> dict[str, str]:
+        return {
+            "AWS_ACCESS_KEY_ID": self.access_key_id,
+            "AWS_SECRET_ACCESS_KEY": self.secret_access_key,
+            "AWS_REGION": self.region,
+        }
+
+    def as_credentials_file(self, profile_name: str = "default") -> str:
+        return textwrap.dedent(
+            f"""\
+            [{profile_name}]
+            aws_access_key_id = {self.access_key_id}
+            aws_secret_access_key = {self.secret_access_key}
+            region = {self.region}
+            """
+        )
+
+    def build_session(self) -> Session:
+        return Session(
+            aws_access_key_id=self.access_key_id,
+            aws_secret_access_key=self.secret_access_key,
+            region_name=self.region,
+        )
+
+
+class AWSTemporaryCredentials(AWSStaticCredentials):
+    """
+    A model representing temporary AWS credentials.
+    """
+
+    session_token: str
+
+    def as_env_vars(self) -> dict[str, str]:
+        env_vars = super().as_env_vars()
+        env_vars["AWS_SESSION_TOKEN"] = self.session_token
+        return env_vars
+
+    def as_credentials_file(self, profile_name: str = "default") -> str:
+        return textwrap.dedent(
+            f"""\
+            [{profile_name}]
+            aws_access_key_id = {self.access_key_id}
+            aws_secret_access_key = {self.secret_access_key}
+            aws_session_token = {self.session_token}
+            region = {self.region}
+            """
+        )
+
+    def build_session(self) -> Session:
+        return Session(
+            aws_access_key_id=self.access_key_id,
+            aws_secret_access_key=self.secret_access_key,
+            aws_session_token=self.session_token,
+            region_name=self.region,
+        )
+
+
+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] = []
+
+    def __enter__(self) -> AWSApi:
+        return self
+
+    def __exit__(self, *exec: Any) -> None:
+        self.close()
+
+    def close(self) -> None:
+        """Close all clients created by this API instance."""
+        for client in self._session_clients:
+            client.close()
+        self._session_clients = []
+
+    def _init_sub_api(self, api_cls: type[SubApi]) -> SubApi:
+        """Return a new or cached sub api client."""
+        match api_cls:
+            case reconcile.utils.aws_api_typed.account.AWSApiAccount:
+                client = self.session.client("account")
+                api = api_cls(client)
+            case reconcile.utils.aws_api_typed.dynamodb.AWSApiDynamoDB:
+                client = self.session.client("dynamodb")
+                api = api_cls(client)
+            case reconcile.utils.aws_api_typed.iam.AWSApiIam:
+                client = self.session.client("iam")
+                api = api_cls(client)
+            case reconcile.utils.aws_api_typed.organization.AWSApiOrganizations:
+                client = self.session.client("organizations")
+                api = api_cls(client)
+            case reconcile.utils.aws_api_typed.s3.AWSApiS3:
+                client = self.session.client("s3")
+                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}")
+
+        self._session_clients.append(client)
+        return api
+
+    @cached_property
+    def account(self) -> AWSApiAccount:
+        """Return an AWS Acount Api client"""
+        return self._init_sub_api(AWSApiAccount)
+
+    @cached_property
+    def dynamodb(self) -> AWSApiDynamoDB:
+        """Return an AWS DynamoDB Api client"""
+        return self._init_sub_api(AWSApiDynamoDB)
+
+    @cached_property
+    def iam(self) -> AWSApiIam:
+        """Return an AWS IAM Api client."""
+        return self._init_sub_api(AWSApiIam)
+
+    @cached_property
+    def organizations(self) -> AWSApiOrganizations:
+        """Return an AWS Organizations Api client."""
+        return self._init_sub_api(AWSApiOrganizations)
+
+    @cached_property
+    def s3(self) -> AWSApiS3:
+        """Return an AWS S3 Api client."""
+        return self._init_sub_api(AWSApiS3)
+
+    @cached_property
+    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."""
+        credentials = self.sts.assume_role(account_id=account_id, role=role)
+        return AWSApi(
+            AWSTemporaryCredentials(
+                access_key_id=credentials.access_key_id,
+                secret_access_key=credentials.secret_access_key,
+                session_token=credentials.session_token,
+                region=self.session.region_name,
+            )
+        )
+
+    def temporary_session(self, duration_seconds: int = 900) -> AWSApi:
+        """Return a new AWSAPI with temporary AWS credentials from a session.
+
+        This is similar to assuming a role, in the sense that the credentials will expire after a certain amount of time.
+
+        These temporary credentials have the same permissions as the session they were built from, except:
+        - they can't be used for anything IAM related
+        - for the STS API only the AssumeRole and GetSessionToken actions are allowed
+        """
+        tmp_creds = self.sts.get_session_token(duration_seconds=duration_seconds)
+        return AWSApi(
+            AWSTemporaryCredentials(
+                access_key_id=tmp_creds.access_key_id,
+                secret_access_key=tmp_creds.secret_access_key,
+                session_token=tmp_creds.session_token,
+                region=self.session.region_name,
+            )
+        )

reconcile/utils/aws_api_typed/dynamodb.py

@@ -0,0 +1,16 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from mypy_boto3_dynamodb import DynamoDBClient
+else:
+    DynamoDBClient = object
+
+
+class AWSApiDynamoDB:
+    def __init__(self, client: DynamoDBClient) -> None:
+        self.client = client
+
+    @property
+    def boto3_client(self) -> DynamoDBClient:
+        """Gets the RAW boto3 DynamoDB client"""
+        return self.client

reconcile/utils/aws_api_typed/iam.py

@@ -0,0 +1,67 @@
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from mypy_boto3_iam import IAMClient
+else:
+    IAMClient = object
+
+
+class AWSAccessKey(BaseModel):
+    access_key_id: str = Field(..., alias="AccessKeyId")
+    secret_access_key: str = Field(..., alias="SecretAccessKey")
+
+
+class AWSUser(BaseModel):
+    user_name: str = Field(..., alias="UserName")
+    user_id: str = Field(..., alias="UserId")
+    arn: str = Field(..., alias="Arn")
+    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
+
+    def create_access_key(self, user_name: str) -> AWSAccessKey:
+        """Create an access key for a given user."""
+        credentials = self.client.create_access_key(
+            UserName=user_name,
+        )
+        return AWSAccessKey(**credentials["AccessKey"])
+
+    def create_user(self, user_name: str) -> AWSUser:
+        """Create a new IAM 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"
+            ) from None
+
+    def attach_user_policy(self, user_name: str, policy_arn: str) -> None:
+        """Attach a policy to a user."""
+        self.client.attach_user_policy(
+            UserName=user_name,
+            PolicyArn=policy_arn,
+        )
+
+    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!"
+                ) from None

reconcile/utils/aws_api_typed/organization.py

@@ -0,0 +1,152 @@
+import functools
+from collections.abc import Iterable, Mapping
+from typing import TYPE_CHECKING, Optional
+
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from mypy_boto3_organizations import OrganizationsClient
+    from mypy_boto3_organizations.literals import CreateAccountFailureReasonType
+else:
+    OrganizationsClient = object
+    CreateAccountFailureReasonType = object
+
+
+class AwsOrganizationOU(BaseModel):
+    id: str = Field(..., alias="Id")
+    arn: str = Field(..., alias="Arn")
+    name: str = Field(..., alias="Name")
+    children: list["AwsOrganizationOU"] = []
+
+    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")
+    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):
+    """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(
+        self, root: AwsOrganizationOU | None = None
+    ) -> AwsOrganizationOU:
+        """List all organizational units for a given root recursively."""
+        if not root:
+            root = AwsOrganizationOU(**self.client.list_roots()["Roots"][0])
+
+        paginator = self.client.get_paginator("list_organizational_units_for_parent")
+        for page in paginator.paginate(ParentId=root.id):
+            for ou_raw in page["OrganizationalUnits"]:
+                ou = AwsOrganizationOU(**ou_raw)
+                root.children.append(ou)
+                self.get_organizational_units_tree(root=ou)
+        return root
+
+    def create_account(
+        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=name,
+            IamUserAccessToBilling="ALLOW" if access_to_billing else "DENY",
+        )
+        status = AWSAccountStatus(**resp["CreateAccountStatus"])
+        if status.state == "FAILED":
+            raise AWSAccountCreationException(
+                f"Account creation failed: {status.failure_reason}"
+            )
+        return status
+
+    def describe_create_account_status(
+        self, create_account_request_id: str
+    ) -> AWSAccountStatus:
+        """Return the status of a create account request."""
+        resp = self.client.describe_create_account_status(
+            CreateAccountRequestId=create_account_request_id
+        )
+        return AWSAccountStatus(**resp["CreateAccountStatus"])
+
+    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=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/s3.py

@@ -0,0 +1,26 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from mypy_boto3_s3 import S3Client
+    from mypy_boto3_s3.literals import BucketLocationConstraintType
+else:
+    S3Client = BucketLocationConstraintType = object
+
+
+class AWSApiS3:
+    def __init__(self, client: S3Client) -> None:
+        self.client = client
+
+    def create_bucket(self, name: str, region: str) -> str:
+        """Create an S3 bucket without any ACLs therefore the creator will be the owner and returns the ARN."""
+        bucket_kwargs = {}
+        if region != "us-east-1":
+            # you can't specify the location if it's us-east-1 :(
+            # see valid values "LocationConstraint" here: https://docs.aws.amazon.com/AmazonS3/latest/API/API_CreateBucketConfiguration.html
+            bucket_kwargs = {
+                "CreateBucketConfiguration": {
+                    "LocationConstraint": region,
+                },
+            }
+        self.client.create_bucket(Bucket=name, **bucket_kwargs)  # type: ignore
+        return f"arn:aws:s3:::{name}"

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."
+            ) from None
+
+    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."
+            ) from None

reconcile/utils/aws_api_typed/sts.py

@@ -0,0 +1,36 @@
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from mypy_boto3_sts import STSClient
+else:
+    STSClient = object
+
+
+class AWSCredentials(BaseModel):
+    access_key_id: str = Field(..., alias="AccessKeyId")
+    secret_access_key: str = Field(..., alias="SecretAccessKey")
+    session_token: str = Field(..., alias="SessionToken")
+    expiration: datetime = Field(..., alias="Expiration")
+
+
+class AWSApiSts:
+    def __init__(self, client: STSClient) -> None:
+        self.client = client
+
+    def assume_role(self, account_id: str, role: str) -> AWSCredentials:
+        """Assume a role and return temporary credentials."""
+        assumed_role_object = self.client.assume_role(
+            RoleArn=f"arn:aws:iam::{account_id}:role/{role}",
+            RoleSessionName=role,
+        )
+        return AWSCredentials(**assumed_role_object["Credentials"])
+
+    def get_session_token(self, duration_seconds: int = 900) -> AWSCredentials:
+        """Return temporary credentials."""
+        assumed_role_object = self.client.get_session_token(
+            DurationSeconds=duration_seconds
+        )
+        return AWSCredentials(**assumed_role_object["Credentials"])