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

aws_annoying/cli/session_manager/port_forward.py

@@ -1,19 +1,20 @@
 from __future__ import annotations
 
+import logging
 import os
-import re
 import signal
+import subprocess
 from pathlib import Path  # noqa: TC003
 
-import boto3
 import typer
-from rich import print  # noqa: A004
 
-from aws_annoying.utils.downloader import TQDMDownloader
+from aws_annoying.utils.ec2 import get_instance_id_by_name
 
 from ._app import session_manager_app
 from ._common import SessionManager
 
+logger = logging.getLogger(__name__)
+
 
 # https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html
 @session_manager_app.command()
@@ -57,44 +58,39 @@ def port_forward(  # noqa: PLR0913
     ),
 ) -> None:
     """Start a port forwarding session using AWS Session Manager."""
-    session_manager = SessionManager(downloader=TQDMDownloader())
+    session_manager = SessionManager()
 
     # 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
-    if re.match(r"m?i-.+", through):
-        target = through
-    else:
-        # If the instance name is provided, get the instance ID
-        instance_id = _get_instance_id_by_name(through)
-        if instance_id:
-            print(f"❗ Instance ID resolved: [bold]{instance_id}[/bold]")
-        else:
-            print(f"🚫 Instance with name '{through}' not found.")
-            raise typer.Exit(1)
-
+    instance_id = get_instance_id_by_name(through)
+    if instance_id:
+        logger.info("Instance ID resolved: [bold]%s[/bold]", instance_id)
         target = instance_id
+    else:
+        logger.info("Instance with name '%s' not found.", through)
+        raise typer.Exit(1)
 
     # Initiate the session
-    proc = session_manager.start(
+    command = session_manager.build_command(
         target=target,
         document_name="AWS-StartPortForwardingSessionToRemoteHost",
         parameters={
@@ -104,23 +100,30 @@ def port_forward(  # noqa: PLR0913
         },
         reason=reason,
     )
-    print(f"✅ Session Manager Plugin started with PID {proc.pid}. Outputs will be logged to {log_file.absolute()}.")
+    stdout: subprocess._FILE
+    if log_file is not None:  # noqa: SIM108
+        stdout = log_file.open(mode="at+", buffering=1)
+    else:
+        stdout = subprocess.DEVNULL
+
+    logger.info(
+        "Starting port forwarding session through [bold]%s[/bold] with reason: [italic]%r[/italic].",
+        through,
+        reason,
+    )
+    proc = subprocess.Popen(  # noqa: S603
+        command,
+        stdout=stdout,
+        stderr=subprocess.STDOUT,
+        text=True,
+        close_fds=False,  # FD inherited from parent process
+    )
+    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()}.")
-
-
-def _get_instance_id_by_name(name: str) -> str | None:
-    """Get the EC2 instance ID by name."""
-    ec2 = boto3.client("ec2")
-    response = ec2.describe_instances(Filters=[{"Name": "tag:Name", "Values": [name]}])
-    reservations = response["Reservations"]
-    if not reservations:
-        return None
-
-    instances = reservations[0]["Instances"]
-    if not instances:
-        return None
-
-    return str(instances[0]["InstanceId"])
+    logger.info("PID file written to %s.", pid_file.absolute())

aws_annoying/cli/session_manager/start.py

@@ -1,9 +1,55 @@
 from __future__ import annotations
 
+import logging
+import os
+
+import typer
+
+from aws_annoying.utils.ec2 import get_instance_id_by_name
+
 from ._app import session_manager_app
+from ._common import SessionManager
+
+logger = logging.getLogger(__name__)
+
+# TODO(lasuillard): ECS support (#24)
+# TODO(lasuillard): Interactive instance selection
 
 
 @session_manager_app.command()
-def start() -> None:
+def start(
+    target: str = typer.Option(
+        ...,
+        show_default=False,
+        help="The name or ID of the EC2 instance to connect to.",
+    ),
+    reason: str = typer.Option(
+        "",
+        help="The reason for starting the session.",
+    ),
+) -> None:
     """Start new session."""
-    # TODO(lasuillard): To be implemented (maybe in #24?)
+    session_manager = SessionManager()
+
+    # Resolve the instance name or ID
+    instance_id = get_instance_id_by_name(target)
+    if instance_id:
+        logger.info("Instance ID resolved: [bold]%s[/bold]", instance_id)
+        target = instance_id
+    else:
+        logger.info("Instance with name '%s' not found.", target)
+        raise typer.Exit(1)
+
+    # Start the session, replacing the current process
+    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",
+        parameters={},
+        reason=reason,
+    )
+    os.execvp(command[0], command)  # noqa: S606

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/__init__.py

@@ -1,4 +1,11 @@
 from .errors import PluginNotInstalledError, SessionManagerError, UnsupportedPlatformError
 from .session_manager import SessionManager
+from .shortcuts import port_forward
 
-__all__ = ("PluginNotInstalledError", "SessionManager", "SessionManagerError", "UnsupportedPlatformError")
+__all__ = (
+    "PluginNotInstalledError",
+    "SessionManager",
+    "SessionManagerError",
+    "UnsupportedPlatformError",
+    "port_forward",
+)