aws-annoying 0.6.0__tar.gz → 0.8.0__tar.gz

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/.devcontainer/devcontainer.json

@@ -7,10 +7,8 @@
   "overrideCommand": true,
   "workspaceFolder": "/workspaces/aws-annoying",
   "features": {
-    "ghcr.io/devcontainers-contrib/features/pre-commit:2": {},
     "ghcr.io/devcontainers/features/docker-in-docker:2": {},
-    "ghcr.io/devcontainers/features/aws-cli:1": {},
-    "ghcr.io/lasuillard/devcontainer-features/aws-ssm-session-manager-plugin:0.1": {}
+    "ghcr.io/devcontainers/features/aws-cli:1": {}
   },
   "onCreateCommand": "./.devcontainer/onCreateCommand.sh",
   "postAttachCommand": "./.devcontainer/postAttachCommand.sh",

aws_annoying-0.8.0/.devcontainer/onCreateCommand.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+
+pipx install pre-commit

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/.github/workflows/ci.yaml

@@ -66,23 +66,27 @@ jobs:
       - name: Run tests (Ubuntu)
         if: ${{ matrix.os == 'ubuntu-latest' }}
         run: |
+          which session-manager-plugin || true
           sudo apt-get remove --yes session-manager-plugin
           uv run pytest
 
       - name: Run tests (macOS)
         if: ${{ matrix.os == 'macos-latest' }}
         run: |
+          which session-manager-plugin || true
           brew uninstall session-manager-plugin
           uv run pytest -m 'not docker'
 
       - name: Run tests (Windows)
         if: ${{ matrix.os == 'windows-latest' }}
         run: |
-          Invoke-WebRequest `
-            -Uri "https://raw.githubusercontent.com/aws/session-manager-plugin/refs/heads/mainline/Tools/src/update/windows/uninstall.bat" `
-            -OutFile "uninstall.bat"
+          Get-Command session-manager-plugin | Out-String -Width 2000
+
+          $uri = 'https://s3.amazonaws.com/session-manager-downloads/plugin/latest/windows/SessionManagerPluginSetup.exe'
+          $outFile = 'C:\Windows\Temp\SessionManagerPluginSetup.exe'
+          Invoke-WebRequest -Uri "$uri" -OutFile "$outFile"
+          & "$outFile" /uninstall /quiet
 
-          Start-Process "cmd.exe" -ArgumentList "/c uninstall.bat" -Wait
           uv run pytest -m 'not docker'
 
       - name: Upload test results to Codecov

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/.vscode/settings.json

@@ -10,7 +10,9 @@
   "cSpell.words": [
     "Deregistering",
     "localstack",
+    "pytestmark",
     "sessionmanagerplugin",
+    "Stubber",
     "Tampermonkey"
   ],
   "editor.formatOnSave": true,

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aws-annoying
-Version: 0.6.0
+Version: 0.8.0
 Summary: Utils to handle some annoying AWS tasks.
 Project-URL: Homepage, https://github.com/lasuillard/aws-annoying
 Project-URL: Repository, https://github.com/lasuillard/aws-annoying.git
@@ -16,7 +16,7 @@ Requires-Dist: tqdm<5,>=4
 Requires-Dist: typer<1,>=0
 Provides-Extra: dev
 Requires-Dist: boto3-stubs[ec2,ecs,secretsmanager,ssm,sts]>=1.37.1; extra == 'dev'
-Requires-Dist: mypy~=1.15.0; extra == 'dev'
+Requires-Dist: mypy<1.17,>=1.15; extra == 'dev'
 Requires-Dist: ruff<0.12.0,>=0.9.9; extra == 'dev'
 Requires-Dist: types-requests>=2.31.0.6; extra == 'dev'
 Provides-Extra: test

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/app.py

@@ -8,6 +8,10 @@ from typing import Optional
 import typer
 from rich import print  # noqa: A004
 from rich.console import Console
+from rich.highlighter import ReprHighlighter
+from rich.theme import Theme
+
+logger = logging.getLogger(__name__)
 
 app = typer.Typer(
     pretty_exceptions_short=True,
@@ -51,7 +55,7 @@ def main(  # noqa: D103
     ),
 ) -> None:
     log_level = logging.DEBUG if verbose else logging.INFO
