aws-annoying 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

aws_annoying/cli/session_manager/port_forward.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 
+import logging
 import os
 import signal
 import subprocess
 from pathlib import Path  # noqa: TC003
 
 import typer
-from rich import print  # noqa: A004
+
+from aws_annoying.utils.ec2 import get_instance_id_by_name
 
 from ._app import session_manager_app
-from ._common import SessionManager, get_instance_id_by_name
+from ._common import SessionManager
+
+logger = logging.getLogger(__name__)
 
 
 # https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html
@@ -59,30 +63,30 @@ def port_forward(  # noqa: PLR0913
     # Check if the PID file already exists
     if pid_file.exists():
         if not terminate_running_process:
-            print("🚫 PID file already exists.")
+            logger.error("PID file already exists.")
             raise typer.Exit(1)
 
         pid_content = pid_file.read_text()
         try:
             existing_pid = int(pid_content)
         except ValueError:
-            print(f"🚫 PID file content is invalid; expected integer, but got: {type(pid_content)}")
+            logger.error("PID file content is invalid; expected integer, but got: %r", type(pid_content))  # noqa: TRY400
             raise typer.Exit(1) from None
 
         try:
-            print(f"⚠️ Terminating running process with PID {existing_pid}.")
+            logger.warning("Terminating running process with PID %d.", existing_pid)
             os.kill(existing_pid, signal.SIGTERM)
             pid_file.write_text("")  # Clear the PID file
         except ProcessLookupError:
-            print(f"⚠️ Tried to terminate process with PID {existing_pid} but does not exist.")
+            logger.warning("Tried to terminate process with PID %d but does not exist.", existing_pid)
 
     # Resolve the instance name or ID
     instance_id = get_instance_id_by_name(through)
     if instance_id:
-        print(f"❗ Instance ID resolved: [bold]{instance_id}[/bold]")
+        logger.info("Instance ID resolved: [bold]%s[/bold]", instance_id)
         target = instance_id
     else:
-        print(f"🚫 Instance with name '{through}' not found.")
+        logger.info("Instance with name '%s' not found.", through)
         raise typer.Exit(1)
 
     # Initiate the session
@@ -102,8 +106,10 @@ def port_forward(  # noqa: PLR0913
     else:
         stdout = subprocess.DEVNULL
 
-    print(
-        f"🚀 Starting port forwarding session through [bold]{through}[/bold] with reason: [italic]{reason!r}[/italic].",
+    logger.info(
+        "Starting port forwarding session through [bold]%s[/bold] with reason: [italic]%r[/italic].",
+        through,
+        reason,
     )
     proc = subprocess.Popen(  # noqa: S603
         command,
@@ -112,8 +118,12 @@ def port_forward(  # noqa: PLR0913
         text=True,
         close_fds=False,  # FD inherited from parent process
     )
-    print(f"✅ Session Manager Plugin started with PID {proc.pid}. Outputs will be logged to {log_file.absolute()}.")
+    logger.info(
+        "Session Manager Plugin started with PID %d. Outputs will be logged to %s.",
+        proc.pid,
+        log_file.absolute(),
+    )
 
     # Write the PID to the file
     pid_file.write_text(str(proc.pid))
-    print(f"💾 PID file written to {pid_file.absolute()}.")
+    logger.info("PID file written to %s.", pid_file.absolute())

aws_annoying/cli/session_manager/start.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
+import logging
 import os
 
 import typer
-from rich import print  # noqa: A004
+
+from aws_annoying.utils.ec2 import get_instance_id_by_name
 
 from ._app import session_manager_app
-from ._common import SessionManager, get_instance_id_by_name
+from ._common import SessionManager
+
+logger = logging.getLogger(__name__)
 
 # TODO(lasuillard): ECS support (#24)
 # TODO(lasuillard): Interactive instance selection
@@ -30,14 +34,18 @@ def start(
     # Resolve the instance name or ID
     instance_id = get_instance_id_by_name(target)
     if instance_id:
-        print(f"❗ Instance ID resolved: [bold]{instance_id}[/bold]")
+        logger.info("Instance ID resolved: [bold]%s[/bold]", instance_id)
         target = instance_id
     else:
-        print(f"🚫 Instance with name '{target}' not found.")
+        logger.info("Instance with name '%s' not found.", target)
         raise typer.Exit(1)
 
     # Start the session, replacing the current process
-    print(f"🚀 Starting session to target [bold]{target}[/bold] with reason: [italic]{reason!r}[/italic].")
+    logger.info(
+        "Starting session to target [bold]%s[/bold] with reason: [italic]%r[/italic].",
+        target,
+        reason,
+    )
     command = session_manager.build_command(
         target=target,
         document_name="SSM-SessionManagerRunShell",

aws_annoying/cli/session_manager/stop.py

@@ -1,14 +1,16 @@
 from __future__ import annotations
 
+import logging
 import os
 import signal
 from pathlib import Path  # noqa: TC003
 
 import typer
-from rich import print  # noqa: A004
 
 from ._app import session_manager_app
 
+logger = logging.getLogger(__name__)
+
 
 @session_manager_app.command()
 def stop(
@@ -24,7 +26,7 @@ def stop(
     """Stop running session for PID file."""
     # Check if PID file exists
     if not pid_file.is_file():
-        print(f"❌ PID file not found: {pid_file}")
+        logger.error("PID file not found: %s", pid_file)
         raise typer.Exit(1)
 
     # Read PID from file
@@ -32,19 +34,19 @@ def stop(
     try:
         pid = int(pid_content)
     except ValueError:
-        print(f"🚫 PID file content is invalid; expected integer, but got: {type(pid_content)}")
+        logger.error("PID file content is invalid; expected integer, but got: %s", type(pid_content))  # noqa: TRY400
         raise typer.Exit(1) from None
 
     # Send SIGTERM to the process
     try:
-        print(f"⚠️ Terminating running process with PID {pid}.")
+        logger.warning("Terminating running process with PID %d.", pid)
         os.kill(pid, signal.SIGTERM)
     except ProcessLookupError:
-        print(f"❗ Tried to terminate process with PID {pid} but does not exist.")
+        logger.warning("Tried to terminate process with PID %d but does not exist.", pid)
 
     # Remove the PID file
     if remove:
-        print(f"✅ Removed the PID file {pid_file}.")
+        logger.info("Removed the PID file %s.", pid_file)
         pid_file.unlink()
 
-    print("✅ Terminated the session successfully.")
+    logger.info("Terminated the session successfully.")

aws_annoying/ecs/__init__.py

@@ -0,0 +1,17 @@
+from .common import ECSServiceRef
+from .deployment_waiter import ECSDeploymentWaiter
+from .errors import (
+    DeploymentFailedError,
+    NoRunningDeploymentError,
+    ServiceTaskDefinitionAssertionError,
+    WaitForDeploymentError,
+)
+
+__all__ = (
+    "DeploymentFailedError",
+    "ECSDeploymentWaiter",
+    "ECSServiceRef",
+    "NoRunningDeploymentError",
+    "ServiceTaskDefinitionAssertionError",
+    "WaitForDeploymentError",
+)

aws_annoying/ecs/common.py

@@ -0,0 +1,8 @@
+from typing import NamedTuple
+
+
+class ECSServiceRef(NamedTuple):
+    """Reference to an ECS service."""
+
+    cluster: str
+    service: str

aws_annoying/ecs/deployment_waiter.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import logging
+from operator import itemgetter
+from time import sleep
+from typing import TYPE_CHECKING, Optional
+
+import boto3
+import botocore.exceptions
+from pydantic import PositiveInt, validate_call
+
+from .errors import DeploymentFailedError, NoRunningDeploymentError, ServiceTaskDefinitionAssertionError
+
+if TYPE_CHECKING:
+    from .common import ECSServiceRef
+
+logger = logging.getLogger(__name__)
+
+
+class ECSDeploymentWaiter:
+    """ECS service deployment waiter."""
+
+    def __init__(self, service_ref: ECSServiceRef, *, session: boto3.session.Session | None = None) -> None:
+        """Initialize instance.
+
+        Args:
+            service_ref: Reference to the ECS service.
+            session: Boto3 session to use for AWS operations.
+
+        """
+        self.service_ref = service_ref
+        self.session = session or boto3.session.Session()
+
+    @validate_call
+    def wait(
+        self,
+        *,
+        wait_for_start: bool,
+        polling_interval: PositiveInt = 5,
+        wait_for_stability: bool,
+        expected_task_definition: Optional[str] = None,
+    ) -> None:
+        """Wait for the ECS deployment to complete.
+
+        Args:
+            wait_for_start: Whether to wait for the deployment to start.
+            polling_interval: The interval between any polling attempts, in seconds.
+            wait_for_stability: Whether to wait for the service to be stable after the deployment.
+            expected_task_definition: The service's task definition expected after deployment.
+        """
+        # Find current deployment for the service
+        logger.info(
+            "Looking up running deployment for service %s",
+            self.service_ref.service,
+        )
+        latest_deployment_arn = self.get_latest_deployment_arn(
+            wait_for_start=wait_for_start,
+            polling_interval=polling_interval,
+        )
+
+        # Polling for the deployment to finish (successfully or unsuccessfully)
+        logger.info(
+            "Start waiting for deployment %s to finish.",
+            latest_deployment_arn,
+        )
+        ok, status = self.wait_for_deployment_complete(latest_deployment_arn, polling_interval=polling_interval)
+        if ok:
+            logger.info(
+                "Deployment succeeded with status %s",
+                status,
+            )
+        else:
+            msg = f"Deployment failed with status: {status}"
+            raise DeploymentFailedError(msg)
+
+        # Wait for the service to be stable
+        if wait_for_stability:
+            logger.debug(
+                "Start waiting for service %s to be stable.",
+                self.service_ref.service,
+            )
+            self.wait_for_service_stability(polling_interval=polling_interval)
+
+        # Check if the service task definition matches the expected one
+        if expected_task_definition:
+            logger.info(
+                "Checking if the service task definition is the expected one: %s",
+                expected_task_definition,
+            )
+            ok, actual = self.check_service_task_definition_is(expect=expected_task_definition)
+            if not ok:
+                msg = f"The service task definition is not the expected one; got: {actual!r}"
+                raise ServiceTaskDefinitionAssertionError(msg)
+
+            logger.info("The service task definition matches the expected one.")
+
+    @validate_call
+    def get_latest_deployment_arn(
+        self,
+        *,
+        wait_for_start: bool,
+        polling_interval: PositiveInt,
+        max_attempts: Optional[PositiveInt] = None,
+    ) -> str:
+        """Get the most recently started deployment ARN for the service.
+
+        Args:
+            wait_for_start: Whether to wait for the deployment to start.
+            polling_interval: The interval between any polling attempts, in seconds.
+            max_attempts: The maximum number of attempts to wait for the deployment to start.
+
+        Raises:
+            NoRunningDeploymentError: If no running deployments are found and `wait_for_start` is False.
+
+        Returns:
+            The ARN of the latest deployment for the service.
+        """
+        ecs = self.session.client("ecs")
+        if wait_for_start:
+            logger.warning("`wait_for_start` is set, will wait for a new deployment to start.")
+
+        attempts = 0
+        while True:  # do-while
+            # Do
+            running_deployments = ecs.list_service_deployments(
+                cluster=self.service_ref.cluster,
+                service=self.service_ref.service,
+                status=["PENDING", "IN_PROGRESS"],
+            )["serviceDeployments"]
+
+            # While
+            if running_deployments:
+                logger.debug("Found %d running deployments for service. Exiting loop.", len(running_deployments))
+                break
+
+            if not wait_for_start:
+                logger.debug("`wait_for_start` is off, no need to wait for a new deployment to start.")
+                break
+
+            if max_attempts and attempts >= max_attempts:
+                logger.debug("Max attempts exceeded while waiting for a new deployment to start.")
+                break
+
+            logger.debug(
+                "(%d-th attempt) No running deployments found for service. Start waiting for a new deployment.",
+                attempts + 1,
+            )
+
+            sleep(polling_interval)
+            attempts += 1
+
+        if not running_deployments:
+            msg = "No running deployments found for service."
+            raise NoRunningDeploymentError(msg)
+
+        latest_deployment = sorted(running_deployments, key=itemgetter("startedAt"))[-1]
+        if len(running_deployments) > 1:
+            logger.warning(
+                "%d running deployments found for service. Using most recently started deployment: %s",
+                len(running_deployments),
+                latest_deployment["serviceDeploymentArn"],
+            )
+
+        return latest_deployment["serviceDeploymentArn"]
+
+    @validate_call
+    def wait_for_deployment_complete(
+        self,
+        deployment_arn: str,
+        *,
+        polling_interval: PositiveInt,
+        max_attempts: Optional[PositiveInt] = None,
+    ) -> tuple[bool, str]:
+        """Wait for the ECS deployment to complete.
+
+        Args:
+            deployment_arn: The ARN of the deployment to wait for.
+            polling_interval: The interval between any polling attempts, in seconds.
+            max_attempts: The maximum number of attempts to wait for the deployment to complete.
+
+        Returns:
+            A tuple containing a boolean indicating whether the deployment succeeded and the status of the deployment.
+        """
+        ecs = self.session.client("ecs")
+
+        attempts = 0
+        while (max_attempts is None) or (attempts <= max_attempts):
+            latest_deployment = ecs.describe_service_deployments(serviceDeploymentArns=[deployment_arn])[
+                "serviceDeployments"
+            ][0]
+            status = latest_deployment["status"]
+            if status == "SUCCESSFUL":
+                return (True, status)
+
+            if status in ("PENDING", "IN_PROGRESS"):
+                logger.debug(
+                    "(%d-th attempt) Deployment in progress... (%s)",
+                    attempts + 1,
+                    status,
+                )
+            else:
+                break
+
+            sleep(polling_interval)
+            attempts += 1
+
+        return (False, status)
+
+    @validate_call
+    def wait_for_service_stability(
+        self,
+        *,
+        polling_interval: PositiveInt,
+        max_attempts: Optional[PositiveInt] = None,
+    ) -> bool:
+        """Wait for the ECS service to be stable.
+
+        Args:
+            polling_interval: The interval between any polling attempts, in seconds.
+            max_attempts: The maximum number of attempts to wait for the service to be stable.
+
+        Returns:
+            A boolean indicating whether the service is stable.
+        """
+        ecs = self.session.client("ecs")
+
+        # TODO(lasuillard): Likely to be a problem in some cases: https://github.com/boto/botocore/issues/3314
+        stability_waiter = ecs.get_waiter("services_stable")
+
+        attempts = 0
+        while (max_attempts is None) or (attempts <= max_attempts):
+            logger.debug(
+                "(%d-th attempt) Waiting for service %s to be stable...",
+                attempts + 1,
+                self.service_ref.service,
+            )
+            try:
+                stability_waiter.wait(
+                    cluster=self.service_ref.cluster,
+                    services=[self.service_ref.service],
+                    WaiterConfig={"Delay": polling_interval, "MaxAttempts": 1},
+                )
+            except botocore.exceptions.WaiterError as err:
+                if err.kwargs["reason"] != "Max attempts exceeded":
+                    raise
+            else:
+                return True
+
+            sleep(polling_interval)
+            attempts += 1
+
+        return False
+
+    @validate_call
+    def check_service_task_definition_is(self, expect: str) -> tuple[bool, str]:
+        """Check the service's current task definition matches the expected one.
+
+        Args:
+            expect: The ARN of expected task definition.
+
+        Returns:
+            A tuple containing a boolean indicating whether the task definition matches the expected one
+            and the current task definition ARN.
+        """
+        ecs = self.session.client("ecs")
+
+        service_detail = ecs.describe_services(cluster=self.service_ref.cluster, services=[self.service_ref.service])[
+            "services"
+        ][0]
+        current_task_definition_arn = service_detail["taskDefinition"]
+        if current_task_definition_arn != expect:
+            return (False, current_task_definition_arn)
+
+        return (True, current_task_definition_arn)

aws_annoying/ecs/errors.py

@@ -0,0 +1,14 @@
+class WaitForDeploymentError(Exception):
+    """Base class for all deployment waiter errors."""
+
+
+class NoRunningDeploymentError(WaitForDeploymentError):
+    """No running deployment found for the service."""
+
+
+class DeploymentFailedError(WaitForDeploymentError):
+    """Deployment failed."""
+
+
+class ServiceTaskDefinitionAssertionError(WaitForDeploymentError):
+    """Service task definition does not match the expected one."""

aws_annoying/{mfa.py → mfa_config.py}

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import configparser
+import logging
 from pathlib import Path  # noqa: TC003
 from typing import Optional
 
 from pydantic import BaseModel, ConfigDict
 
-# TODO(lasuillard): Need some refactoring (configurator class)
-# TODO(lasuillard): Put some logging
+logger = logging.getLogger(__name__)
 
 
 class MfaConfig(BaseModel):
@@ -30,9 +30,12 @@ class MfaConfig(BaseModel):
         with path.open("w") as f:
             config_ini.write(f)
 
+        logger.debug("Saved config to %s with section %s", path, section_key)
+
     @classmethod
     def from_ini_file(cls, path: Path, section_key: str) -> tuple[MfaConfig, bool]:
         """Load configuration from an AWS config file, with boolean indicating if the config already exists."""
+        logger.debug("Loading config from %s with section %s", path, section_key)
         config_ini = configparser.ConfigParser()
         config_ini.read(path)
         if config_ini.has_section(section_key):
@@ -52,3 +55,5 @@ def update_credentials(path: Path, profile: str, *, access_key: str, secret_key:
     credentials_ini[profile]["aws_session_token"] = session_token
     with path.open("w") as f:
         credentials_ini.write(f)
+
+    logger.debug("Updated credentials file %s with profile %s", path, profile)

aws_annoying/session_manager/session_manager.py

@@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, NamedTuple
 
 import boto3
 
-from aws_annoying.utils.platform import command_as_root, is_root, os_release
+from aws_annoying.utils.platform import command_as_root, is_root, is_windows, os_release
 
 from .errors import PluginNotInstalledError, UnsupportedPlatformError
 
@@ -22,8 +22,6 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-# TODO(lasuillard): Platform checking is spread everywhere, should be moved to a single place
-
 
 class SessionManager:
     """AWS Session Manager plugin manager."""
@@ -65,7 +63,7 @@ class SessionManager:
         """Get the path to the session-manager-plugin binary."""
         binary_path_str = shutil.which("session-manager-plugin")
         if not binary_path_str:
-            if platform.system() == "Windows":
+            if is_windows():
                 # Windows: use the default installation path
                 binary_path = (
                     Path(os.environ["ProgramFiles"])  # noqa: SIM112

aws_annoying/session_manager/shortcuts.py

@@ -5,6 +5,8 @@ import subprocess
 from contextlib import contextmanager
 from typing import TYPE_CHECKING
 
+from aws_annoying.utils.timeout import Timeout
+
 from .session_manager import SessionManager
 
 if TYPE_CHECKING:
@@ -15,13 +17,14 @@ logger = logging.getLogger(__name__)
 
 
 @contextmanager
-def port_forward(
+def port_forward(  # noqa: PLR0913
     *,
     through: str,
     local_port: int,
     remote_host: str,
     remote_port: int,
     reason: str | None = None,
+    start_timeout: int | None = None,
 ) -> Iterator[subprocess.Popen[str]]:
     """Context manager for port forwarding sessions.
 
@@ -31,6 +34,7 @@ def port_forward(
         remote_host: The remote host to connect to.
         remote_port: The remote port to connect to.
         reason: The reason for starting the session.
+        start_timeout: The timeout in seconds to wait for the session to start.
 
     Returns:
         The command to start the session.
@@ -61,11 +65,11 @@ def port_forward(
 
         # Wait for the session to start
         # ? Not sure this is trustworthy health check
-        # TODO(lasuillard): Need timeout to avoid hanging forever
-        for line in proc.stdout:
-            if "Waiting for connections..." in line:
-                logger.info("Session started successfully.")
-                break
+        with Timeout(start_timeout):
+            for line in proc.stdout:
+                if "Waiting for connections..." in line:
+                    logger.info("Session started successfully.")
+                    break
 
         yield proc
     finally:

aws_annoying/utils/ec2.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import re
+
+import boto3
+
+
+def get_instance_id_by_name(name_or_id: str, *, session: boto3.session.Session | None = None) -> str | None:
+    """Get the EC2 instance ID by name or ID.
+
+    Be aware that this function will only return the first instance found
+    with the given name, no matter how many instances are found.
+
+    Args:
+        name_or_id: The name or ID of the EC2 instance.
+        session: The boto3 session to use. If not provided, a new session will be created.
+
+    Returns:
+        The instance ID if found, otherwise `None`.
+    """
+    if re.match(r"^m?i-[0-9a-f]+$", name_or_id):
+        return name_or_id
+
+    session = session or boto3.session.Session()
+    ec2 = session.client("ec2")
+
+    response = ec2.describe_instances(Filters=[{"Name": "tag:Name", "Values": [name_or_id]}])
+    reservations = response["Reservations"]
+    if not reservations:
+        return None
+
+    instances = reservations[0]["Instances"]
+    if not instances:
+        return None
+
+    return str(instances[0]["InstanceId"])

aws_annoying/utils/platform.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import platform
 from pathlib import Path
 
 
@@ -25,3 +26,13 @@ def os_release() -> dict[str, str]:
         key.strip('"'): value.strip('"')
         for key, value in (line.split("=", 1) for line in content.splitlines() if "=" in line)
     }
+
+
+def is_macos() -> bool:
+    """Check if the current OS is macOS."""
+    return platform.system() == "Darwin"
+
+
+def is_windows() -> bool:
+    """Check if the current OS is Windows."""
+    return platform.system() == "Windows"