aws-annoying 0.2.1__tar.gz → 0.3.0__tar.gz

{aws_annoying-0.2.1 → aws_annoying-0.3.0}/.github/workflows/ci.yaml

@@ -36,16 +36,24 @@ jobs:
         if: always()
 
   test:
-    name: Test (Python ${{ matrix.python-version }})
-    runs-on: ubuntu-latest
+    name: Test (${{ matrix.os }}, Python ${{ matrix.python-version }})
+    runs-on: ${{ matrix.os }}
     permissions:
       contents: read
       id-token: write
     strategy:
       fail-fast: false
       matrix:
+        os: [ubuntu-latest]
         python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
 
+        # TODO(lasuillard): Simplified matrix for macOS and Windows for testing
+        include:
+          - os: macos-latest
+            python-version: "3.9"
+          - os: windows-latest
+            python-version: "3.9"
+
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -62,8 +70,27 @@ jobs:
       - name: Install deps
         run: uv sync --all-extras
 
-      - name: Run tests
-        run: uv run pytest
+      - name: Run tests (Ubuntu)
+        if: ${{ matrix.os == 'ubuntu-latest' }}
+        run: |
+          sudo apt-get remove --yes session-manager-plugin
+          uv run pytest
+
+      - name: Run tests (macOS)
+        if: ${{ matrix.os == 'macos-latest' }}
+        run: |
+          brew uninstall session-manager-plugin
+          uv run pytest -m macos
+
+      - 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"
+
+          Start-Process "cmd.exe" -ArgumentList "/c uninstall.bat" -Wait
+          uv run pytest -m windows
 
       - name: Upload test results to Codecov
         uses: codecov/codecov-action@v5
@@ -73,6 +100,7 @@ jobs:
           report_type: test_results
           files: junit.xml
           flags: >-
+            ${{ matrix.os }}
             python-${{ matrix.python-version }}
 
       - name: Upload coverage report
@@ -82,4 +110,5 @@ jobs:
           fail_ci_if_error: false
           files: coverage.xml
           flags: >-
+            ${{ matrix.os }}
             python-${{ matrix.python-version }}

{aws_annoying-0.2.1 → aws_annoying-0.3.0}/.vscode/settings.json

@@ -8,7 +8,9 @@
   },
   "autoDocstring.docstringFormat": "google-notypes",
   "cSpell.words": [
+    "Deregistering",
     "localstack",
+    "sessionmanagerplugin",
     "Tampermonkey"
   ],
   "editor.formatOnSave": true,

