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

aws_annoying/{ecs_task_definition_lifecycle.py → cli/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/cli/load_variables.py

@@ -0,0 +1,139 @@
+# flake8: noqa: B008
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import NoReturn, Optional
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from aws_annoying.variables import VariableLoader
+
+from .app import app
+
+
+@app.command(
+    context_settings={
+        # Allow extra arguments for user provided command
+        "allow_extra_args": True,
+        "ignore_unknown_options": True,
+    },
+)
+def load_variables(  # noqa: PLR0913
+    *,
+    ctx: typer.Context,
+    arns: list[str] = typer.Option(
+        [],
+        metavar="ARN",
+        help=(
+            "ARNs of the secret or parameter to load."
+            " The variables are loaded in the order of the ARNs,"
+            " overwriting the variables with the same name in the order of the ARNs."
+        ),
+    ),
+    env_prefix: Optional[str] = typer.Option(
+        None,
+        help="Prefix of the environment variables to load the ARNs from.",
+        show_default=False,
+    ),
+    overwrite_env: bool = typer.Option(
+        False,  # noqa: FBT003
+        help="Overwrite the existing environment variables with the same name.",
+    ),
+    quiet: bool = typer.Option(
+        False,  # noqa: FBT003
+        help="Suppress all outputs from this command.",
+    ),
+    dry_run: bool = typer.Option(
+        False,  # noqa: FBT003
+        help="Print the progress only. Neither load variables nor run the command.",
+    ),
+    replace: bool = typer.Option(
+        True,  # noqa: FBT003
+        help=(
+            "Replace the current process (`os.execvpe`) with the command."
+            " If disabled, run the command as a `subprocess`."
+        ),
+    ),
+) -> NoReturn:
+    """Wrapper command to run command with variables from AWS resources injected as environment variables.
+
+    This script is intended to be used in the ECS environment, where currently AWS does not support
+    injecting whole JSON dictionary of secrets or parameters as environment variables directly.
+
+    It first loads the variables from the AWS sources then runs the command with the variables injected as environment variables.
+
+    In addition to `--arns` option, you can provide ARNs as the environment variables by providing `--env-prefix`.
+    For example, if you have the following environment variables:
+
+    ```shell
+    export LOAD_AWS_CONFIG__001_app_config=arn:aws:secretsmanager:...
+    export LOAD_AWS_CONFIG__002_db_config=arn:aws:ssm:...
+    ```
+
+    You can run the following command:
+
+    ```shell
+    aws-annoying load-variables --env-prefix LOAD_AWS_CONFIG__ -- ...
+    ```
+
+    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
+    console = Console(quiet=quiet, emoji=False)
+
+    command = ctx.args
+    if not command:
+        console.print("⚠️ No command provided. Exiting...")
+        raise typer.Exit(0)
+
+    # Mapping of the ARNs by index (index used for ordering)
+    map_arns_by_index = {str(idx): arn for idx, arn in enumerate(arns)}
+    if env_prefix:
+        console.print(f"🔍 Loading ARNs from environment variables with prefix: {env_prefix!r}")
+        arns_env = {
+            key.removeprefix(env_prefix): value for key, value in os.environ.items() if key.startswith(env_prefix)
+        }
+        console.print(f"🔍 Found {len(arns_env)} sources from environment variables.")
+        map_arns_by_index = arns_env | map_arns_by_index
+
+    # Briefly show the ARNs
+    table = Table("Index", "ARN")
+    for idx, arn in sorted(map_arns_by_index.items()):
+        table.add_row(idx, arn)
+
+    console.print(table)
+
+    # Retrieve the variables
+    loader = VariableLoader(dry_run=dry_run)
+    console.print("🔍 Retrieving variables from AWS resources...")
+    if dry_run:
+        console.print("⚠️ 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
+        console.print(f"❌ Failed to load the variables: {exc!s}")
+        raise typer.Exit(1) from None
+
+    console.print(f"✅ Retrieved {load_stats['secrets']} secrets and {load_stats['parameters']} parameters.")
+
+    # Prepare the environment variables
+    env = os.environ.copy()
+    if overwrite_env:
+        env.update(variables)
+    else:
+        # Update variables, preserving the existing ones
+        for key, value in variables.items():
+            env.setdefault(key, str(value))
+
+    # Run the command with the variables injected as environment variables, replacing current process
+    console.print(f"🚀 Running the command: [bold orchid]{' '.join(command)}[/bold orchid]")
+    if replace:  # pragma: no cover (not coverable)
+        os.execvpe(command[0], command, env=env)  # noqa: S606
+        # The above line should never return
+
+    result = subprocess.run(command, env=env, check=False)  # noqa: S603
+    raise typer.Exit(result.returncode)

aws_annoying/{main.py → cli/main.py}

@@ -1,9 +1,10 @@
 # flake8: noqa: F401
 from __future__ import annotations
 
-import aws_annoying.ecs_task_definition_lifecycle
-import aws_annoying.load_variables
-import aws_annoying.mfa
+import aws_annoying.cli.ecs_task_definition_lifecycle
+import aws_annoying.cli.load_variables
+import aws_annoying.cli.mfa
+import aws_annoying.cli.session_manager
 from aws_annoying.utils.debugger import input_as_args
 
 # App with all commands registered

aws_annoying/{mfa → cli/mfa}/_app.py

@@ -1,6 +1,6 @@
 import typer
 
-from aws_annoying.app import app
+from aws_annoying.cli.app import app
 
 mfa_app = typer.Typer(
     no_args_is_help=True,

aws_annoying/{mfa → cli/mfa}/configure.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
-import configparser
 from pathlib import Path  # noqa: TC003
 from typing import Optional
 
 import boto3
 import typer
-from pydantic import BaseModel, ConfigDict
 from rich import print  # noqa: A004
 from rich.prompt import Prompt
 
+from aws_annoying.mfa import MfaConfig, update_credentials
+
 from ._app import mfa_app
 
 _CONFIG_INI_SECTION = "aws-annoying:mfa"
@@ -55,7 +55,7 @@ def configure(  # noqa: PLR0913
     aws_config = aws_config.expanduser()
 
     # Load configuration
-    mfa_config, exists = _MfaConfig.from_ini_file(aws_config, _CONFIG_INI_SECTION)
+    mfa_config, exists = MfaConfig.from_ini_file(aws_config, _CONFIG_INI_SECTION)
     if exists:
         print(f"⚙️ Loaded MFA configuration from AWS config ({aws_config}).")
 
@@ -94,7 +94,7 @@ def configure(  # noqa: PLR0913
 
     # Update MFA profile in AWS credentials
     print(f"✅ Updating MFA profile ([bold]{mfa_profile}[/bold]) to AWS credentials ({aws_credentials})")
-    _update_credentials(
+    update_credentials(
         aws_credentials,
         mfa_profile,  # type: ignore[arg-type]
         access_key=credentials["AccessKeyId"],
@@ -114,44 +114,3 @@ def configure(  # noqa: PLR0913
         mfa_config.save_ini_file(aws_config, _CONFIG_INI_SECTION)
     else:
         print("⚠️ MFA configuration not persisted.")
-
-
-class _MfaConfig(BaseModel):
-    model_config = ConfigDict(extra="ignore")
-
-    mfa_profile: Optional[str] = None
-    mfa_source_profile: Optional[str] = None
-    mfa_serial_number: Optional[str] = None
-
-    def save_ini_file(self, path: Path, section_key: str) -> None:
-        """Save configuration to an AWS config file."""
-        config_ini = configparser.ConfigParser()
-        config_ini.read(path)
-        config_ini.setdefault(section_key, {})
-        for k, v in self.model_dump(exclude_none=True).items():
-            config_ini[section_key][k] = v
-
-        with path.open("w") as f:
-            config_ini.write(f)
-
-    @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."""
-        config_ini = configparser.ConfigParser()
-        config_ini.read(path)
-        if config_ini.has_section(section_key):
-            section = dict(config_ini.items(section_key))
-            return cls.model_validate(section), True
-
-        return cls(), False
-
-
-def _update_credentials(path: Path, profile: str, *, access_key: str, secret_key: str, session_token: str) -> None:
-    credentials_ini = configparser.ConfigParser()
-    credentials_ini.read(path)
-    credentials_ini.setdefault(profile, {})
-    credentials_ini[profile]["aws_access_key_id"] = access_key
-    credentials_ini[profile]["aws_secret_access_key"] = secret_key
-    credentials_ini[profile]["aws_session_token"] = session_token
-    with path.open("w") as f:
-        credentials_ini.write(f)

aws_annoying/cli/session_manager/__init__.py

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

aws_annoying/cli/session_manager/_app.py

@@ -0,0 +1,9 @@
+import typer
+
+from aws_annoying.cli.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/cli/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 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/cli/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/cli/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"])

aws_annoying/cli/session_manager/start.py

@@ -0,0 +1,9 @@
+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/cli/session_manager/stop.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import os
+import signal
+from pathlib import Path  # noqa: TC003
+
+import typer
+from rich import print  # noqa: A004
+
+from ._app import session_manager_app
+
+
+@session_manager_app.command()
+def stop(
+    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
+        True,  # noqa: FBT003
+        help="Remove the PID file after stopping the session.",
+    ),
+) -> None:
+    """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}")
+        raise typer.Exit(1)
+
+    # Read PID from file
+    pid_content = pid_file.read_text()
+    try:
+        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
+
+    # Send SIGTERM to the process
+    try:
+        print(f"⚠️ Terminating running process with PID {pid}.")
+        os.kill(pid, signal.SIGTERM)
+    except ProcessLookupError:
+        print(f"❗ Tried to terminate process with PID {pid} but does not exist.")
+
+    # Remove the PID file
+    if remove:
+        print(f"✅ Removed the PID file {pid_file}.")
+        pid_file.unlink()
+
+    print("✅ Terminated the session successfully.")

aws_annoying/mfa.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import configparser
+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
+
+
+class MfaConfig(BaseModel):
+    """MFA configuration for AWS profiles."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    mfa_profile: Optional[str] = None
+    mfa_source_profile: Optional[str] = None
+    mfa_serial_number: Optional[str] = None
+
+    def save_ini_file(self, path: Path, section_key: str) -> None:
+        """Save configuration to an AWS config file."""
+        config_ini = configparser.ConfigParser()
+        config_ini.read(path)
+        config_ini.setdefault(section_key, {})
+        for k, v in self.model_dump(exclude_none=True).items():
+            config_ini[section_key][k] = v
+
+        with path.open("w") as f:
+            config_ini.write(f)
+
+    @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."""
+        config_ini = configparser.ConfigParser()
+        config_ini.read(path)
+        if config_ini.has_section(section_key):
+            section = dict(config_ini.items(section_key))
+            return cls.model_validate(section), True
+
+        return cls(), False
+
+
+def update_credentials(path: Path, profile: str, *, access_key: str, secret_key: str, session_token: str) -> None:
+    """Update AWS credentials file with the provided profile and credentials."""
+    credentials_ini = configparser.ConfigParser()
+    credentials_ini.read(path)
+    credentials_ini.setdefault(profile, {})
+    credentials_ini[profile]["aws_access_key_id"] = access_key
+    credentials_ini[profile]["aws_secret_access_key"] = secret_key
+    credentials_ini[profile]["aws_session_token"] = session_token
+    with path.open("w") as f:
+        credentials_ini.write(f)

aws_annoying/session_manager/__init__.py

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

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."""