aws-annoying 0.4.0__tar.gz → 0.5.0__tar.gz

{aws_annoying-0.4.0 → aws_annoying-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aws-annoying
-Version: 0.4.0
+Version: 0.5.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

aws_annoying-0.5.0/aws_annoying/cli/session_manager/_common.py

@@ -0,0 +1,54 @@
+# TODO(lasuillard): Using this file until split CLI from library codebase
+from __future__ import annotations
+
+import re
+from typing import Any
+
+import boto3
+import typer
+from rich.prompt import Confirm
+
+from aws_annoying.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)
+
+
+def get_instance_id_by_name(name_or_id: str) -> 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.
+
+    Returns:
+        The instance ID if found, otherwise `None`.
+    """
+    if re.match(r"m?i-.+", name_or_id):
+        return name_or_id
+
+    ec2 = boto3.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-0.4.0 → aws_annoying-0.5.0}/aws_annoying/cli/session_manager/install.py

@@ -18,7 +18,7 @@ def install(
     ),
 ) -> None:
     """Install AWS Session Manager plugin."""
-    session_manager = SessionManager(downloader=TQDMDownloader())
+    session_manager = SessionManager()
 
     # Check session-manager-plugin already installed
     is_installed, binary_path, version = session_manager.verify_installation()
@@ -28,7 +28,7 @@ def install(
 
     # Install session-manager-plugin
     print("⬇️ Installing AWS Session Manager plugin. You could be prompted for admin privileges request.")
-    session_manager.install(confirm=yes)
+    session_manager.install(confirm=yes, downloader=TQDMDownloader())
 
     # Verify installation
     is_installed, binary_path, version = session_manager.verify_installation()

{aws_annoying-0.4.0 → aws_annoying-0.5.0}/aws_annoying/cli/session_manager/port_forward.py

@@ -1,18 +1,15 @@
 from __future__ import annotations
 
 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 ._app import session_manager_app
-from ._common import SessionManager
+from ._common import SessionManager, get_instance_id_by_name
 
 
 # https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html
@@ -57,7 +54,7 @@ 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():
@@ -80,21 +77,16 @@ def port_forward(  # noqa: PLR0913
             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)
-
+    instance_id = get_instance_id_by_name(through)
+    if instance_id:
+        print(f"❗ Instance ID resolved: [bold]{instance_id}[/bold]")
         target = instance_id
+    else:
+        print(f"🚫 Instance with name '{through}' not found.")
+        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 +96,24 @@ def port_forward(  # noqa: PLR0913
         },
         reason=reason,
     )
+    stdout: subprocess._FILE
+    if log_file is not None:  # noqa: SIM108
+        stdout = log_file.open(mode="at+", buffering=1)
+    else:
+        stdout = subprocess.DEVNULL
+
+    print(
+        f"🚀 Starting port forwarding session through [bold]{through}[/bold] with reason: [italic]{reason!r}[/italic].",
+    )
+    proc = subprocess.Popen(  # noqa: S603
+        command,
+        stdout=stdout,
+        stderr=subprocess.STDOUT,
+        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()}.")
 
     # 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"])

aws_annoying-0.5.0/aws_annoying/cli/session_manager/start.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import os
+
+import typer
+from rich import print  # noqa: A004
+
+from ._app import session_manager_app
+from ._common import SessionManager, get_instance_id_by_name
+
+# TODO(lasuillard): ECS support (#24)
+# TODO(lasuillard): Interactive instance selection
+
+
+@session_manager_app.command()
+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."""
+    session_manager = SessionManager()
+
+    # 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]")
+        target = instance_id
+    else:
+        print(f"🚫 Instance with name '{target}' not found.")
+        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].")
+    command = session_manager.build_command(
+        target=target,
+        document_name="SSM-SessionManagerRunShell",
+        parameters={},
+        reason=reason,
+    )
+    os.execvp(command[0], command)  # noqa: S606

aws_annoying-0.5.0/aws_annoying/session_manager/__init__.py

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

{aws_annoying-0.4.0 → aws_annoying-0.5.0}/aws_annoying/session_manager/session_manager.py

@@ -28,15 +28,13 @@ logger = logging.getLogger(__name__)
 class SessionManager:
     """AWS Session Manager plugin manager."""
 
-    def __init__(self, *, session: boto3.session.Session | None = None, downloader: AbstractDownloader) -> None:
+    def __init__(self, *, session: boto3.session.Session | None = None) -> None:
         """Initialize SessionManager.
 
         Args:
             session: Boto3 session to use for AWS operations.
-            downloader: File downloader to use for downloading the plugin.
         """
         self.session = session or boto3.session.Session()
-        self.downloader = downloader
 
     # ------------------------------------------------------------------------
     # Installation
@@ -90,6 +88,7 @@ class SessionManager:
         linux_distribution: _LinuxDistribution | None = None,
         arch: str | None = None,
         root: bool | None = None,
