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

aws_annoying/cli/session_manager/install.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
+import logging
+
 import typer
-from rich import print  # noqa: A004
 
 from aws_annoying.utils.downloader import TQDMDownloader
 
 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()
@@ -23,17 +26,17 @@ def install(
     # Check session-manager-plugin already installed
     is_installed, binary_path, version = session_manager.verify_installation()
     if is_installed:
-        print(f"✅ Session Manager plugin is already installed at {binary_path} (version: {version})")
+        logger.info("Session Manager plugin is already installed at %s (version: %s)", binary_path, version)
         return
 
     # Install session-manager-plugin
-    print("⬇️ Installing AWS Session Manager plugin. You could be prompted for admin privileges request.")
+    logger.warning("Installing AWS Session Manager plugin. You could be prompted for admin privileges request.")
     session_manager.install(confirm=yes, downloader=TQDMDownloader())
 
     # Verify installation
     is_installed, binary_path, version = session_manager.verify_installation()
     if not is_installed:
-        print("❌ Installation failed. Session Manager plugin not found.")
+        logger.error("Installation failed. Session Manager plugin not found.")
         raise typer.Exit(1)
 
-    print(f"✅ Session Manager plugin successfully installed at {binary_path} (version: {version})")
+    logger.info("Session Manager plugin successfully installed at %s (version: %s)", binary_path, version)

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,25 @@
+from .check import check_service_task_definition
+from .common import ECSServiceRef
+from .errors import (
+    DeploymentFailedError,
+    NoRunningDeploymentError,
+    ServiceTaskDefinitionAssertionError,
+    WaitForDeploymentError,
+)
+from .wait_for import (
+    wait_for_deployment_complete,
+    wait_for_deployment_start,
+    wait_for_service_stability,
+)
+
+__all__ = (
+    "DeploymentFailedError",
+    "ECSServiceRef",
+    "NoRunningDeploymentError",
+    "ServiceTaskDefinitionAssertionError",
+    "WaitForDeploymentError",
+    "check_service_task_definition",
+    "wait_for_deployment_complete",
+    "wait_for_deployment_start",
+    "wait_for_service_stability",
+)

aws_annoying/ecs/check.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import boto3
+
+if TYPE_CHECKING:
+    from .common import ECSServiceRef
+
+logger = logging.getLogger(__name__)
+
+
+def check_service_task_definition(
+    service_ref: ECSServiceRef,
+    *,
+    session: boto3.session.Session | None = None,
+    expect: str,
+) -> tuple[bool, str]:
+    """Check the service's current task definition matches the expected one.
+
+    Args:
+        service_ref: The ECS service reference containing the cluster and service names.
+        session: The boto3 session to use for the ECS client.
+        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.
+    """
+    session = session or boto3.session.Session()
+    ecs = session.client("ecs")
+
+    service_detail = ecs.describe_services(cluster=service_ref.cluster, services=[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/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/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/ecs/wait_for.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from time import sleep
+from typing import TYPE_CHECKING
+
+import boto3
+import botocore.exceptions
+
+from .errors import NoRunningDeploymentError
+
+if TYPE_CHECKING:
+    from .common import ECSServiceRef
+
+logger = logging.getLogger(__name__)
+
+
+def wait_for_deployment_start(
+    service_ref: ECSServiceRef,
+    *,
+    session: boto3.session.Session | None = None,
+    wait_for_start: bool,
+    polling_interval: int = 5,
+    max_attempts: int | None = None,
+) -> str:
+    """Wait for the ECS deployment to start.
+
+    Args:
+        service_ref: The ECS service reference containing the cluster and service names.
+        session: The boto3 session to use for the ECS client.
+        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.
+    """
+    session = session or boto3.session.Session()
+    ecs = 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=service_ref.cluster,
+            service=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 = max(
+        running_deployments,
+        key=lambda dep: dep.get(
+            "startedAt",
+            datetime.min.replace(tzinfo=timezone.utc),
+        ),
+    )
+    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"]
+
+
+def wait_for_deployment_complete(
+    deployment_arn: str,
+    *,
+    session: boto3.session.Session | None = None,
+    polling_interval: int = 5,
+    max_attempts: int | None = None,
+) -> tuple[bool, str]:
+    """Wait for the ECS deployment to complete.
+
+    Args:
+        deployment_arn: The ARN of the deployment to wait for.
+        session: The boto3 session to use for the ECS client.
+        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.
+    """
+    session = session or boto3.session.Session()
+    ecs = 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)
+
+
+def wait_for_service_stability(
+    service_ref: ECSServiceRef,
+    *,
+    session: boto3.session.Session | None = None,
+    polling_interval: int = 5,
+    max_attempts: int | None = None,
+) -> bool:
+    """Wait for the ECS service to be stable.
+
+    Args:
+        service_ref: The ECS service reference containing the cluster and service names.
+        session: The boto3 session to use for the ECS client.
+        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.
+    """
+    session = session or boto3.session.Session()
+    ecs = 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,
+            service_ref.service,
+        )
+        try:
+            stability_waiter.wait(
+                cluster=service_ref.cluster,
+                services=[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

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

@@ -42,14 +42,7 @@ class TQDMDownloader(AbstractDownloader):
             total_size = int(response.headers.get("content-length", 0))
             with (
                 to.open("wb") as f,
-                tqdm(
-                    # Make the URL less verbose in the progress bar
-                    desc=url.replace("https://s3.amazonaws.com/session-manager-downloads/plugin", "..."),
-                    total=total_size,
-                    unit="iB",
-                    unit_scale=True,
-                    unit_divisor=1_024,
-                ) as pbar,
+                tqdm(desc=url, total=total_size, unit="iB", unit_scale=True, unit_divisor=1_024) as pbar,
             ):
                 for chunk in response.iter_content(chunk_size=8_192):
                     size = f.write(chunk)

aws_annoying/utils/ec2.py

@@ -0,0 +1,33 @@
+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 or not reservations[0]["Instances"]:
+        return None
+
+    instances = reservations[0]["Instances"]
+    return str(instances[0]["InstanceId"])