aws-ssh-utils 0.3.3__tar.gz → 0.4.0__tar.gz

{aws_ssh_utils-0.3.3 → aws_ssh_utils-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: aws-ssh-utils
-Version: 0.3.3
+Version: 0.4.0
 Summary: Easy AWS SSHing
 Project-URL: Homepage, https://github.com/mvanderlee/aws-ssh-utils
 Author-email: Michiel Vanderlee <jmt.vanderlee@gmail.com>

aws_ssh_utils-0.4.0/aws_ssh_utils/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.4.0"

{aws_ssh_utils-0.3.3 → aws_ssh_utils-0.4.0}/aws_ssh_utils/commands/app.py

@@ -1,378 +1,492 @@
-import sys
-import time
-from collections.abc import Callable
-from dataclasses import dataclass
-from itertools import count
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
-
-import boto3
-import click
-from botocore.exceptions import BotoCoreError, ClientError
-from loguru import logger
-from rich.console import Console
-from rich.markup import escape
-from rich.spinner import Spinner
-from rich.text import Text
-from textual import work
-from textual.app import ComposeResult
-from textual.binding import Binding, BindingType
-from textual.containers import Horizontal
-from textual.widgets import Footer, Input, OptionList, RichLog, Static, TabbedContent, TabPane, Tree
-from textual.widgets.option_list import Option
-from textual.widgets.tree import TreeNode
-
-from aws_ssh_utils.cli_utils import detect_environment, detect_scrollback, render_attempts, scrollback_option
-from aws_ssh_utils.connection import Attempt, Environment, Target, connect, open_shell_channel
-from aws_ssh_utils.ec2_utils import ec2_name, ec2_target, get_running_ec2_instances
-from aws_ssh_utils.emr_utils import emr_target, get_emr_clusters, get_emr_instances, group_role, group_sort_key
-from aws_ssh_utils.logging_utils import terminal_logging
-from aws_ssh_utils.terminal import ShellStatus, Terminal, TerminalApp
-
-if TYPE_CHECKING:
-    from mypy_boto3_ec2 import EC2Client
-    from mypy_boto3_ec2.type_defs import InstanceTypeDef as EC2InstanceTypeDef
-    from mypy_boto3_emr import EMRClient
-    from mypy_boto3_emr.type_defs import InstanceTypeDef as EMRInstanceTypeDef
-
-Status = Literal['connecting', 'waiting', 'connected', 'failed']
-STATUS_ICONS: dict[Status, Text] = {
-    'waiting': Text('⚠', style='yellow'),
-    'connected': Text('●', style='green'),
-    'failed': Text('●', style='red'),
-}
-ANIMATION_INTERVAL = 1 / 12
-
-
-@click.command("app")
-@click.option('-p', '--profile', default=None, help='Which AWS profile to use')
-@click.option('-r', '--region', default=None, help='Which AWS region to use')
-@scrollback_option
-def app(profile: str | None = None, region: str | None = None, scrollback: int | None = None, **kwargs: Any):
-    """Browse EC2 instances and EMR clusters, and open SSM or SSH shells in tabs."""
-    try:
-        session = boto3.Session(profile_name=profile, region_name=region)
-        session.client('sts').get_caller_identity()
-    except (BotoCoreError, ClientError) as e:
-        logger.bind(file_only=True).opt(exception=e).error('AWS authentication failed')
-        Console().print(f"[red]✗[/] {escape(str(e))}")
-        sys.exit(1)
-
-    env = detect_environment(profile, session.region_name)
-    scrollback = detect_scrollback() if scrollback is None else scrollback
-    with terminal_logging():
-        SSHApp(session.client('ec2'), session.client('emr'), env, scrollback).run()
-
-
-def muted_label(name: str, instance_id: str) -> Text:
-    return Text.assemble(name, (f" - {instance_id}", 'dim'))
-
-
-@dataclass(frozen=True)
-class Cluster:
-    id: str
-    name: str
-
-
-@dataclass(frozen=True)
-class EMRNode:
-    cluster_id: str
-    group_name: str
-    instance: "EMRInstanceTypeDef"
-
-
-class BouncingBall(Static):
-    def on_mount(self):
-        spinner = Spinner('bouncingBall')
-        self.set_interval(ANIMATION_INTERVAL, lambda: self.update(spinner.render(time.monotonic())))
-
-
-class ShellPane(TabPane):
-    """Connects in the background, then swaps its log for a Terminal."""
-
-    def __init__(self, title: str, resolve: Callable[[], Target], env: Environment, scrollback: int, **kwargs: Any):
-        super().__init__(title, **kwargs)
-        self.title_ = title
-        self.resolve = resolve
-        self.env = env
-        self.scrollback = scrollback
-        self.status: Status = 'connecting'
-        self.spinner = Spinner('dots')
-        self.attempts: list[Attempt] = []
-
-    def compose(self) -> ComposeResult:
-        yield Static(classes='attempts')
-        yield RichLog(wrap=True, markup=True)
-
-    def on_mount(self):
-        self.animation = self.set_interval(ANIMATION_INTERVAL, self.update_label)
-        self.connect()
-
-    def update_label(self):
-        icon = cast(Text, self.spinner.render(time.monotonic())) if self.status == 'connecting' else STATUS_ICONS[self.status]
-        self.query_ancestor(TabbedContent).get_tab(self).label = Text.assemble(icon, ' ', self.title_)
-
-    def set_attempts(self, attempts: list[Attempt]):
-        self.attempts = attempts
-        self.query_one('.attempts', Static).update(render_attempts(self.attempts))
-        self.update_label()
-
-    def set_status(self, status: Status):
-        self.status = status
-        if status != 'connecting':
-            self.animation.stop()
-        self.update_label()
-
-    def write(self, markup: str):
-        self.query_one(RichLog).write(markup)
-
-    @work(thread=True, exit_on_error=False)
-    def connect(self):
-        def write(line: str):
-            self.app.call_from_thread(self.write, escape(line))
-
-        try:
-            result = connect(
-                self.resolve(),
-                self.env,
-                on_output=write,
-                on_waiting=lambda: self.app.call_from_thread(self.set_status, 'waiting'),
-                on_progress=lambda attempts: self.app.call_from_thread(self.set_attempts, attempts),
-            )
-            if result is None:
-                self.app.call_from_thread(self.set_status, 'failed')
-                return
-
-            channel = open_shell_channel(result.client)
-        except Exception as e:
-            logger.exception("Connection failed")
-            self.app.call_from_thread(self.write, f"[red]✗ {escape(str(e) or type(e).__name__)}[/]")
-            self.app.call_from_thread(self.set_status, 'failed')
-            return
-
-        self.app.call_from_thread(self.attach, result.via, Terminal(result.client, channel, self.scrollback))
-
-    async def attach(self, via: str, terminal: Terminal):
-        logger.info('Shell tab {} connected via {}', self.id, via)
-        # Stop animation before yielding to widget removal or mounting a reader
-        # that may immediately close the session.
-        self.set_status('connected')
-        await self.query_one(RichLog).remove()
-        await self.query_one('.attempts').remove()
-        await self.mount(ShellStatus(via), terminal)
-
-    def on_terminal_closed(self):
-        logger.info('Shell tab {} closed', self.id)
-        self.query_ancestor(TabbedContent).remove_pane(self.id or '')
-
-
-class SSHApp(TerminalApp):
-    """The textual app"""
-
-    TITLE = "AWS SSH"
-    CSS = """
-    #sidebar { width: 60; }
-    #shells { width: 1fr; }
-    Footer { background: $surface; }
-    TabbedContent, TabbedContent > ContentSwitcher, TabPane { height: 1fr; }
-    #ec2-list, #emr-tree { height: 1fr; }
-    #emr-tree {
-        border: tall $border-blurred;
-        padding: 0 1;
-        background: $surface;
-        &:focus { border: tall $border; }
-    }
-    .error { color: $error; }
-    """
-    BINDINGS: ClassVar[list[BindingType]] = [
-        Binding('ctrl+q', 'quit', 'Quit', priority=True),
-        Binding('f12', 'toggle_sidebar', 'Toggle sidebar', priority=True),
-        Binding('ctrl+w', 'close_tab', 'Close tab'),
-    ]
-
-    def __init__(self, ec2: "EC2Client", emr: "EMRClient", env: Environment, scrollback: int):
-        super().__init__()
-        self.ec2 = ec2
-        self.emr = emr
-        self.env = env
-        self.scrollback = scrollback
-        self.ec2_instances: list[EC2InstanceTypeDef] = []
-        self.clusters: list[Cluster] = []
-        self.cluster_instances: dict[str, dict[str, list[EMRInstanceTypeDef]]] = {}
-        self.emr_loading = 0
-        self.pane_ids = count()
-
-    def compose(self) -> ComposeResult:
-        with Horizontal():
-            with TabbedContent(id='sidebar'):
-                with TabPane('EC2'):
-                    yield Input(placeholder='Filter by name', id='ec2-filter')
-                    yield BouncingBall(id='ec2-loading')
-                    yield Static(id='ec2-error', classes='error')
-                    yield OptionList(id='ec2-list')
-                with TabPane('EMR'):
-                    yield Input(placeholder='Filter by name', id='emr-filter')
-                    yield BouncingBall(id='emr-loading')
-                    yield Static(id='emr-error', classes='error')
-                    yield Tree('Clusters', id='emr-tree')
-            yield TabbedContent(id='shells')
-        yield Footer()
-
-    def on_mount(self):
-        self.query_one('#ec2-error').display = False
-        self.query_one('#emr-error').display = False
-        self.query_one('#emr-tree', Tree).show_root = False
-        self.load_ec2()
-        self.load_emr()
-
-    def show_error(self, prefix: str, error: Exception):
-        self.query_one(f'#{prefix}-loading').display = False
-        error_widget = self.query_one(f'#{prefix}-error', Static)
-        error_widget.update(Text(str(error)))
-        error_widget.display = True
-
-    def set_emr_loading(self, delta: int):
-        self.emr_loading += delta
-        self.query_one('#emr-loading').display = self.emr_loading > 0
-
-    def on_input_changed(self, event: Input.Changed):
-        if event.input.id == 'ec2-filter':
-            self.filter_ec2()
-        else:
-            self.filter_emr()
-
-    def action_toggle_sidebar(self):
-        sidebar = self.query_one('#sidebar')
-        sidebar.display = not sidebar.display
-        if not sidebar.display and self.focused and sidebar in self.focused.ancestors_with_self:
-            pane = self.query_one('#shells', TabbedContent).active_pane
-            self.set_focus(pane.query(Terminal).first() if pane and pane.query(Terminal) else None)
-
-    def action_close_tab(self):
-        shells = self.query_one('#shells', TabbedContent)
-        if shells.active:
-            shells.remove_pane(shells.active)
-
-    def on_tabbed_content_tab_activated(self, event: TabbedContent.TabActivated):
-        if event.tabbed_content.id != 'shells' or event.pane.id != event.tabbed_content.active:
-            return
-        if terminals := event.pane.query(Terminal):
-            terminals.first().focus()
-
-    def open_shell(self, title: str, resolve: Callable[[], Target]):
-        shells = self.query_one('#shells', TabbedContent)
-        pane = ShellPane(title, resolve, self.env, self.scrollback, id=f'shell-{next(self.pane_ids)}')
-        logger.info('Opening shell tab {} ({})', pane.id, title)
-        shells.add_pane(pane)
-        shells.active = pane.id or ''
-
-    # region - EC2
-    @work(thread=True, exit_on_error=False)
-    def load_ec2(self):
-        try:
-            instances = get_running_ec2_instances(self.ec2)
-        except (BotoCoreError, ClientError) as e:
-            self.call_from_thread(self.show_error, 'ec2', e)
-            return
-        self.call_from_thread(self.set_ec2_instances, sorted(instances, key=lambda i: ec2_name(i).lower()))
-
-    def set_ec2_instances(self, instances: list["EC2InstanceTypeDef"]):
-        self.ec2_instances = instances
-        self.query_one('#ec2-loading').display = False
-        self.filter_ec2()
-
-    def filter_ec2(self):
-        name_filter = self.query_one('#ec2-filter', Input).value.lower()
-        option_list = self.query_one('#ec2-list', OptionList)
-        option_list.clear_options()
-        option_list.add_options(
-            Option(muted_label(ec2_name(i), i.get('InstanceId', '')), id=i.get('InstanceId', ''))
-            for i in self.ec2_instances
-            if name_filter in ec2_name(i).lower()
-        )
-
-    def on_option_list_option_selected(self, event: OptionList.OptionSelected):
-        instance = next(i for i in self.ec2_instances if i.get('InstanceId', '') == event.option_id)
-        self.open_shell(ec2_name(instance), lambda: ec2_target(self.ec2, instance))
-
-    # endregion - EC2
-
-    # region - EMR
-    @work(thread=True, exit_on_error=False)
-    def load_emr(self):
-        self.call_from_thread(self.set_emr_loading, 1)
-        try:
-            clusters = [Cluster(cluster_id, name) for cluster_id, name in get_emr_clusters(self.emr).values()]
-        except (BotoCoreError, ClientError) as e:
-            self.call_from_thread(self.show_error, 'emr', e)
-            return
-        finally:
-            self.call_from_thread(self.set_emr_loading, -1)
-        self.call_from_thread(self.set_clusters, sorted(clusters, key=lambda c: c.name.lower()))
-
-    def set_clusters(self, clusters: list[Cluster]):
-        self.clusters = clusters
-        self.filter_emr()
-
-    def filter_emr(self):
-        name_filter = self.query_one('#emr-filter', Input).value.lower()
-        tree = self.query_one('#emr-tree', Tree)
-        tree.clear()
-        for cluster in self.clusters:
-            if name_filter in cluster.name.lower():
-                node = tree.root.add(muted_label(cluster.name, cluster.id), data=cluster)
-                if cluster.id in self.cluster_instances:
-                    self.add_groups(node, cluster.id)
-
-    def on_tree_node_expanded(self, event: Tree.NodeExpanded[Any]):
-        cluster = event.node.data
-        if isinstance(cluster, Cluster) and cluster.id not in self.cluster_instances and not event.node.children:
-            self.load_cluster_instances(cluster.id)
-
-    @work(thread=True, exit_on_error=False)
-    def load_cluster_instances(self, cluster_id: str):
-        self.call_from_thread(self.set_emr_loading, 1)
-        try:
-            groups = get_emr_instances(self.emr, cluster_id)
-        except (BotoCoreError, ClientError, ValueError) as e:
-            self.call_from_thread(self.show_cluster_error, cluster_id, e)
-            return
-        finally:
-            self.call_from_thread(self.set_emr_loading, -1)
-        self.call_from_thread(self.set_cluster_instances, cluster_id, groups)
-
-    def cluster_node(self, cluster_id: str) -> TreeNode[Any] | None:
-        tree = self.query_one('#emr-tree', Tree)
-        return next((n for n in tree.root.children if isinstance(n.data, Cluster) and n.data.id == cluster_id), None)
-
-    def show_cluster_error(self, cluster_id: str, error: Exception):
-        if node := self.cluster_node(cluster_id):
-            node.add_leaf(Text(str(error), style='red'))
-
-    def set_cluster_instances(self, cluster_id: str, groups: dict[str, list["EMRInstanceTypeDef"]]):
-        self.cluster_instances[cluster_id] = groups
-        if node := self.cluster_node(cluster_id):
-            self.add_groups(node, cluster_id)
-
-    def add_groups(self, node: TreeNode[Any], cluster_id: str):
-        groups = self.cluster_instances[cluster_id]
-        for group_name in sorted(groups, key=group_sort_key):
-            group = node.add(group_name, expand=True)
-            for instance in sorted(groups[group_name], key=lambda i: i.get('PrivateIpAddress', '')):
-                label = muted_label(instance.get('PrivateIpAddress', ''), instance.get('Ec2InstanceId', ''))
-                group.add_leaf(label, data=EMRNode(cluster_id, group_name, instance))
-
-    def on_tree_node_selected(self, event: Tree.NodeSelected[Any]):
-        selected = event.node.data
-        if isinstance(selected, EMRNode):
-            title = f"{group_role(selected.group_name)} {selected.instance.get('PrivateIpAddress', '')}"
-            self.open_shell(
-                title,
-                lambda: emr_target(
-                    self.emr,
-                    selected.cluster_id,
-                    selected.instance.get('Ec2InstanceId', ''),
-                    selected.instance.get('PrivateIpAddress', ''),
-                ),
-            )
-
-    # endregion - EMR
+import sys
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, replace
+from itertools import count
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
+
+import boto3
+import click
+from botocore.exceptions import BotoCoreError, ClientError
+from loguru import logger
+from rich.console import Console, Group
+from rich.markup import escape
+from rich.spinner import Spinner
+from rich.table import Table
+from rich.text import Text
+from textual import events, work
+from textual.app import ComposeResult
+from textual.binding import Binding, BindingType
+from textual.containers import Horizontal
+from textual.geometry import Offset
+from textual.screen import ModalScreen
+from textual.widgets import Footer, Input, OptionList, RichLog, Static, TabbedContent, TabPane, Tree
+from textual.widgets.option_list import Option
+from textual.widgets.tree import TreeNode
+
+from aws_ssh_utils.cli_utils import detect_environment, detect_scrollback, render_attempts, scrollback_option
+from aws_ssh_utils.connection import Attempt, Environment, Target, connect, open_shell_channel
+from aws_ssh_utils.ec2_utils import ec2_name, ec2_target, get_running_ec2_instances, get_vpc_names
+from aws_ssh_utils.emr_utils import emr_target, get_emr_clusters, get_emr_instances, group_role, group_sort_key
+from aws_ssh_utils.logging_utils import terminal_logging
+from aws_ssh_utils.terminal import ShellStatus, Terminal, TerminalApp
+
+if TYPE_CHECKING:
+    from mypy_boto3_ec2 import EC2Client
+    from mypy_boto3_ec2.type_defs import InstanceTypeDef as EC2InstanceTypeDef
+    from mypy_boto3_emr import EMRClient
+    from mypy_boto3_emr.type_defs import InstanceTypeDef as EMRInstanceTypeDef
+
+Status = Literal['connecting', 'waiting', 'connected', 'failed']
+STATUS_ICONS: dict[Status, Text] = {
+    'waiting': Text('⚠', style='yellow'),
+    'connected': Text('●', style='green'),
+    'failed': Text('●', style='red'),
+}
+ANIMATION_INTERVAL = 1 / 12
+RIGHT_MOUSE_BUTTON = 3
+PUBLIC_IP_ICON = '\U000f01e7'  # nf-md-earth
+PRIVATE_IP_ICON = ''  # nf-oct-lock
+INSTANCE_TYPE_ICON = ''  # nf-oct-cpu
+
+
+@click.command("app")
+@click.option('-p', '--profile', default=None, help='Which AWS profile to use')
+@click.option('-r', '--region', default=None, help='Which AWS region to use')
+@scrollback_option
+def app(profile: str | None = None, region: str | None = None, scrollback: int | None = None, **kwargs: Any):
+    """Browse EC2 instances and EMR clusters, and open SSM or SSH shells in tabs."""
+    try:
+        session = boto3.Session(profile_name=profile, region_name=region)
+        session.client('sts').get_caller_identity()
+    except (BotoCoreError, ClientError) as e:
+        logger.bind(file_only=True).opt(exception=e).error('AWS authentication failed')
+        Console().print(f"[red]✗[/] {escape(str(e))}")
+        sys.exit(1)
+
+    env = detect_environment(profile, session.region_name)
+    scrollback = detect_scrollback() if scrollback is None else scrollback
+    with terminal_logging():
+        SSHApp(session.client('ec2'), session.client('emr'), env, scrollback).run()
+
+
+def muted_label(name: str, instance_id: str) -> Text:
+    return Text.assemble(name, (f" - {instance_id}", 'dim'))
+
+
+def ec2_label(instance: "EC2InstanceTypeDef") -> Table:
+    """Name and id, with a right-aligned globe when the instance has a public IP."""
+    grid = Table.grid(expand=True)
+    grid.add_column(ratio=1)
+    grid.add_column(justify='right')
+    grid.add_row(
+        muted_label(ec2_name(instance), instance.get('InstanceId', '')),
+        PUBLIC_IP_ICON if instance.get('PublicIpAddress') else '',
+    )
+    return grid
+
+
+@dataclass(frozen=True)
+class Cluster:
+    id: str
+    name: str
+
+
+@dataclass(frozen=True)
+class EMRNode:
+    cluster_id: str
+    group_name: str
+    instance: "EMRInstanceTypeDef"
+
+
+class BouncingBall(Static):
+    def on_mount(self):
+        spinner = Spinner('bouncingBall')
+        self.set_interval(ANIMATION_INTERVAL, lambda: self.update(spinner.render(time.monotonic())))
+
+
+class EC2List(OptionList):
+    """Only the left button selects; right clicks bubble up to open a context menu."""
+
+    async def _on_click(self, event: events.Click) -> None:
+        if event.button == RIGHT_MOUSE_BUTTON:
+            event.prevent_default()
+
+
+class ContextMenu(ModalScreen[str | None]):
+    """Options at the mouse position. Dismisses with the chosen option id, or None."""
+
+    DEFAULT_CSS = """
+    ContextMenu { background: transparent; }
+    ContextMenu > OptionList {
+        position: absolute;
+        constrain: inside inside;
+        width: auto;
+        height: auto;
+        padding: 0;
+        &, &:focus { border: round $border; }
+    }
+    """
+    BINDINGS: ClassVar[list[BindingType]] = [Binding('escape', 'dismiss', 'Close')]
+
+    def __init__(self, offset: Offset, options: list[Option]):
+        super().__init__()
+        self.menu_offset = offset
+        self.menu_options = options
+
+    def compose(self) -> ComposeResult:
+        yield OptionList(*self.menu_options)
+
+    def on_mount(self):
+        menu = self.query_one(OptionList)
+        menu.styles.offset = self.menu_offset
+        menu.focus()
+
+    def on_option_list_option_selected(self, event: OptionList.OptionSelected):
+        event.stop()
+        self.dismiss(event.option_id)
+
+    def on_click(self, event: events.Click):
+        if event.widget is self:
+            self.dismiss()
+
+
+class ShellPane(TabPane):
+    """Connects in the background, then swaps its log for a Terminal."""
+
+    def __init__(self, title: str, resolve: Callable[[], Target], env: Environment, scrollback: int, **kwargs: Any):
+        super().__init__(title, **kwargs)
+        self.title_ = title
+        self.resolve = resolve
+        self.env = env
+        self.scrollback = scrollback
+        self.status: Status = 'connecting'
+        self.spinner = Spinner('dots')
+        self.attempts: list[Attempt] = []
+
+    def compose(self) -> ComposeResult:
+        yield Static(classes='attempts')
+        yield RichLog(wrap=True, markup=True)
+
+    def on_mount(self):
+        self.animation = self.set_interval(ANIMATION_INTERVAL, self.update_label)
+        self.connect()
+
+    def update_label(self):
+        icon = cast(Text, self.spinner.render(time.monotonic())) if self.status == 'connecting' else STATUS_ICONS[self.status]
+        self.query_ancestor(TabbedContent).get_tab(self).label = Text.assemble(icon, ' ', self.title_)
+
+    def set_attempts(self, attempts: list[Attempt]):
+        self.attempts = attempts
+        self.query_one('.attempts', Static).update(render_attempts(self.attempts))
+        self.update_label()
+
+    def set_status(self, status: Status):
+        self.status = status
+        if status != 'connecting':
+            self.animation.stop()
+        self.update_label()
+
+    def write(self, markup: str):
+        self.query_one(RichLog).write(markup)
+
+    @work(thread=True, exit_on_error=False)
+    def connect(self):
+        def write(line: str):
+            self.app.call_from_thread(self.write, escape(line))
+
+        try:
+            result = connect(
+                self.resolve(),
+                self.env,
+                on_output=write,
+                on_waiting=lambda: self.app.call_from_thread(self.set_status, 'waiting'),
+                on_progress=lambda attempts: self.app.call_from_thread(self.set_attempts, attempts),
+            )
+            if result is None:
+                self.app.call_from_thread(self.set_status, 'failed')
+                return
+
+            channel = open_shell_channel(result.client)
+        except Exception as e:
+            logger.exception("Connection failed")
+            self.app.call_from_thread(self.write, f"[red]✗ {escape(str(e) or type(e).__name__)}[/]")
+            self.app.call_from_thread(self.set_status, 'failed')
+            return
+
+        self.app.call_from_thread(self.attach, result.via, Terminal(result.client, channel, self.scrollback))
+
+    async def attach(self, via: str, terminal: Terminal):
+        logger.info('Shell tab {} connected via {}', self.id, via)
+        # Stop animation before yielding to widget removal or mounting a reader
+        # that may immediately close the session.
+        self.set_status('connected')
+        await self.query_one(RichLog).remove()
+        await self.query_one('.attempts').remove()
+        await self.mount(ShellStatus(via), terminal)
+
+    def on_terminal_closed(self):
+        logger.info('Shell tab {} closed', self.id)
+        self.query_ancestor(TabbedContent).remove_pane(self.id or '')
+
+
+class SSHApp(TerminalApp):
+    """The textual app"""
+
+    TITLE = "AWS SSH"
+    CSS = """
+    #sidebar { width: 60; }
+    #shells { width: 1fr; }
+    Footer { background: $surface; }
+    TabbedContent, TabbedContent > ContentSwitcher, TabPane { height: 1fr; }
+    #ec2-list, #emr-tree { height: 1fr; }
+    #ec2-details { height: auto; }
+    #emr-tree, #ec2-details {
+        border: tall $border-blurred;
+        padding: 0 1;
+        background: $surface;
+        &:focus { border: tall $border; }
+    }
+    .error { color: $error; }
+    """
+    BINDINGS: ClassVar[list[BindingType]] = [
+        Binding('ctrl+q', 'quit', 'Quit', priority=True),
+        Binding('f12', 'toggle_sidebar', 'Toggle sidebar', priority=True),
+        Binding('ctrl+w', 'close_tab', 'Close tab'),
+    ]
+
+    def __init__(self, ec2: "EC2Client", emr: "EMRClient", env: Environment, scrollback: int):
+        super().__init__()
+        self.ec2 = ec2
+        self.emr = emr
+        self.env = env
+        self.scrollback = scrollback
+        self.ec2_instances: list[EC2InstanceTypeDef] = []
+        self.vpc_names: dict[str, str] = {}
+        self.clusters: list[Cluster] = []
+        self.cluster_instances: dict[str, dict[str, list[EMRInstanceTypeDef]]] = {}
+        self.emr_loading = 0
+        self.pane_ids = count()
+
+    def compose(self) -> ComposeResult:
+        with Horizontal():
+            with TabbedContent(id='sidebar'):
+                with TabPane('EC2'):
+                    yield Input(placeholder='Filter by name', id='ec2-filter')
+                    yield BouncingBall(id='ec2-loading')
+                    yield Static(id='ec2-error', classes='error')
+                    yield EC2List(id='ec2-list')
+                    yield Static(id='ec2-details')
+                with TabPane('EMR'):
+                    yield Input(placeholder='Filter by name', id='emr-filter')
+                    yield BouncingBall(id='emr-loading')
+                    yield Static(id='emr-error', classes='error')
+                    yield Tree('Clusters', id='emr-tree')
+            yield TabbedContent(id='shells')
+        yield Footer()
+
+    def on_mount(self):
+        self.query_one('#ec2-error').display = False
+        self.query_one('#emr-error').display = False
+        self.query_one('#emr-tree', Tree).show_root = False
+        self.load_ec2()
+        self.load_emr()
+
+    def show_error(self, prefix: str, error: Exception):
+        self.query_one(f'#{prefix}-loading').display = False
+        error_widget = self.query_one(f'#{prefix}-error', Static)
+        error_widget.update(Text(str(error)))
+        error_widget.display = True
+
+    def set_emr_loading(self, delta: int):
+        self.emr_loading += delta
+        self.query_one('#emr-loading').display = self.emr_loading > 0
+
+    def on_input_changed(self, event: Input.Changed):
+        if event.input.id == 'ec2-filter':
+            self.filter_ec2()
+        else:
+            self.filter_emr()
+
+    def action_toggle_sidebar(self):
+        sidebar = self.query_one('#sidebar')
+        sidebar.display = not sidebar.display
+        if not sidebar.display and self.focused and sidebar in self.focused.ancestors_with_self:
+            pane = self.query_one('#shells', TabbedContent).active_pane
+            self.set_focus(pane.query(Terminal).first() if pane and pane.query(Terminal) else None)
+
+    def action_close_tab(self):
+        shells = self.query_one('#shells', TabbedContent)
+        if shells.active:
+            shells.remove_pane(shells.active)
+
+    def on_tabbed_content_tab_activated(self, event: TabbedContent.TabActivated):
+        if event.tabbed_content.id != 'shells' or event.pane.id != event.tabbed_content.active:
+            return
+        if terminals := event.pane.query(Terminal):
+            terminals.first().focus()
+
+    def open_shell(self, title: str, resolve: Callable[[], Target]):
+        shells = self.query_one('#shells', TabbedContent)
+        pane = ShellPane(title, resolve, self.env, self.scrollback, id=f'shell-{next(self.pane_ids)}')
+        logger.info('Opening shell tab {} ({})', pane.id, title)
+        shells.add_pane(pane)
+        shells.active = pane.id or ''
+
+    # region - EC2
+    @work(thread=True, exit_on_error=False)
+    def load_ec2(self):
+        try:
+            instances = get_running_ec2_instances(self.ec2)
+        except (BotoCoreError, ClientError) as e:
+            self.call_from_thread(self.show_error, 'ec2', e)
+            return
+        try:
+            vpc_names = get_vpc_names(self.ec2, sorted({vpc_id for i in instances if (vpc_id := i.get('VpcId'))}))
+        except (BotoCoreError, ClientError) as e:
+            logger.warning('Could not load VPC names: {!r}', e)
+            vpc_names = {}
+        self.call_from_thread(self.set_ec2_instances, sorted(instances, key=lambda i: ec2_name(i).lower()), vpc_names)
+
+    def set_ec2_instances(self, instances: list["EC2InstanceTypeDef"], vpc_names: dict[str, str] | None = None):
+        self.ec2_instances = instances
+        self.vpc_names = vpc_names or {}
+        self.query_one('#ec2-loading').display = False
+        self.filter_ec2()
+
+    def filter_ec2(self):
+        name_filter = self.query_one('#ec2-filter', Input).value.lower()
+        option_list = self.query_one('#ec2-list', OptionList)
+        option_list.clear_options()
+        self.query_one('#ec2-details', Static).update('')
+        option_list.add_options(
+            Option(ec2_label(i), id=i.get('InstanceId', '')) for i in self.ec2_instances if name_filter in ec2_name(i).lower()
+        )
+
+    def ec2_instance(self, instance_id: str | None) -> "EC2InstanceTypeDef":
+        return next(i for i in self.ec2_instances if i.get('InstanceId', '') == instance_id)
+
+    def open_ec2_shell(self, instance: "EC2InstanceTypeDef", public: bool = False):
+        def resolve() -> Target:
+            target = ec2_target(self.ec2, instance)
+            return replace(target, ip=instance.get('PublicIpAddress', '')) if public else target
+
+        self.open_shell(f'{ec2_name(instance)} (public)' if public else ec2_name(instance), resolve)
+
+    def on_option_list_option_selected(self, event: OptionList.OptionSelected):
+        self.open_ec2_shell(self.ec2_instance(event.option_id))
+
+    def on_option_list_option_highlighted(self, event: OptionList.OptionHighlighted):
+        if not isinstance(event.option_list, EC2List):
+            return
+        instance = self.ec2_instance(event.option_id)
+        vpc_id = instance.get('VpcId', '')
+        fields = Table.grid(padding=(0, 1))
+        for label, value in (
+            (INSTANCE_TYPE_ICON, instance.get('InstanceType', '')),
+            ('VPC:', self.vpc_names.get(vpc_id, vpc_id)),
+            (PRIVATE_IP_ICON, instance.get('PrivateIpAddress', '')),
+            (PUBLIC_IP_ICON, instance.get('PublicIpAddress', '')),
+        ):
+            if value:
+                fields.add_row(label, value)
+        self.query_one('#ec2-details', Static).update(Group(Text(ec2_name(instance), style='bold'), fields))
+
+    def on_click(self, event: events.Click):
+        index = event.style.meta.get('option')
+        if event.button != RIGHT_MOUSE_BUTTON or not isinstance(event.widget, EC2List) or index is None:
+            return
+        instance = self.ec2_instance(event.widget.get_option_at_index(index).id)
+        options = [
+            Option('Connect via private IP', id='private'),
+            Option('Connect via public IP', id='public', disabled=not instance.get('PublicIpAddress')),
+        ]
+
+        def on_choice(choice: str | None):
+            if choice:
+                self.open_ec2_shell(instance, public=choice == 'public')
+
+        self.push_screen(ContextMenu(event.screen_offset, options), on_choice)
+
+    # endregion - EC2
+
+    # region - EMR
+    @work(thread=True, exit_on_error=False)
+    def load_emr(self):
+        self.call_from_thread(self.set_emr_loading, 1)
+        try:
+            clusters = [Cluster(cluster_id, name) for cluster_id, name in get_emr_clusters(self.emr).values()]
+        except (BotoCoreError, ClientError) as e:
+            self.call_from_thread(self.show_error, 'emr', e)
+            return
+        finally:
+            self.call_from_thread(self.set_emr_loading, -1)
+        self.call_from_thread(self.set_clusters, sorted(clusters, key=lambda c: c.name.lower()))
+
+    def set_clusters(self, clusters: list[Cluster]):
+        self.clusters = clusters
+        self.filter_emr()
+
+    def filter_emr(self):
+        name_filter = self.query_one('#emr-filter', Input).value.lower()
+        tree = self.query_one('#emr-tree', Tree)
+        tree.clear()
+        for cluster in self.clusters:
+            if name_filter in cluster.name.lower():
+                node = tree.root.add(muted_label(cluster.name, cluster.id), data=cluster)
+                if cluster.id in self.cluster_instances:
+                    self.add_groups(node, cluster.id)
+
+    def on_tree_node_expanded(self, event: Tree.NodeExpanded[Any]):
+        cluster = event.node.data
+        if isinstance(cluster, Cluster) and cluster.id not in self.cluster_instances and not event.node.children:
+            self.load_cluster_instances(cluster.id)
+
+    @work(thread=True, exit_on_error=False)
+    def load_cluster_instances(self, cluster_id: str):
+        self.call_from_thread(self.set_emr_loading, 1)
+        try:
+            groups = get_emr_instances(self.emr, cluster_id)
+        except (BotoCoreError, ClientError, ValueError) as e:
+            self.call_from_thread(self.show_cluster_error, cluster_id, e)
+            return
+        finally:
+            self.call_from_thread(self.set_emr_loading, -1)
+        self.call_from_thread(self.set_cluster_instances, cluster_id, groups)
+
+    def cluster_node(self, cluster_id: str) -> TreeNode[Any] | None:
+        tree = self.query_one('#emr-tree', Tree)
+        return next((n for n in tree.root.children if isinstance(n.data, Cluster) and n.data.id == cluster_id), None)
+
+    def show_cluster_error(self, cluster_id: str, error: Exception):
+        if node := self.cluster_node(cluster_id):
+            node.add_leaf(Text(str(error), style='red'))
+
+    def set_cluster_instances(self, cluster_id: str, groups: dict[str, list["EMRInstanceTypeDef"]]):
+        self.cluster_instances[cluster_id] = groups
+        if node := self.cluster_node(cluster_id):
+            self.add_groups(node, cluster_id)
+
+    def add_groups(self, node: TreeNode[Any], cluster_id: str):
+        groups = self.cluster_instances[cluster_id]
+        for group_name in sorted(groups, key=group_sort_key):
+            group = node.add(group_name, expand=True)
+            for instance in sorted(groups[group_name], key=lambda i: i.get('PrivateIpAddress', '')):
+                label = muted_label(instance.get('PrivateIpAddress', ''), instance.get('Ec2InstanceId', ''))
+                group.add_leaf(label, data=EMRNode(cluster_id, group_name, instance))
+
+    def on_tree_node_selected(self, event: Tree.NodeSelected[Any]):
+        selected = event.node.data
+        if isinstance(selected, EMRNode):
+            title = f"{group_role(selected.group_name)} {selected.instance.get('PrivateIpAddress', '')}"
+            self.open_shell(
+                title,
+                lambda: emr_target(
+                    self.emr,
+                    selected.cluster_id,
+                    selected.instance.get('Ec2InstanceId', ''),
+                    selected.instance.get('PrivateIpAddress', ''),
+                ),
+            )
+
+    # endregion - EMR

{aws_ssh_utils-0.3.3 → aws_ssh_utils-0.4.0}/aws_ssh_utils/ec2_utils.py

@@ -1,68 +1,77 @@
-from typing import TYPE_CHECKING
-
-import questionary
-
-from aws_ssh_utils.cli_utils import ShellError, spinner
-from aws_ssh_utils.connection import Target, opkssh_provider
-
-if TYPE_CHECKING:
-    from mypy_boto3_ec2 import EC2Client
-    from mypy_boto3_ec2.type_defs import InstanceTypeDef
-
-
-def get_running_ec2_instances(ec2: "EC2Client") -> list["InstanceTypeDef"]:
-    pages = ec2.get_paginator('describe_instances').paginate(
-        Filters=[{'Name': 'instance-state-name', 'Values': ['running']}],
-    )
-    return [i for page in pages for r in page['Reservations'] for i in r.get('Instances', [])]
-
-
-def ec2_name(instance: "InstanceTypeDef") -> str:
-    """The Name tag, or the instance id when there is none."""
-    return next((t.get('Value', '') for t in instance.get('Tags', []) if t.get('Key') == 'Name'), instance.get('InstanceId', ''))
-
-
-def guess_ec2_user(ec2: "EC2Client", instance: "InstanceTypeDef") -> str:
-    image_id = instance.get('ImageId')
-    images = ec2.describe_images(ImageIds=[image_id]).get('Images', []) if image_id else []
-    image_name = images[0].get('Name', '') if images else ''
-    return 'ubuntu' if 'ubuntu' in image_name.lower() else 'ec2-user'
-
-
-def ec2_target(ec2: "EC2Client", instance: "InstanceTypeDef", user: str | None = None) -> Target:
-    """Connection target on the private IP. Guesses the user from the AMI name when not given."""
-    tags = {t.get('Key', '').lower(): t.get('Value', '') for t in instance.get('Tags', [])}
-    return Target(
-        instance_id=instance.get('InstanceId', ''),
-        ip=instance.get('PrivateIpAddress', ''),
-        user=user or guess_ec2_user(ec2, instance),
-        aliases=(ec2_name(instance),),
-        key_name=instance.get('KeyName'),
-        opkssh_provider=opkssh_provider(tags),
-    )
-
-
-def prompt_for_ec2_instance(
-    ec2: "EC2Client",
-    prompt: str = 'Which EC2 instance do you want to connect to?',
-) -> "InstanceTypeDef":
-    """Asks for a name filter, then which of the matching running instances to use."""
-    with spinner():
-        instances = get_running_ec2_instances(ec2)
-
-    name_contains = questionary.text('Provide a name filter or leave blank to show all').unsafe_ask().lower()
-
-    by_display_name: dict[str, InstanceTypeDef] = {}
-    for instance in instances:
-        name = ec2_name(instance)
-        if name_contains and name_contains not in name.lower():
-            continue
-        # Disambiguate duplicate Name tags (e.g. ASG members) by appending the instance id
-        display = name if name not in by_display_name else f'{name} ({instance.get("InstanceId")})'
-        by_display_name[display] = instance
-
-    if not by_display_name:
-        raise ShellError('No matching EC2 instances found!')
-
-    choice = questionary.select(prompt, choices=sorted(by_display_name)).unsafe_ask()
-    return by_display_name[choice]
+from typing import TYPE_CHECKING
+
+import questionary
+
+from aws_ssh_utils.cli_utils import ShellError, spinner
+from aws_ssh_utils.connection import Target, opkssh_provider
+
+if TYPE_CHECKING:
+    from mypy_boto3_ec2 import EC2Client
+    from mypy_boto3_ec2.type_defs import InstanceTypeDef
+
+
+def get_running_ec2_instances(ec2: "EC2Client") -> list["InstanceTypeDef"]:
+    pages = ec2.get_paginator('describe_instances').paginate(
+        Filters=[{'Name': 'instance-state-name', 'Values': ['running']}],
+    )
+    return [i for page in pages for r in page['Reservations'] for i in r.get('Instances', [])]
+
+
+def get_vpc_names(ec2: "EC2Client", vpc_ids: list[str]) -> dict[str, str]:
+    """VPC id to its Name tag, or the id when there is none."""
+    vpcs = ec2.describe_vpcs(VpcIds=vpc_ids).get('Vpcs', []) if vpc_ids else []
+    return {
+        v.get('VpcId', ''): next((t.get('Value', '') for t in v.get('Tags', []) if t.get('Key') == 'Name'), v.get('VpcId', ''))
+        for v in vpcs
+    }
+
+
+def ec2_name(instance: "InstanceTypeDef") -> str:
+    """The Name tag, or the instance id when there is none."""
+    return next((t.get('Value', '') for t in instance.get('Tags', []) if t.get('Key') == 'Name'), instance.get('InstanceId', ''))
+
+
+def guess_ec2_user(ec2: "EC2Client", instance: "InstanceTypeDef") -> str:
+    image_id = instance.get('ImageId')
+    images = ec2.describe_images(ImageIds=[image_id]).get('Images', []) if image_id else []
+    image_name = images[0].get('Name', '') if images else ''
+    return 'ubuntu' if 'ubuntu' in image_name.lower() else 'ec2-user'
+
+
+def ec2_target(ec2: "EC2Client", instance: "InstanceTypeDef", user: str | None = None) -> Target:
+    """Connection target on the private IP. Guesses the user from the AMI name when not given."""
+    tags = {t.get('Key', '').lower(): t.get('Value', '') for t in instance.get('Tags', [])}
+    return Target(
+        instance_id=instance.get('InstanceId', ''),
+        ip=instance.get('PrivateIpAddress', ''),
+        user=user or guess_ec2_user(ec2, instance),
+        aliases=(ec2_name(instance),),
+        key_name=instance.get('KeyName'),
+        opkssh_provider=opkssh_provider(tags),
+    )
+
+
+def prompt_for_ec2_instance(
+    ec2: "EC2Client",
+    prompt: str = 'Which EC2 instance do you want to connect to?',
+) -> "InstanceTypeDef":
+    """Asks for a name filter, then which of the matching running instances to use."""
+    with spinner():
+        instances = get_running_ec2_instances(ec2)
+
+    name_contains = questionary.text('Provide a name filter or leave blank to show all').unsafe_ask().lower()
+
+    by_display_name: dict[str, InstanceTypeDef] = {}
+    for instance in instances:
+        name = ec2_name(instance)
+        if name_contains and name_contains not in name.lower():
+            continue
+        # Disambiguate duplicate Name tags (e.g. ASG members) by appending the instance id
+        display = name if name not in by_display_name else f'{name} ({instance.get("InstanceId")})'
+        by_display_name[display] = instance
+
+    if not by_display_name:
+        raise ShellError('No matching EC2 instances found!')
+
+    choice = questionary.select(prompt, choices=sorted(by_display_name)).unsafe_ask()
+    return by_display_name[choice]

aws_ssh_utils-0.3.3/aws_ssh_utils/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.3.3"