+        downloader: AbstractDownloader,
     ) -> None:
         """Install AWS Session Manager plugin.
 
@@ -99,17 +98,23 @@ class SessionManager:
                 If `None` and current `os` is `"Linux"`, will try to detect the distribution from current system.
             arch: The architecture to install the plugin on. If `None`, will use the current architecture.
             root: Whether to run the installation as root. If `None`, will check if the current user is root.
+            downloader: File downloader to use for downloading the plugin.
         """
         os = os or platform.system()
         arch = arch or platform.machine()
 
         if os == "Windows":
-            self._install_windows()
+            self._install_windows(downloader=downloader)
         elif os == "Darwin":
-            self._install_macos(arch=arch, root=root or is_root())
+            self._install_macos(arch=arch, root=root or is_root(), downloader=downloader)
         elif os == "Linux":
             linux_distribution = linux_distribution or _detect_linux_distribution()
-            self._install_linux(linux_distribution=linux_distribution, arch=arch, root=root or is_root())
+            self._install_linux(
+                linux_distribution=linux_distribution,
+                arch=arch,
+                root=root or is_root(),
+                downloader=downloader,
+            )
         else:
             msg = f"Unsupported operating system: {os}"
             raise UnsupportedPlatformError(msg)
@@ -118,20 +123,20 @@ class SessionManager:
         """Hook to run before invoking plugin installation command."""
 
     # https://docs.aws.amazon.com/systems-manager/latest/userguide/install-plugin-windows.html
-    def _install_windows(self) -> None:
+    def _install_windows(self, *, downloader: AbstractDownloader) -> None:
         """Install session-manager-plugin on Windows via EXE installer."""
         download_url = (
             "https://s3.amazonaws.com/session-manager-downloads/plugin/latest/windows/SessionManagerPluginSetup.exe"
         )
         with tempfile.TemporaryDirectory() as temp_dir:
             p = Path(temp_dir)
-            exe_installer = self.downloader.download(download_url, to=p / "SessionManagerPluginSetup.exe")
+            exe_installer = downloader.download(download_url, to=p / "SessionManagerPluginSetup.exe")
             command = [str(exe_installer), "/quiet"]
             self.before_install(command)
             subprocess.call(command, cwd=p)  # noqa: S603
 
     # https://docs.aws.amazon.com/systems-manager/latest/userguide/install-plugin-macos-overview.html
-    def _install_macos(self, *, arch: str, root: bool) -> None:
+    def _install_macos(self, *, arch: str, root: bool, downloader: AbstractDownloader) -> None:
         """Install session-manager-plugin on macOS via signed installer."""
         # ! Intel chip will not be supported
         if arch == "x86_64":
@@ -148,7 +153,7 @@ class SessionManager:
 
         with tempfile.TemporaryDirectory() as temp_dir:
             p = Path(temp_dir)
-            pkg_installer = self.downloader.download(download_url, to=p / "session-manager-plugin.pkg")
+            pkg_installer = downloader.download(download_url, to=p / "session-manager-plugin.pkg")
 
             # Run installer
             command = command_as_root(
@@ -179,6 +184,7 @@ class SessionManager:
         linux_distribution: _LinuxDistribution,
         arch: str,
         root: bool,
+        downloader: AbstractDownloader,
     ) -> None:
         name = linux_distribution.name
         version = linux_distribution.version
@@ -200,7 +206,7 @@ class SessionManager:
 
             with tempfile.TemporaryDirectory() as temp_dir:
                 p = Path(temp_dir)
-                deb_installer = self.downloader.download(download_url, to=p / "session-manager-plugin.deb")
+                deb_installer = downloader.download(download_url, to=p / "session-manager-plugin.deb")
 
                 # Invoke installation command
                 command = command_as_root(["dpkg", "--install", str(deb_installer)], root=root)
@@ -241,28 +247,25 @@ class SessionManager:
             raise UnsupportedPlatformError(msg)
 
     # ------------------------------------------------------------------------
-    # Session
+    # Command
     # ------------------------------------------------------------------------
