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

aws_annoying/cli/session_manager/_common.py

@@ -1,8 +1,10 @@
 # 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
 
@@ -22,3 +24,31 @@ class SessionManager(_SessionManager):
     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/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/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/cli/session_manager/start.py

@@ -1,9 +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() -> 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:
+        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/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",
+)

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/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.dist-info → aws_annoying-0.5.0.dist-info}/METADATA

@@ -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.4.0.dist-info → aws_annoying-0.5.0.dist-info}/RECORD

@@ -11,20 +11,21 @@ aws_annoying/cli/mfa/_app.py,sha256=Ub7gxb6kGF3Ve1ucQSOjHmc4jAu8mxgegcXsIbOzLLQ,
 aws_annoying/cli/mfa/configure.py,sha256=vsoHfTVFF2dPgiYsp2L-EkMwtAA0_-tVwFd6Wv6DscU,3746
 aws_annoying/cli/session_manager/__init__.py,sha256=FkT6jT6OXduOURN61d-U6hgd-XluQbvuVtKXXiXgSEk,105
 aws_annoying/cli/session_manager/_app.py,sha256=OVOHW0iyKzunvaqLhjoseHw1-WxJ1gGb7QmiyAEezyY,221
-aws_annoying/cli/session_manager/_common.py,sha256=0fCm6Zx6P1NcyyiHlSoB9PgIdrxzUXLVPqKJWQnJ8I4,792
-aws_annoying/cli/session_manager/install.py,sha256=zX2cu-IBYlf9yl8z9P0CTQz72BP6aAMo-KnolC_PGsU,1443
-aws_annoying/cli/session_manager/port_forward.py,sha256=uMcsafTNAHZh-0E1yWXCliMivWyuWV2K9DdVhBZ9pG8,4327
-aws_annoying/cli/session_manager/start.py,sha256=1yaMuy-7IWrrDoPHjgygOP_-g_tajfnpaVPZanxsZxs,215
+aws_annoying/cli/session_manager/_common.py,sha256=u23F4mJOHWHphLYL1gOAh8J3a_Odyk4VVvI175KWmzg,1616
+aws_annoying/cli/session_manager/install.py,sha256=zcQi91xVFKhbSOD4VBc6YG9-fDnhymVyK63PlITdxug,1445
+aws_annoying/cli/session_manager/port_forward.py,sha256=J8_CIrTsbcOYRYOdHkdz81dkaNWxI_l70E8gyJ-Ukh8,4192
+aws_annoying/cli/session_manager/start.py,sha256=pPS0jKuURGTX-WTix3owqqisX-bmCydqfyqC0kFGnt8,1358
 aws_annoying/cli/session_manager/stop.py,sha256=ttU6nlbVgBkZDtY-DwUyCstv5TFtat5TljkyuY8QICU,1482
-aws_annoying/session_manager/__init__.py,sha256=swbRdFh9CdO6tzyBjSxN0KC7MMASvptIpbsijUPIZgI,243
+aws_annoying/session_manager/__init__.py,sha256=IENviL3ux2LF7o9xFGYEiqaGw03hxnyNX2btbB1xyEU,318
 aws_annoying/session_manager/errors.py,sha256=YioKlRtZ-GUP0F_ts_ebw7-HYkxe8mTes6HK821Kuiw,353
-aws_annoying/session_manager/session_manager.py,sha256=7tqfFPA4bWF0Me_VbUDcq5Cpl1sw3bevGyT8QWfLTOk,12454
+aws_annoying/session_manager/session_manager.py,sha256=myZxY_WE4akdlTsH1mOvf0Ublwg-hf1vEkEcmdZyYSU,12147
+aws_annoying/session_manager/shortcuts.py,sha256=uFRPGia_5gqfBDxwOjmLg7UFzhvkSFUqopWuzN5_kbA,1973
 aws_annoying/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aws_annoying/utils/debugger.py,sha256=UFllDCGI2gPtwo1XS5vqw0qyR6bYr7XknmBwSxalKIc,754
 aws_annoying/utils/downloader.py,sha256=aB5RzT-LpbFX24-2HXlAkdgVowc4TR9FWT_K8WwZ1BE,1923
 aws_annoying/utils/platform.py,sha256=h3DUWmTMM-_4TfTWNqY0uNqyVsBjAuMm2DEbG-daxe8,742
-aws_annoying-0.4.0.dist-info/METADATA,sha256=9u5wofuLZXXPrAj1YcOhbqPlPH46jtub8AUvs2KNZ9s,1916
-aws_annoying-0.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-aws_annoying-0.4.0.dist-info/entry_points.txt,sha256=DcKE5V0WvVJ8wUOHxyUz1yLAJOuuJUgRPlMcQ4O7jEs,66
-aws_annoying-0.4.0.dist-info/licenses/LICENSE,sha256=Q5GkvYijQ2KTQ-QWhv43ilzCno4ZrzrEuATEQZd9rYo,1067
-aws_annoying-0.4.0.dist-info/RECORD,,
+aws_annoying-0.5.0.dist-info/METADATA,sha256=kHaGAHqfkZ8Ip4NzcLbE7Bh01oA1CvoTsUa3Ljx1ctU,1916
+aws_annoying-0.5.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+aws_annoying-0.5.0.dist-info/entry_points.txt,sha256=DcKE5V0WvVJ8wUOHxyUz1yLAJOuuJUgRPlMcQ4O7jEs,66
+aws_annoying-0.5.0.dist-info/licenses/LICENSE,sha256=Q5GkvYijQ2KTQ-QWhv43ilzCno4ZrzrEuATEQZd9rYo,1067
+aws_annoying-0.5.0.dist-info/RECORD,,