{aws_annoying-0.2.1 → aws_annoying-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aws-annoying
-Version: 0.2.1
+Version: 0.3.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
@@ -11,11 +11,14 @@ License-File: LICENSE
 Requires-Python: <4.0,>=3.9
 Requires-Dist: boto3>=1.37.1
 Requires-Dist: pydantic>=2.10.6
+Requires-Dist: requests>=2.32.3
+Requires-Dist: tqdm>=4.67.1
 Requires-Dist: typer>=0.15.1
 Provides-Extra: dev
-Requires-Dist: boto3-stubs[ecs,secretsmanager,ssm,sts]>=1.37.1; 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: ruff<0.12.0,>=0.9.9; extra == 'dev'
+Requires-Dist: types-requests>=2.31.0.6; extra == 'dev'
 Provides-Extra: test
 Requires-Dist: coverage<7.9,>=7.6; extra == 'test'
 Requires-Dist: moto[ecs,secretsmanager,server,ssm]~=5.1.1; extra == 'test'

{aws_annoying-0.2.1 → aws_annoying-0.3.0}/aws_annoying/ecs_task_definition_lifecycle.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import boto3
 import typer
 from rich import print  # noqa: A004
 
 from .app import app
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+_DELETE_CHUNK_SIZE = 10
+
 
 @app.command()
 def ecs_task_definition_lifecycle(
@@ -22,6 +29,10 @@ def ecs_task_definition_lifecycle(
         min=1,
         max=100,
     ),
+    delete: bool = typer.Option(
+        False,  # noqa: FBT003
+        help="Delete the task definition after deregistering it.",
+    ),
     dry_run: bool = typer.Option(
         False,  # noqa: FBT003
         help="Do not perform any changes, only show what would be done.",
@@ -48,6 +59,7 @@ def ecs_task_definition_lifecycle(
 
     # Keep the latest N task definitions
     expired_taskdef_arns = task_definition_arns[:-keep_latest]
+    print(f"⚠️ Deregistering {len(expired_taskdef_arns)} task definitions...")
     for arn in expired_taskdef_arns:
         if not dry_run:
             ecs.deregister_task_definition(taskDefinition=arn)
@@ -55,3 +67,18 @@ def ecs_task_definition_lifecycle(
         # ARN like: "arn:aws:ecs:<region>:<account-id>:task-definition/<family>:<revision>"
         _, family_revision = arn.split(":task-definition/")
         print(f"✅ Deregistered task definition [yellow]{family_revision!r}[/yellow]")
+
+    if delete and expired_taskdef_arns:
+        # Delete the expired task definitions in chunks due to API limitation
+        print(f"⚠️ Deleting {len(expired_taskdef_arns)} task definitions in chunks of size {_DELETE_CHUNK_SIZE}...")
+        for idx, chunk in enumerate(_chunker(expired_taskdef_arns, _DELETE_CHUNK_SIZE)):
+            if not dry_run:
+                ecs.delete_task_definitions(taskDefinitions=chunk)
+
+            print(f"✅ Deleted {len(chunk)} task definitions in {idx}-th batch.")
+
+
+def _chunker(sequence: list, size: int) -> Iterator[list]:
+    """Yield successive chunks of a given size from the sequence."""
+    for i in range(0, len(sequence), size):
+        yield sequence[i : i + size]

{aws_annoying-0.2.1 → aws_annoying-0.3.0}/aws_annoying/main.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 import aws_annoying.ecs_task_definition_lifecycle
 import aws_annoying.load_variables
 import aws_annoying.mfa
+import aws_annoying.session_manager
 from aws_annoying.utils.debugger import input_as_args
 
 # App with all commands registered

aws_annoying-0.3.0/aws_annoying/session_manager/__init__.py

@@ -0,0 +1,3 @@
+from . import install, port_forward, start, stop
+
+__all__ = ("install", "port_forward", "start", "stop")

aws_annoying-0.3.0/aws_annoying/session_manager/_app.py

@@ -0,0 +1,9 @@
+import typer
+
+from aws_annoying.app import app
+
+session_manager_app = typer.Typer(
+    no_args_is_help=True,
+    help="AWS Session Manager CLI utilities.",
+)
+app.add_typer(session_manager_app, name="session-manager")

aws_annoying-0.3.0/aws_annoying/session_manager/_common.py

@@ -0,0 +1,24 @@
+# TODO(lasuillard): Using this file until split CLI from library codebase
+from __future__ import annotations
+
+from typing import Any
+
+import typer
+from rich.prompt import Confirm
+
+from .session_manager import SessionManager as _SessionManager
+
+
+# Custom session manager with console interactivity
+class SessionManager(_SessionManager):
+    def before_install(self, command: list[str]) -> None:
+        if self._confirm:
+            return
+
+        confirm = Confirm.ask(f"⚠️ Will run the following command: [bold red]{' '.join(command)}[/bold red]. Proceed?")
+        if not confirm:
+            raise typer.Abort
+
+    def install(self, *args: Any, confirm: bool = False, **kwargs: Any) -> None:
+        self._confirm = confirm
+        return super().install(*args, **kwargs)

aws_annoying-0.3.0/aws_annoying/session_manager/errors.py

@@ -0,0 +1,10 @@
+class SessionManagerError(Exception):
+    """Base exception for all errors related to Session Manager."""
+
+
+class UnsupportedPlatformError(SessionManagerError):
+    """Exception raised when the platform is not supported."""
+
+
+class PluginNotInstalledError(SessionManagerError):
+    """Trying to use the Session Manager plugin before it is installed."""

aws_annoying-0.3.0/aws_annoying/session_manager/install.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+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
+
+
+# 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
+        False,  # noqa: FBT003
+        help="Do not ask confirmation for installation.",
+    ),
+) -> None:
+    """Install AWS Session Manager plugin."""
+    session_manager = SessionManager(downloader=TQDMDownloader())
+
+    # 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})")
+        return
+
+    # Install session-manager-plugin
+    print("⬇️ Installing AWS Session Manager plugin. You could be prompted for admin privileges request.")
+    session_manager.install(confirm=yes)
+
+    # Verify installation
+    is_installed, binary_path, version = session_manager.verify_installation()
+    if not is_installed:
+        print("❌ Installation failed. Session Manager plugin not found.")
+        raise typer.Exit(1)
+
+    print(f"✅ Session Manager plugin successfully installed at {binary_path} (version: {version})")

aws_annoying-0.3.0/aws_annoying/session_manager/port_forward.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import os
+import re
+import signal
+from pathlib import Path  # noqa: TC003
+
+import boto3
+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
+
+
+# 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
+    # TODO(lasuillard): Add `--local-host` option, redirect the traffic to non-localhost bind (unsupported by AWS)
+    local_port: int = typer.Option(
+        ...,
+        show_default=False,
+        help="The local port to use for port forwarding.",
+    ),
+    through: str = typer.Option(
+        ...,
+        show_default=False,
+        help="The name or ID of the EC2 instance to use as a proxy for port forwarding.",
+    ),
+    remote_host: str = typer.Option(
+        ...,
+        show_default=False,
+        help="The remote host to connect to.",
+    ),
+    remote_port: int = typer.Option(
+        ...,
+        show_default=False,
+        help="The remote port to connect to.",
+    ),
+    reason: str = typer.Option(
+        "",
+        help="The reason for starting the port forwarding session.",
+    ),
+    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.",
+    ),
+    terminate_running_process: bool = typer.Option(  # noqa: FBT001
+        False,  # noqa: FBT003
+        help="Terminate the process in the PID file if it already exists.",
+    ),
+    log_file: Path = typer.Option(  # noqa: B008
+        "./session-manager-plugin.log",
+        help="The path to the log file to store the output of the session manager plugin.",
+    ),
+) -> None:
+    """Start a port forwarding session using AWS Session Manager."""
+    session_manager = SessionManager(downloader=TQDMDownloader())
+
+    # Check if the PID file already exists
+    if pid_file.exists():
+        if not terminate_running_process:
+            print("🚫 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)}")
+            raise typer.Exit(1) from None
+
+        try:
+            print(f"⚠️ Terminating running process with PID {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.")
+
+    # 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)
+
+        target = instance_id
+
+    # Initiate the session
+    proc = session_manager.start(
+        target=target,
+        document_name="AWS-StartPortForwardingSessionToRemoteHost",
+        parameters={
+            "host": [remote_host],
+            "portNumber": [str(remote_port)],
+            "localPortNumber": [str(local_port)],
+        },
+        reason=reason,
+    )
+    print(f"✅ Session Manager Plugin started with PID {proc.pid}. Outputs will be logged to {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"])