-    def start(
+    def build_command(
         self,
-        *,
         target: str,
         document_name: str,
         parameters: dict[str, Any],
         reason: str | None = None,
-        log_file: Path | None = None,
-    ) -> subprocess.Popen:
-        """Start new session.
+    ) -> list[str]:
+        """Build command for starting a session.
 
         Args:
-            target: The target instance ID or name.
+            target: The target instance ID.
             document_name: The SSM document name to use for the session.
             parameters: The parameters to pass to the SSM document.
             reason: The reason for starting the session.
-            log_file: Optional file to log output to.
 
         Returns:
-            Process ID of the session.
+            The command to start the session.
         """
         is_installed, binary_path, version = self.verify_installation()
         if not is_installed:
@@ -279,7 +282,7 @@ class SessionManager:
         )
 
         region = self.session.region_name
-        command = [
+        return [
             str(binary_path),
             json.dumps(response),
             region,
@@ -289,20 +292,6 @@ class SessionManager:
             f"https://ssm.{region}.amazonaws.com",
         ]
 
-        stdout: subprocess._FILE
-        if log_file is not None:  # noqa: SIM108
-            stdout = log_file.open(mode="at+", buffering=1)
-        else:
-            stdout = subprocess.DEVNULL
-
-        return subprocess.Popen(  # noqa: S603
-            command,
-            stdout=stdout,
-            stderr=subprocess.STDOUT,
-            text=True,
-            close_fds=False,  # FD inherited from parent process
-        )
-
 
 # ? Could be moved to utils, but didn't because it's too specific to this module
 class _LinuxDistribution(NamedTuple):

aws_annoying-0.5.0/aws_annoying/session_manager/shortcuts.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import logging
+import subprocess
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+from .session_manager import SessionManager
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def port_forward(
+    *,
+    through: str,
+    local_port: int,
+    remote_host: str,
+    remote_port: int,
+    reason: str | None = None,
+) -> Iterator[subprocess.Popen[str]]:
+    """Context manager for port forwarding sessions.
+
+    Args:
+        through: The instance ID to use as port-forwarding proxy.
+        local_port: The local port to listen to.
+        remote_host: The remote host to connect to.
+        remote_port: The remote port to connect to.
+        reason: The reason for starting the session.
+
+    Returns:
+        The command to start the session.
+    """
+    session_manager = SessionManager()
+    command = session_manager.build_command(
+        target=through,
+        document_name="AWS-StartPortForwardingSessionToRemoteHost",
+        parameters={
+            "localPortNumber": [str(local_port)],
+            "host": [remote_host],
+            "portNumber": [str(remote_port)],
+        },
+        reason=reason,
+    )
+    try:
+        proc = subprocess.Popen(  # noqa: S603
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+        )
+
+        # * Must be unreachable
+        if proc.stdout is None:
+            msg = "Standard output is not available"
+            raise RuntimeError(msg)
+
+        # 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
+
+        yield proc
+    finally:
+        proc.terminate()

{aws_annoying-0.4.0 → aws_annoying-0.5.0}/console/ecs-exec/ecs-exec.user.js

@@ -1,7 +1,7 @@
 // ==UserScript==
 // @name         AWS ECS Exec
 // @namespace    mailto:lasuillard@gmail.com
-// @version      2025.03.21
+// @version      2025.04.16
 // @description  Add link to AWS SSM Session Manager for ECS container
 // @author       lasuillard
 // @source       https://raw.githubusercontent.com/lasuillard/aws-annoying/refs/heads/main/console/ecs-exec/ecs-exec.user.js
@@ -106,11 +106,11 @@
 
   // Get task info from the detail page
   function getTaskInfoForDetailPage() {
-    const arnNeighbor = document.evaluate(`//div[text()="ARN"]`, document).iterateNext();
+    const arnNeighbor = document.evaluate(`//*[text()="ARN"]`, document).iterateNext();
     if (!arnNeighbor) {
       return null;
     }
-    const arn = arnNeighbor.parentNode.children[1].textContent;
+    const arn = arnNeighbor.parentNode.parentNode.children[1].textContent;
     const [, , , region, , taskPart] = arn.split(":");
     const [, clusterName, taskId] = taskPart.split("/");
 

{aws_annoying-0.4.0 → aws_annoying-0.5.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "aws-annoying"
 description = "Utils to handle some annoying AWS tasks."
-version = "0.4.0"
+version = "0.5.0"
 authors = [{ name = "Yuchan Lee", email = "lasuillard@gmail.com" }]
 readme = "README.md"
 license = "MIT"

{aws_annoying-0.4.0 → aws_annoying-0.5.0}/tests/session_manager/test_session_manager.py

@@ -5,7 +5,6 @@ from typer.testing import CliRunner
 
 from aws_annoying.cli.main import app
 from aws_annoying.session_manager import SessionManager
-from aws_annoying.utils.downloader import DummyDownloader
 
 runner = CliRunner()
 
@@ -15,7 +14,7 @@ sentinel = object()
 @pytest.mark.macos
 def test_macos_session_manager_install() -> None:
     # Arrange
-    session_manager = SessionManager(downloader=DummyDownloader())
+    session_manager = SessionManager()
     assert session_manager.verify_installation() == (False, None, None)
 
     # Act
@@ -36,7 +35,7 @@ def test_macos_session_manager_install() -> None:
 @pytest.mark.windows
 def test_windows_session_manager_install() -> None:
     # Arrange
-    session_manager = SessionManager(downloader=DummyDownloader())
+    session_manager = SessionManager()
     assert session_manager.verify_installation() == (False, None, None)
 
     # Act

{aws_annoying-0.4.0 → aws_annoying-0.5.0}/uv.lock

@@ -35,7 +35,7 @@ wheels = [
 
 [[package]]
 name = "aws-annoying"
-version = "0.4.0"
+version = "0.5.0"
 source = { editable = "." }
 dependencies = [
     { name = "boto3" },

aws_annoying-0.4.0/aws_annoying/cli/session_manager/_common.py

@@ -1,24 +0,0 @@
-# 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 aws_annoying.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.4.0/aws_annoying/cli/session_manager/start.py

@@ -1,9 +0,0 @@
-from __future__ import annotations
-
-from ._app import session_manager_app
-
-
-@session_manager_app.command()
-def start() -> None:
-    """Start new session."""
-    # TODO(lasuillard): To be implemented (maybe in #24?)

aws_annoying-0.4.0/aws_annoying/session_manager/__init__.py

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