-    console = Console(soft_wrap=True, emoji=False)
+    console = _get_console()
     logging_config: logging.config._DictConfigArgs = {
         "version": 1,
         "disable_existing_loggers": False,
@@ -89,3 +93,28 @@ def main(  # noqa: D103
 
     # Global flags
     ctx.meta["dry_run"] = dry_run
+    if dry_run:
+        logger.warning("Dry run mode enabled. Some operation may behave differently to avoid making changes.")
+
+
+def _get_console() -> Console:
+    theme = Theme(
+        {
+            "repr.arn": "bold orange3",
+            "repr.constant": "bold blue",
+        },
+    )
+    return Console(soft_wrap=True, emoji=False, highlighter=CustomHighlighter(), theme=theme)
+
+
+class CustomHighlighter(ReprHighlighter):
+    """Custom highlighter to handle additional patterns."""
+
+    highlights = [  # noqa: RUF012
+        *ReprHighlighter.highlights,
+        # AWS Resource Name; https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html
+        # NOTE: Quite simplified regex, may not cover all cases.
+        r"(?P<arn>\barn:[0-9a-zA-Z/+=,\.@_\-:]+\b)",
+        # Constants
+        r"(?P<constant>\b[A-Z_]+\b)",
+    ]

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/ecs/task_definition_lifecycle.py

@@ -40,9 +40,6 @@ def task_definition_lifecycle(
 ) -> None:
     """Execute ECS task definition lifecycle."""
     dry_run = ctx.meta["dry_run"]
-    if dry_run:
-        logger.info("Dry run mode enabled. Will not perform any actual changes.")
-
     ecs = boto3.client("ecs")
 
     # Get all task definitions for the family

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/ecs/wait_for_deployment.py

@@ -6,7 +6,15 @@ from typing import Optional
 
 import typer
 
-from aws_annoying.ecs import DeploymentFailedError, ECSDeploymentWaiter, ECSServiceRef
+from aws_annoying.ecs import (
+    DeploymentFailedError,
+    ECSServiceRef,
+    ServiceTaskDefinitionAssertionError,
+    check_service_task_definition,
+    wait_for_deployment_complete,
+    wait_for_deployment_start,
+    wait_for_service_stability,
+)
 from aws_annoying.utils.timeout import OperationTimeoutError, Timeout
 
 from ._app import ecs_app
@@ -63,10 +71,10 @@ def wait_for_deployment(  # noqa: PLR0913
 ) -> None:
     """Wait for ECS deployment to complete."""
     start = datetime.now(tz=timezone.utc)
-    waiter = ECSDeploymentWaiter(ECSServiceRef(cluster=cluster, service=service))
     try:
         with Timeout(timeout_seconds):
-            waiter.wait(
+            _wait_for_deployment(
+                ECSServiceRef(cluster=cluster, service=service),
                 wait_for_start=wait_for_start,
                 polling_interval=polling_interval,
                 wait_for_stability=wait_for_stability,
@@ -92,3 +100,59 @@ def wait_for_deployment(  # noqa: PLR0913
             "Deployment completed in [bold]%.2f[/bold] seconds.",
             elapsed.total_seconds(),
         )
+
+
+def _wait_for_deployment(
+    service_ref: ECSServiceRef,
+    *,
+    wait_for_start: bool,
+    polling_interval: int = 5,
+    wait_for_stability: bool,
+    expected_task_definition: str | None = None,
+) -> None:
+    # Find current deployment for the service
+    logger.info(
+        "Looking up running deployment for service %s",
+        service_ref.service,
+    )
+    latest_deployment_arn = wait_for_deployment_start(
+        service_ref,
+        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 = 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.info(
+            "Start waiting for service %s to be stable.",
+            service_ref.service,
+        )
+        wait_for_service_stability(service_ref, 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 = check_service_task_definition(service_ref, 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.")

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/load_variables.py

@@ -78,8 +78,6 @@ def load_variables(
     The variables are loaded in the order of option provided, overwriting the variables with the same name in the order of the ARNs.
     Existing environment variables are preserved by default, unless `--overwrite-env` is provided.
     """  # noqa: E501
-    dry_run = ctx.meta["dry_run"]
-
     command = ctx.args
     if not command:
         logger.warning("No command provided. Exiting...")
@@ -109,9 +107,6 @@ def load_variables(
     # Retrieve the variables
     loader = VariableLoader()
     logger.info("Retrieving variables from AWS resources...")
-    if dry_run:
-        logger.warning("Dry run mode enabled. Variables won't be loaded from AWS.")
-
     try:
         variables, load_stats = loader.load(map_arns_by_index)
     except Exception as exc:  # noqa: BLE001

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/logging_handler.py

@@ -15,10 +15,10 @@ class RichLogHandler(logging.Handler):
 
     _level_emojis: Final[dict[str, str]] = {
         "DEBUG": "🔍",
-        "INFO": "ℹ️",  # noqa: RUF001
+        "INFO": "🔔",
         "WARNING": "⚠️",
-        "ERROR": "❗",
-        "CRITICAL": "🚨",
+        "ERROR": "🚨",
+        "CRITICAL": "🔥",
     }
 
     def __init__(self, console: Console, *args: Any, **kwargs: Any) -> None:

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/mfa/configure.py

@@ -17,6 +17,7 @@ logger = logging.getLogger(__name__)
 
 @mfa_app.command()
 def configure(  # noqa: PLR0913
+    ctx: typer.Context,
     *,
     mfa_profile: Optional[str] = typer.Option(
         None,
@@ -54,6 +55,8 @@ def configure(  # noqa: PLR0913
     ),
 ) -> None:
     """Configure AWS profile for MFA."""
+    dry_run = ctx.meta["dry_run"]
+
     # Expand user home directory
     aws_credentials = aws_credentials.expanduser()
     aws_config = aws_config.expanduser()
@@ -102,13 +105,14 @@ def configure(  # noqa: PLR0913
         mfa_profile,
         aws_credentials,
     )
-    update_credentials(
-        aws_credentials,
-        mfa_profile,  # type: ignore[arg-type]
-        access_key=credentials["AccessKeyId"],
-        secret_key=credentials["SecretAccessKey"],
-        session_token=credentials["SessionToken"],
-    )
+    if not dry_run:
+        update_credentials(
+            aws_credentials,
+            mfa_profile,  # type: ignore[arg-type]
+            access_key=credentials["AccessKeyId"],
+            secret_key=credentials["SecretAccessKey"],
+            session_token=credentials["SessionToken"],
+        )
 
     # Persist MFA configuration
     if persist:
@@ -120,6 +124,7 @@ def configure(  # noqa: PLR0913
         mfa_config.mfa_profile = mfa_profile
         mfa_config.mfa_source_profile = mfa_source_profile
         mfa_config.mfa_serial_number = mfa_serial_number
-        mfa_config.save_ini_file(aws_config, aws_config_section)
+        if not dry_run:
+            mfa_config.save_ini_file(aws_config, aws_config_section)
     else:
         logger.warning("MFA configuration not persisted.")

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/session_manager/install.py

@@ -15,12 +15,15 @@ logger = logging.getLogger(__name__)
 # https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html
 @session_manager_app.command()
 def install(
-    yes: bool = typer.Option(  # noqa: FBT001
+    ctx: typer.Context,
+    *,
+    yes: bool = typer.Option(
         False,  # noqa: FBT003
         help="Do not ask confirmation for installation.",
     ),
 ) -> None:
     """Install AWS Session Manager plugin."""
+    dry_run = ctx.meta["dry_run"]
     session_manager = SessionManager()
 
     # Check session-manager-plugin already installed
@@ -31,7 +34,8 @@ def install(
 
     # Install session-manager-plugin
     logger.warning("Installing AWS Session Manager plugin. You could be prompted for admin privileges request.")
-    session_manager.install(confirm=yes, downloader=TQDMDownloader())
+    if not dry_run:
+        session_manager.install(confirm=yes, downloader=TQDMDownloader())
 
     # Verify installation
     is_installed, binary_path, version = session_manager.verify_installation()

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/session_manager/port_forward.py

@@ -19,6 +19,8 @@ logger = logging.getLogger(__name__)
 # https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html
 @session_manager_app.command()
 def port_forward(  # noqa: PLR0913
+    ctx: typer.Context,
+    *,
     # TODO(lasuillard): Add `--local-host` option, redirect the traffic to non-localhost bind (unsupported by AWS)
     local_port: int = typer.Option(
         ...,
@@ -48,7 +50,7 @@ def port_forward(  # noqa: PLR0913
         "./session-manager-plugin.pid",
         help="The path to the PID file to store the process ID of the session manager plugin.",
     ),
-    terminate_running_process: bool = typer.Option(  # noqa: FBT001
+    terminate_running_process: bool = typer.Option(
         False,  # noqa: FBT003
         help="Terminate the process in the PID file if it already exists.",
     ),
@@ -58,6 +60,7 @@ def port_forward(  # noqa: PLR0913
     ),
 ) -> None:
     """Start a port forwarding session using AWS Session Manager."""
+    dry_run = ctx.meta["dry_run"]
     session_manager = SessionManager()
 
     # Check if the PID file already exists
@@ -86,7 +89,7 @@ def port_forward(  # noqa: PLR0913
         logger.info("Instance ID resolved: [bold]%s[/bold]", instance_id)
         target = instance_id
     else:
-        logger.info("Instance with name '%s' not found.", through)
+        logger.error("Instance with name '%s' not found.", through)
         raise typer.Exit(1)
 
     # Initiate the session
@@ -111,19 +114,26 @@ def port_forward(  # noqa: PLR0913
         through,
         reason,
     )
-    proc = subprocess.Popen(  # noqa: S603
-        command,
-        stdout=stdout,
-        stderr=subprocess.STDOUT,
-        text=True,
-        close_fds=False,  # FD inherited from parent process
-    )
+    if not dry_run:
+        proc = subprocess.Popen(  # noqa: S603
+            command,
+            stdout=stdout,
+            stderr=subprocess.STDOUT,
+            text=True,
+            close_fds=False,  # FD inherited from parent process
+        )
+        pid = proc.pid
+    else:
+        pid = -1
+
     logger.info(
         "Session Manager Plugin started with PID %d. Outputs will be logged to %s.",
-        proc.pid,
+        pid,
         log_file.absolute(),
     )
 
     # Write the PID to the file
-    pid_file.write_text(str(proc.pid))
+    if not dry_run:
+        pid_file.write_text(str(pid))
+
     logger.info("PID file written to %s.", pid_file.absolute())

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/session_manager/start.py

@@ -18,6 +18,8 @@ logger = logging.getLogger(__name__)
 
 @session_manager_app.command()
 def start(
+    ctx: typer.Context,
+    *,
     target: str = typer.Option(
         ...,
         show_default=False,
@@ -29,6 +31,7 @@ def start(
     ),
 ) -> None:
     """Start new session."""
+    dry_run = ctx.meta["dry_run"]
     session_manager = SessionManager()
 
     # Resolve the instance name or ID
@@ -52,4 +55,5 @@ def start(
         parameters={},
         reason=reason,
     )
-    os.execvp(command[0], command)  # noqa: S606
+    if not dry_run:
+        os.execvp(command[0], command)  # noqa: S606

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/cli/session_manager/stop.py

@@ -14,16 +14,20 @@ logger = logging.getLogger(__name__)
 
 @session_manager_app.command()
 def stop(
+    ctx: typer.Context,
+    *,
     pid_file: Path = typer.Option(  # noqa: B008
         "./session-manager-plugin.pid",
         help="The path to the PID file to store the process ID of the session manager plugin.",
     ),
-    remove: bool = typer.Option(  # noqa: FBT001
+    remove: bool = typer.Option(
         True,  # noqa: FBT003
         help="Remove the PID file after stopping the session.",
     ),
 ) -> None:
     """Stop running session for PID file."""
+    dry_run = ctx.meta["dry_run"]
+
     # Check if PID file exists
     if not pid_file.is_file():
         logger.error("PID file not found: %s", pid_file)
@@ -40,13 +44,15 @@ def stop(
     # Send SIGTERM to the process
     try:
         logger.warning("Terminating running process with PID %d.", pid)
-        os.kill(pid, signal.SIGTERM)
+        if not dry_run:
+            os.kill(pid, signal.SIGTERM)
     except ProcessLookupError:
         logger.warning("Tried to terminate process with PID %d but does not exist.", pid)
 
     # Remove the PID file
     if remove:
         logger.info("Removed the PID file %s.", pid_file)
-        pid_file.unlink()
+        if not dry_run:
+            pid_file.unlink()
 
     logger.info("Terminated the session successfully.")

{aws_annoying-0.6.0 → aws_annoying-0.8.0}/aws_annoying/ecs/__init__.py

@@ -1,17 +1,25 @@
+from .check import check_service_task_definition
 from .common import ECSServiceRef
-from .deployment_waiter import ECSDeploymentWaiter
 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",
-    "ECSDeploymentWaiter",
     "ECSServiceRef",
     "NoRunningDeploymentError",
     "ServiceTaskDefinitionAssertionError",
     "WaitForDeploymentError",
+    "check_service_task_definition",
+    "wait_for_deployment_complete",
+    "wait_for_deployment_start",
+    "wait_for_service_stability",
 )

aws_annoying-0.8.0/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)