aws-ssh-utils 0.1.1__tar.gz → 0.2.1__tar.gz

aws_ssh_utils-0.2.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 mvanderlee
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{aws_ssh_utils-0.1.1 → aws_ssh_utils-0.2.1}/PKG-INFO

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: aws-ssh-utils
-Version: 0.1.1
+Version: 0.2.1
 Summary: Easy AWS SSHing
-Project-URL: Homepage, https://github.com/mvanderlee/aws_ssh_utils
+Project-URL: Homepage, https://github.com/mvanderlee/aws-ssh-utils
 Author-email: Michiel Vanderlee <jmt.vanderlee@gmail.com>
 License-Expression: MIT
+License-File: LICENSE
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Information Technology
 Classifier: Intended Audience :: System Administrators
@@ -28,6 +29,7 @@ Requires-Dist: environs
 Requires-Dist: loguru
 Requires-Dist: paramiko
 Requires-Dist: questionary
+Requires-Dist: typing-extensions
 Provides-Extra: dev
 Requires-Dist: boto3-stubs[ec2,emr]; extra == 'dev'
 Provides-Extra: publish
@@ -37,6 +39,10 @@ Description-Content-Type: text/markdown
 
 # AWS SSH Utils
 
+[![Version](https://img.shields.io/pypi/v/aws-ssh-utils.svg)](https://pypi.org/project/aws-ssh-utils/)
+[![License](https://img.shields.io/pypi/l/aws-ssh-utils.svg)](#)
+[![Supported Python Versions](https://img.shields.io/pypi/pyversions/aws-ssh-utils.svg)](https://pypi.org/project/aws-ssh-utils/)
+
 ```shell
 pip install aws-ssh-utils
 

{aws_ssh_utils-0.1.1 → aws_ssh_utils-0.2.1}/README.md

@@ -1,5 +1,9 @@
 # AWS SSH Utils
 
+[![Version](https://img.shields.io/pypi/v/aws-ssh-utils.svg)](https://pypi.org/project/aws-ssh-utils/)
+[![License](https://img.shields.io/pypi/l/aws-ssh-utils.svg)](#)
+[![Supported Python Versions](https://img.shields.io/pypi/pyversions/aws-ssh-utils.svg)](https://pypi.org/project/aws-ssh-utils/)
+
 ```shell
 pip install aws-ssh-utils
 

aws_ssh_utils-0.2.1/aws_ssh_utils/__init__.py

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

{aws_ssh_utils-0.1.1 → aws_ssh_utils-0.2.1}/aws_ssh_utils/emr_utils.py

@@ -1,25 +1,27 @@
+from collections.abc import MutableMapping, Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, MutableMapping
+from typing import TYPE_CHECKING, Literal
 
 import click_spinner
 import questionary
 
 if TYPE_CHECKING:
     from mypy_boto3_emr import EMRClient
-    from mypy_boto3_emr.type_defs import ClusterSummaryTypeDef, InstanceTypeDef
+    from mypy_boto3_emr.literals import ClusterStateType
+    from mypy_boto3_emr.type_defs import ClusterSummaryTypeDef, InstanceFleetTypeDef, InstanceGroupTypeDef, InstanceTypeDef
 
 
 @dataclass
 class IP:
-    private: str
-    public: str = None
+    private: str | None = None
+    public: str | None = None
 
 
 def prompt_for_emr_cluster(
     emr: "EMRClient",
-    prompt="Which cluster do you want to connect to?",
+    prompt: str = "Which cluster do you want to connect to?",
     # Either a list of application names, or a dict[name, version]
-    applications: dict[str, str] | list[str] = None,
+    applications: dict[str, str] | list[str] | None = None,
 ) -> tuple[str, str]:
     """Discovers the available EMR clusters and asks the user which to use.
 
@@ -41,7 +43,8 @@ def prompt_for_emr_cluster(
 
 def prompt_for_emr_instance_group(
     emr: "EMRClient",
-    cluster_id: str, prompt="Which instance group do you want to connect to?"
+    cluster_id: str,
+    prompt: str = "Which instance group do you want to connect to?",
 ) -> tuple[list[IP], str]:
     """Discovers the available instance groups for the selected EMR cluster.
         Note that instance group and instance fleet are both treated the same here.
@@ -64,16 +67,20 @@ def prompt_for_emr_instance_group(
 
 def does_cluster_have_applications(
     emr: "EMRClient",
-    cluster_id: str,
-    applications: dict[str, str] | list[str] = None,
+    cluster_id: str | None = None,
+    applications: dict[str, str] | list[str] | None = None,
 ) -> bool:
-    cluster_details = emr.describe_cluster(ClusterId=cluster_id).get("Cluster")
+    if applications is None or cluster_id is None:
+        return True
 
+    cluster_details = emr.describe_cluster(ClusterId=cluster_id).get("Cluster")
     if not cluster_details:
         return False
 
-    cluster_applications = {a["Name"]: a["Version"] for a in cluster_details["Applications"]}
-
+    cluster_applications = {
+        a.get("Name"): a.get("Version")
+        for a in cluster_details.get("Applications", [])
+    }
     if isinstance(applications, MutableMapping):
         # If a map, check that all the required applications are in the cluster and at the version
         return all(v == cluster_applications.get(k) for k, v in applications.items())
@@ -83,9 +90,9 @@ def does_cluster_have_applications(
 
 def get_emr_clusters(
     emr: "EMRClient",
-    states: list[str] = None,
+    states: Sequence["ClusterStateType"] | None = None,
     # Either a list of application names, or a dict[name, version]
-    applications: dict[str, str] | list[str] = None,
+    applications: dict[str, str] | list[str] | None = None,
 ) -> dict[str, tuple[str, str]]:
     """Discover the available EMR cluster
 
@@ -99,25 +106,37 @@ def get_emr_clusters(
     ).get("Clusters", [])
 
     if applications:
-        clusters = [c for c in clusters if does_cluster_have_applications(emr, c["Id"], applications=applications)]
+        clusters = [
+            c for c in clusters
+            if does_cluster_have_applications(emr, c.get("Id"), applications=applications)
+        ]
 
     def get_display_name(c: "ClusterSummaryTypeDef") -> str:
-        created = c["Status"]["Timeline"]["CreationDateTime"].strftime("%Y-%m-%d %H:%M")
-        return f"{c['Name']} - {c['Id']} - {created}"
+        created = c.get("Status", {}).get(
+            "Timeline", {},
+        ).get("CreationDateTime")
+        if created:
+            created = created.strftime("%Y-%m-%d %H:%M")
+        return f"{c.get('Name')} - {c.get('Id')} - {created}"
 
-    return {f"{get_display_name(c)}": (c["Id"], c["Name"]) for c in clusters}
+    return {
+        f"{get_display_name(c)}": (c.get("Id", ""), c.get("Name", ""))
+        for c in clusters
+    }
 
 
 def get_emr_groups(
     emr: "EMRClient",
     cluster_id: str,
-):
+) -> tuple[list["InstanceFleetTypeDef"] | list["InstanceGroupTypeDef"], Literal['InstanceFleetId', 'InstanceGroupId']]:
     """Discover the EMR groups.
 
         Returns the groups and the id_name to use for list_instances
     """
     cluster_details = emr.describe_cluster(ClusterId=cluster_id)
-    instance_collection_type = cluster_details["Cluster"].get("InstanceCollectionType")
+    instance_collection_type = cluster_details["Cluster"].get(
+        "InstanceCollectionType",
+    )
 
     if instance_collection_type == "INSTANCE_FLEET":
         id_name = "InstanceFleetId"
@@ -129,6 +148,11 @@ def get_emr_groups(
         groups = emr.list_instance_groups(ClusterId=cluster_id)
         groups = groups["InstanceGroups"]
 
+    else:
+        raise ValueError(
+            f"Unknown instance collection type: {instance_collection_type}",
+        )
+
     return groups, id_name
 
 
@@ -139,12 +163,14 @@ def get_emr_instances(
     """Discover the running EMR instances per instance group"""
     groups, id_name = get_emr_groups(emr, cluster_id)
 
-    groups = {g["Id"]: g["Name"] for g in groups}
+    groups = {g.get("Id", ""): g.get("Name", "") for g in groups}
     grouped_instances = {g: [] for g in groups.values()}
 
-    instances = emr.list_instances(ClusterId=cluster_id, InstanceStates=["RUNNING"])
+    instances = emr.list_instances(
+        ClusterId=cluster_id, InstanceStates=["RUNNING"],
+    )
     for instance in instances["Instances"]:
-        grouped_instances[groups[instance[id_name]]].append(instance)
+        grouped_instances[groups[instance.get(id_name, "")]].append(instance)
 
     return grouped_instances
 
@@ -156,18 +182,8 @@ def get_emr_instance_ips(
     '''Returns a dict of IPs per group'''
     return {
         k: [
-            IP(x["PrivateIpAddress"], x.get("PublicIpAddress"))
+            IP(x.get("PrivateIpAddress"), x.get("PublicIpAddress"))
             for x in v
         ]
         for k, v in get_emr_instances(emr, cluster_id).items()
     }
-
-
-def get_instance_key_name(
-    emr: "EMRClient",
-    cluster_id: str,
-) -> str:
-    cluster = emr.describe_cluster(ClusterId=cluster_id)
-
-    key_name = cluster["Cluster"]["Ec2InstanceAttributes"]["Ec2KeyName"]
-    return key_name

{aws_ssh_utils-0.1.1 → aws_ssh_utils-0.2.1}/aws_ssh_utils/interactive_ssh.py

@@ -20,8 +20,10 @@
 import os
 import socket
 import sys
+from importlib.util import find_spec
 
 import paramiko
+from loguru import logger
 
 # https://invisible-island.net/xterm/ctlseqs/ctlseqs.html#h2-Bracketed-Paste-Mode
 START_PASTE = "\x1B\x5B\x32\x30\x30\x7E"  # ESC[200~
@@ -43,24 +45,15 @@ ALL_CODECS = [
     'kz1048', 'mac_cyrillic', 'mac_greek', 'mac_iceland', 'mac_latin2', 'mac_roman',
     'mac_turkish', 'ptcp154', 'shift_jis', 'shift_jis_2004', 'shift_jisx0213',
     'utf_32', 'utf_32_be', 'utf_32_le', 'utf_16', 'utf_16_be', 'utf_16_le', 'utf_7',
-    'utf_8', 'utf_8_sig'
+    'utf_8', 'utf_8_sig',
 ]
 
 # windows does not have termios...
-try:
-    import termios
-    import tty
+has_termios = find_spec("termios") is not None and find_spec("tty") is not None
 
-    has_termios = True
-except ImportError:
-    has_termios = False
 
-if os.getenv('TMUX'):
-    _TITLE_START = '\x1bk'
-    _TITLE_END = '\x1b\\'
-else:
-    _TITLE_START = '\x1b]0;'
-    _TITLE_END = '\x07'
+_TITLE_START = '\x1bk' if os.getenv('TMUX') else '\x1b]0;'
+_TITLE_END = '\x1b\\' if os.getenv('TMUX') else '\x07'
 
 
 def is_int(val: str) -> bool:
@@ -105,19 +98,21 @@ def decode(chars: bytes) -> str:
         return chars.decode('utf-8')
     except Exception:
         # Attempt to decode each individual character
-        def decode_char(byte: bytes) -> str:
+        def decode_char(byte: int) -> str:
             for codec in ALL_CODECS:
                 try:
-                    c = byte.decode(codec)
+                    c = byte.to_bytes(1, 'big').decode(codec)
                     return c
                 except Exception:
-                    pass
+                    logger.debug(
+                        f'Failed to decode character {byte} as {codec}',
+                    )
             else:
                 # Failed to decode character, return replacement character.
                 # https://www.fileformat.info/info/unicode/char/fffd/index.htm
                 return '\uFFFD'
 
-        return ''.join(decode_char(b for b in chars))
+        return ''.join(decode_char(b) for b in chars)
 
 
 def posix_readkey() -> str:
@@ -133,7 +128,7 @@ def posix_readkey() -> str:
             Reads one character and handles encoding errors
             `htop` scrolling for example uses cp037
         '''
-        return decode(sys.stdin.buffer.raw.read(1))
+        return decode(sys.stdin.buffer.raw.read(1))  # pyright: ignore[reportAttributeAccessIssue, reportUnknownArgumentType]
 
     c1 = read()
 
@@ -183,20 +178,25 @@ def windows_readkey() -> str:
     return "\x00" + ch2
 
 
-def posix_shell(chan: paramiko.Channel, allow_title_changes: bool = True):  # noqa: C901
+def posix_shell(chan: paramiko.Channel, allow_title_changes: bool = True):
+    if not has_termios:
+        raise RuntimeError("Termios is not available on this system")
+
     import select
+    import termios
+    import tty
 
-    oldtty = termios.tcgetattr(sys.stdin)
+    oldtty = termios.tcgetattr(sys.stdin)  # pyright: ignore[reportAttributeAccessIssue]
 
     # input_history = []
     # output_history = []
 
     try:
-        tty.setraw(sys.stdin.fileno())
-        tty.setcbreak(sys.stdin.fileno())
+        tty.setraw(sys.stdin.fileno())  # pyright: ignore[reportAttributeAccessIssue]
+        tty.setcbreak(sys.stdin.fileno())  # pyright: ignore[reportAttributeAccessIssue]
         chan.settimeout(0.0)
         while True:
-            r, w, e = select.select([chan, sys.stdin], [], [])
+            r, _, _ = select.select([chan, sys.stdin], [], [])
             if chan in r:
                 try:
                     data = decode(chan.recv(1024))
@@ -210,7 +210,7 @@ def posix_shell(chan: paramiko.Channel, allow_title_changes: bool = True):  # no
                     # output_history.append(data)
                     sys.stdout.write(data)
                     sys.stdout.flush()
-                except socket.timeout:
+                except TimeoutError:
                     pass
             if sys.stdin in r:
                 key = posix_readkey()
@@ -223,7 +223,7 @@ def posix_shell(chan: paramiko.Channel, allow_title_changes: bool = True):  # no
                     key = posix_readkey()
                     # Until we reach the end of the pasted text
                     while key != END_PASTE:
-                        chan.send(key)
+                        chan.send(key.encode())
                         # input_history.append(key)
                         key = posix_readkey()
                     # We've exhausted the paste event, wait for next event
@@ -231,11 +231,11 @@ def posix_shell(chan: paramiko.Channel, allow_title_changes: bool = True):  # no
 
                 if len(key) == 0:
                     break
-                chan.send(key)
+                chan.send(key.encode())
                 # input_history.append(key)
 
     finally:
-        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
+        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)  # pyright: ignore[reportAttributeAccessIssue]
 
     # Useful in debugging how control characters were send
     # from pprint import pprint
@@ -248,15 +248,17 @@ def windows_shell(chan: paramiko.Channel, allow_title_changes: bool = True):
     import threading
 
     sys.stdout.write(
-        "Line-buffered terminal emulation. Press F6 or ^Z to send EOF.\r\n\r\n"
+        "Line-buffered terminal emulation. Press F6 or ^Z to send EOF.\r\n\r\n",
     )
 
-    def writeall(sock):
+    def writeall(sock: socket.socket):
         while True:
             data = sock.recv(256).decode()
             if not data:
                 # Need user to input any character so we sys.stdin.read(1) completes and unblocks
-                sys.stdout.write("\r\n Connection closed. Press Enter to continue...\r\n")
+                sys.stdout.write(
+                    "\r\n Connection closed. Press Enter to continue...\r\n",
+                )
                 sys.stdout.flush()
                 break
 
@@ -274,7 +276,7 @@ def windows_shell(chan: paramiko.Channel, allow_title_changes: bool = True):
             d = windows_readkey()
             if not d or chan.closed:
                 break
-            chan.send(d)
+            chan.send(d.encode())
     except EOFError:
         # user hit ^Z or F6
         pass

{aws_ssh_utils-0.1.1 → aws_ssh_utils-0.2.1}/aws_ssh_utils/ssh.py

@@ -1,23 +1,28 @@
+import datetime as dt
 import logging
 import os
+import subprocess
 import sys
 import textwrap
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from dataclasses import dataclass, field
+from hashlib import sha256
+from types import TracebackType
+from typing import TYPE_CHECKING, Any
 
 import boto3
 import click
 import click_spinner
 import paramiko
+import paramiko.client
 import questionary
 from botocore.exceptions import ClientError
 from environs import Env
 from loguru import logger
+from typing_extensions import override
 
 from .emr_utils import (
     IP,
     get_emr_instance_ips,
-    get_instance_key_name,
     prompt_for_emr_cluster,
     prompt_for_emr_instance_group,
 )
@@ -28,19 +33,24 @@ if TYPE_CHECKING:
     from mypy_boto3_emr import EMRClient
 
 Env().read_env()  # Load .env file
+# CSV of issuer,client
+OPKSSH_PROVIDER_TAG = 'opkssh_provider'
+# EMR doesn't support comma's in tags.
+OPKSSH_ISSUER_TAG = 'opkssh_issuer'
+OPKSSH_CLIENT_TAG = 'opkssh_client'
 
 
 class ShellError(Exception):
     def __init__(self, message: str, exit_code: int = 1):
         super().__init__(message, exit_code)
-        self.message = message
-        self.exit_code = exit_code
+        self.message: str = message
+        self.exit_code: int = exit_code
 
 
 def set_terminal_title(title: str = ''):
     if os.name == 'nt':
         # Windows - CMD
-        os.system(f'title "{title}"')
+        os.system(f'title "{title}"')  # noqa: S605
 
         # Windows - Powershell - But it seems that the scripts are always ran inside CMD anyway.
         # os.system(f'$host.UI.RawUI.WindowTitle = {title}')
@@ -60,7 +70,7 @@ class SelectedEMRInstance:
     cluster_id: str
     cluster_name: str
     group_name: str
-    group_idx: str
+    group_idx: int
     ip: IP
 
 
@@ -72,9 +82,9 @@ def cli(
     long_log: bool = False,
     verbose: bool = False,
     quiet: bool = False,
-    **kwargs,
+    **kwargs: Any,
 ):
-    format = (
+    log_format = (
         "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
         "<level>{level: <8}</level> | "
         "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
@@ -82,7 +92,7 @@ def cli(
     level = "DEBUG" if verbose else 100 if quiet else "INFO"
 
     logger.remove()
-    logger.add(sys.stdout, level=level, format=format, colorize=True)
+    logger.add(sys.stdout, level=level, format=log_format, colorize=True)
 
 
 @cli.command('ec2')
@@ -92,12 +102,12 @@ def cli(
 @click.option('--private/--public', default=True, help="Connect to the instance's private or public IP")
 @click.option('-k/', '--key-file', default=None, help="Which key file to use to connect")
 def ec2_ssh(
-    profile: str = None,
-    region: str = None,
-    user: str = None,
+    profile: str | None = None,
+    region: str | None = None,
+    user: str | None = None,
     private: bool = True,
-    key_file: str = None,
-    **kwargs,
+    key_file: str | None = None,
+    **kwargs: Any,
 ):
     '''
         Asks user which EC2 instance they want to connect to,
@@ -107,16 +117,23 @@ def ec2_ssh(
     try:
         b3s = boto3.Session(profile_name=profile, region_name=region)
 
-        ip, user, key_file, instance_name = get_ec2_ssh_options(b3s=b3s, user=user, use_private_ip=private, key_file=key_file)
+        ip, user, key_file, instance_name = get_ec2_ssh_options(
+            b3s=b3s, user=user, use_private_ip=private, key_file=key_file,
+        )
         terminal_title = f'{user}@{instance_name}'
 
-        SSHShell(hostname=ip, username=user, key_filename=key_file, terminal_title=terminal_title).connect()
+        SSHShell(
+            hostname=ip, username=user, key_filename=key_file,
+            terminal_title=terminal_title,
+        ).connect()
     except ShellError as e:
         logger.error(e.message)
         exit(e.exit_code)
     except ClientError as e:
-        if e.response.get('Error') and e.response['Error'].get('Code') == 'ExpiredTokenException':
-            logger.log(logging.CRITICAL, 'Your AWS Token has expired. Please update and try again.')
+        if e.response.get('Error') and e.response.get('Error', {}).get('Code') == 'ExpiredTokenException':
+            logger.log(
+                logging.CRITICAL, 'Your AWS Token has expired. Please update and try again.',
+            )
             exit(1)
 
 
@@ -127,12 +144,12 @@ def ec2_ssh(
 @click.option('--private/--public', default=True, help="Connect to the instance's private or public IP")
 @click.option('-k/', '--key-file', default=None, help="Which key file to use to connect")
 def emr_ssh(
-    profile: str = None,
-    region: str = None,
-    user: str = None,
+    profile: str | None = None,
+    region: str | None = None,
+    user: str | None = None,
     private: bool = True,
-    key_file: str = None,
-    **kwargs
+    key_file: str | None = None,
+    **kwargs: Any,
 ):
     '''
         Asks user which Cluster and EC2 instance they want to connect to,
@@ -148,23 +165,31 @@ def emr_ssh(
             user = 'hadoop'
 
         if key_file is None:
-            key_file = get_emr_ssh_key_file_for_cluster(emr, emr_instance.cluster_id)
+            key_file = get_emr_ssh_key_file_for_cluster(
+                emr, emr_instance.cluster_id,
+            )
 
         if private:
             ip = emr_instance.ip.private
+            if ip is None:
+                raise ShellError(
+                    'The selected instance does not have a private IP',
+                )
         else:
             ip = emr_instance.ip.public
             if ip is None:
-                raise ShellError('The selected instance does not have a public IP')
+                raise ShellError(
+                    'The selected instance does not have a public IP',
+                )
 
         group_name = emr_instance.group_name.split(' ')[0]
-        postfix = f'[{emr_instance.group_idx}]' if emr_instance.group_idx is not None else ''
+        postfix = f'[{emr_instance.group_idx}]' or ''
         terminal_title = f'{emr_instance.cluster_name} - {group_name}{postfix}'
 
         with SSHShell(hostname=ip, username=user, key_filename=key_file, terminal_title=terminal_title) as shell:
             if group_name.lower().startswith('core'):
-                shell.send('sudo su\r\n')
-                shell.send('cd /var/log/hadoop-yarn/containers\r\n')
+                shell.send(b'sudo su\r\n')
+                shell.send(b'cd /var/log/hadoop-yarn/containers\r\n')
 
         set_terminal_title()
 
@@ -172,8 +197,10 @@ def emr_ssh(
         logger.error(e.message)
         exit(e.exit_code)
     except ClientError as e:
-        if e.response.get('Error') and e.response['Error'].get('Code') == 'ExpiredTokenException':
-            logger.log(logging.CRITICAL, 'Your AWS Token has expired. Please update and try again.')
+        if e.response.get('Error') and e.response.get('Error', {}).get('Code') == 'ExpiredTokenException':
+            logger.log(
+                logging.CRITICAL, 'Your AWS Token has expired. Please update and try again.',
+            )
             exit(1)
 
 
@@ -184,12 +211,12 @@ def emr_ssh(
 @click.option('--private/--public', default=True, help="Connect to the instance's private or public IP")
 @click.option('-k/', '--key-file', default=None, help="Which key file to use to connect")
 def emr_ssh_all(
-    profile: str = None,
-    region: str = None,
-    user: str = None,
+    profile: str | None = None,
+    region: str | None = None,
+    user: str | None = None,
     private: bool = True,
-    key_file: str = None,
-    **kwargs,
+    key_file: str | None = None,
+    **kwargs: Any,
 ):
     '''
         Asks user which Cluster and EC2 instance they want to connect to,
@@ -213,18 +240,23 @@ def emr_ssh_all(
 
         window_order = ['Master', 'Primary', 'Core', 'Task']
 
-        def get_ip(ip: IP) -> str:
+        def get_ip(ip: IP) -> str | None:
             if private:
                 return ip.private
             else:
                 if ip.public is None:
-                    raise ShellError(f'The instance with private IP {ip.private} does not have a public IP')
+                    raise ShellError(
+                        f'The instance with private IP {ip.private} does not have a public IP',
+                    )
                 return ip.public
 
         window_cmds = [
             # Add "|| $SHELL -i" so that if the ssh session fails, we still have a window so we can look at the error.
             # without this, tmux will close the window.
-            (f'{group_name} - {instance_num}', f'ssh -i {key_file} {user}@{get_ip(instance_ip)} || $SHELL -i')
+            (
+                f'{group_name} - {instance_num}',
+                f'ssh -i {key_file} {user}@{get_ip(instance_ip)} || $SHELL -i',
+            )
             for group_name, instance_ips in sorted(grouped_instances.items(), key=lambda t: window_order.index(t[0]))
             for instance_num, instance_ip in enumerate(instance_ips, start=1)
         ]
@@ -232,10 +264,14 @@ def emr_ssh_all(
         session_name = cluster_name
         first_window_name, first_cmd = window_cmds[0]
         # Create a new tmux session with the EMR cluster name as the session name and open the first ssh connection
-        tmux(f'new-session -d -s "{session_name}" -n "{first_window_name}" "{first_cmd}"')
+        tmux(
+            f'new-session -d -s "{session_name}" -n "{first_window_name}" "{first_cmd}"',
+        )
         # Create a new tmux window and open the ssh connections
         for window_cmd in window_cmds[1:]:
-            tmux(f'new-window -t "{session_name}:" -n "{window_cmd[0]}" "{window_cmd[1]}"')
+            tmux(
+                f'new-window -t "{session_name}:" -n "{window_cmd[0]}" "{window_cmd[1]}"',
+            )
         # Now switch to the session's first window
         tmux(f'switch-client -t "{session_name}:{first_window_name}"')
 
@@ -243,32 +279,43 @@ def emr_ssh_all(
         logger.error(e.message)
         exit(e.exit_code)
     except ClientError as e:
-        if e.response.get('Error') and e.response['Error'].get('Code') == 'ExpiredTokenException':
-            logger.log(logging.CRITICAL, 'Your AWS Token has expired. Please update and try again.')
+        if e.response.get('Error') and e.response.get('Error', {}).get('Code') == 'ExpiredTokenException':
+            logger.log(
+                logging.CRITICAL, 'Your AWS Token has expired. Please update and try again.',
+            )
             exit(1)
 
 
-@cli.command('test')
-def test(**kwargs):
-    tmux('new-session -d -s "test" -n "window 1" "echo hello && $SHELL -i"')
-    tmux('new-window -n "window 2" -t "test:" "echo hello && $SHELL -i"')
-    tmux('switch-client -t "test:window 1"')
-
-
 @dataclass
 class SSHShell:
     hostname: str
     username: str
-    key_filename: str
-    terminal_title: str = None
+    key_filename: str | None = None
+    terminal_title: str | None = None
+
+    private_key: paramiko.PKey = field(init=False)
+    _ssh_client: paramiko.SSHClient = field(init=False)
+    _channel: paramiko.Channel = field(init=False)
+
+    def __post_init__(self):
+        if self.key_filename is not None:
+            if not os.path.isfile(self.key_filename):
+                raise ValueError(f'File {self.key_filename} does not exist')
+
+            self.private_key = paramiko.PKey.from_path(self.key_filename)
+            public_key_path = f'{self.key_filename}.pub'
+            if os.path.isfile(public_key_path):
+                self.private_key.load_certificate(public_key_path)
 
     def connect(self):
         self._open()
         self._launch()
         self._close()
 
-    def _open(self):
-        logger.info(f'Opening SSH: ssh -i {self.key_filename} {self.username}@{self.hostname}')
+    def _open(self) -> paramiko.Channel:
+        logger.info(
+            f'Opening SSH: ssh -i {self.key_filename} {self.username}@{self.hostname}',
+        )
 
         terminal_size = os.get_terminal_size()
 
@@ -278,7 +325,11 @@ class SSHShell:
         ssh_client.load_host_keys(host_key_path)
         ssh_client.set_missing_host_key_policy(ConfirmAddPolicy())
         try:
-            ssh_client.connect(hostname=self.hostname, username=self.username, key_filename=self.key_filename)
+            ssh_client.connect(
+                hostname=self.hostname,
+                username=self.username,
+                pkey=self.private_key,
+            )
         except paramiko.BadHostKeyException as e:
             error_message = f'''
                 @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@ -301,8 +352,11 @@ class SSHShell:
             '''
             raise ShellError(textwrap.dedent(error_message), 255) from None
 
-        channel = ssh_client.get_transport().open_session()
-        channel.get_pty(term=os.getenv('TERM', 'xterm-256color'), width=terminal_size.columns, height=terminal_size.lines)
+        channel = ssh_client.get_transport().open_session()  # pyright: ignore[reportOptionalMemberAccess]
+        channel.get_pty(
+            term=os.getenv('TERM', 'xterm-256color'),
+            width=terminal_size.columns, height=terminal_size.lines,
+        )
         channel.invoke_shell()
 
         self._ssh_client = ssh_client
@@ -313,7 +367,9 @@ class SSHShell:
         return channel
 
     def _launch(self):
-        interactive_shell(self._channel, allow_title_changes=not self.terminal_title)
+        interactive_shell(
+            self._channel, allow_title_changes=not self.terminal_title,
+        )
 
     def _close(self):
         self._ssh_client.close()
@@ -324,16 +380,16 @@ class SSHShell:
     def __enter__(self):
         return self._open()
 
-    def __exit__(self, exception_type, exception_value, traceback):
+    def __exit__(self, exception_type: type[BaseException] | None, exception_value: BaseException | None, traceback: TracebackType | None):
         self._launch()
         self._close()
 
 
 def tmux(command: str):
-    os.system(f'tmux {command}')
+    os.system(f'tmux {command}')  # noqa: S605
 
 
-def try_to_find_ssh_key_file(key_name: str) -> str:
+def try_to_find_ssh_key_file(key_name: str) -> str | None:
     '''Recursively iterate over the `~/.ssh/` folder to find the matching key'''
     if key_name:
         for dirpath, _, filenames in os.walk(os.path.expanduser('~/.ssh/')):
@@ -345,7 +401,7 @@ def try_to_find_ssh_key_file(key_name: str) -> str:
     return None
 
 
-def get_ec2_image_name(instance) -> str:
+def get_ec2_image_name(instance: "Instance") -> str | None:
     try:
         return instance.image.name
     except AttributeError:
@@ -354,21 +410,34 @@ def get_ec2_image_name(instance) -> str:
 
 def get_ec2_ssh_options(
     b3s: boto3.Session,
-    user: str = None,
+    user: str | None = None,
     use_private_ip: bool = True,
-    key_file: str = None
-) -> tuple[str, str, str, str]:
+    key_file: str | None = None,
+) -> tuple[str, str, str | None, str]:
     instance = prompt_for_ec2_instance(b3s)
+    instance_tags = {
+        tag.get('Key', '').lower(): tag.get('Value', '')
+        for tag in instance.tags
+    }
+
     if key_file is None:
-        key_file = try_to_find_ssh_key_file(instance.key_name)
+        # Support https://github.com/openpubkey/opkssh via tags
+        if OPKSSH_PROVIDER_TAG in instance_tags:
+            key_file = get_opkssh_key_file(instance_tags[OPKSSH_PROVIDER_TAG])
+        elif OPKSSH_ISSUER_TAG in instance_tags and OPKSSH_CLIENT_TAG in instance_tags:
+            key_file = get_opkssh_key_file(
+                ','.join([
+                    instance_tags[OPKSSH_ISSUER_TAG],
+                    instance_tags[OPKSSH_CLIENT_TAG],
+                ]),
+            )
+        else:
+            key_file = try_to_find_ssh_key_file(instance.key_name)
 
     if user is None:
         logger.info('No user specified, attempting to detect required user...')
         image_name = get_ec2_image_name(instance)
-        if 'ubuntu' in image_name:
-            user = 'ubuntu'
-        else:
-            user = 'ec2-user'
+        user = 'ubuntu' if image_name and 'ubuntu' in image_name.lower() else 'ec2-user'
 
     if use_private_ip:
         ip = instance.private_ip_address
@@ -386,10 +455,10 @@ def get_emr_ssh_options(
     cluster_id, cluster_name = prompt_for_emr_cluster(emr)
     instance_ips, group_name = prompt_for_emr_instance_group(emr, cluster_id)
 
-    instance_ips = sorted(instance_ips, key=lambda ips: ips.private)
+    instance_ips = sorted(instance_ips, key=lambda ips: ips.private or "")
 
     instance_options = [
-        f'{ip.private} ({ip.public})' if ip.public else ip.private
+        f'{ip.private} ({ip.public})' if ip.public else f'{ip.private}'
         for ip in instance_ips
     ]
     instance_ip = questionary.select(
@@ -404,15 +473,38 @@ def get_emr_ssh_options(
 def get_emr_ssh_key_file_for_cluster(
     emr: "EMRClient",
     cluster_id: str,
-) -> str:
+) -> str | None:
     with click_spinner.spinner():
-        key_name = get_instance_key_name(emr, cluster_id)
-        key_file = try_to_find_ssh_key_file(key_name)
+        cluster = emr.describe_cluster(ClusterId=cluster_id)
+        cluster_tags = {
+            tag.get('Key', '').lower(): tag.get('Value', '')
+            for tag in cluster['Cluster'].get('Tags', [])
+        }
 
-    if key_file is None:
-        should_continue = questionary.confirm(f'Could not find the ssh key {key_name}, would you like to continue?').unsafe_ask()
-        if not should_continue:
-            exit(1)
+        key_file = key_name = None
+        # Support https://github.com/openpubkey/opkssh via tags
+        if OPKSSH_PROVIDER_TAG in cluster_tags:
+            key_file = get_opkssh_key_file(cluster_tags[OPKSSH_PROVIDER_TAG])
+        elif OPKSSH_ISSUER_TAG in cluster_tags and OPKSSH_CLIENT_TAG in cluster_tags:
+            key_file = get_opkssh_key_file(
+                ','.join([
+                    cluster_tags[OPKSSH_ISSUER_TAG],
+                    cluster_tags[OPKSSH_CLIENT_TAG],
+                ]),
+            )
+
+        if key_file is None:
+            key_name = cluster["Cluster"].get(
+                "Ec2InstanceAttributes", {},
+            ).get("Ec2KeyName", "")
+            key_file = try_to_find_ssh_key_file(key_name)
+
+        if key_file is None:
+            should_continue = questionary.confirm(
+                f'Could not find the ssh key {key_name}, would you like to continue?',
+            ).unsafe_ask()
+            if not should_continue:
+                exit(1)
 
     return key_file
 
@@ -466,28 +558,61 @@ def prompt_for_ec2_instance(
 def get_ec2_name(ec2_instance: "Instance") -> str:
     '''Takes in the boto3 EC2 resource instance object'''
     for tag in ec2_instance.tags:
-        if tag['Key'] == 'Name':
-            return tag['Value']
+        if tag.get('Key') == 'Name':
+            return tag.get('Value', '')
 
     return ec2_instance.instance_id
 
 
+def get_opkssh_key_file(provider: str) -> str:
+    opkssh_key_file = os.path.join(
+        os.path.expanduser(
+            '~/.ssh/',
+        ), f'opkssh_{sha256(provider.encode()).hexdigest()}',
+    )
+
+    if os.path.exists(opkssh_key_file) and os.path.getmtime(opkssh_key_file) > dt.datetime.now().timestamp() - 86400:
+        logger.info('Found existing opkssh key')
+    else:
+        logger.info('Detected opkssh, logging in')
+        cmd = f"opkssh login --provider '{provider}' -i '{opkssh_key_file}'"
+        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)  # noqa: S602
+        process.wait()
+        if process.returncode != 0:
+            logger.error('Failed to authenticate using opkssh')
+            sys.exit(process.returncode)
+
+    return opkssh_key_file
+
+
 class ConfirmAddPolicy(paramiko.client.MissingHostKeyPolicy):
     """
     Policy for automatically adding the hostname and new host key to the
     local `.HostKeys` object, and saving it.  This is used by `.SSHClient`.
     """
 
-    def missing_host_key(self, client, hostname, key):
-        logger.warning(f"Unknown {key.get_name()} host key for {hostname}: {key.fingerprint}")
-        should_add = questionary.confirm("Continue and add host key?").unsafe_ask()
+    @override
+    def missing_host_key(self, client: paramiko.SSHClient, hostname: str, key: paramiko.PKey):
+        logger.warning(
+            f"Unknown {key.get_name()} host key for {hostname}: {key.fingerprint}",
+        )
+        should_add = questionary.confirm(
+            "Continue and add host key?",
+        ).unsafe_ask()
         if should_add:
-            client._host_keys.add(hostname, key.get_name(), key)
-            if client._host_keys_filename is not None:
-                client.save_host_keys(client._host_keys_filename)
+            client.get_host_keys().add(hostname, key.get_name(), key)
+            host_key_filename: str | None = (
+                client._host_keys_filename  # pyright: ignore[reportAttributeAccessIssue]
+            )
+            if host_key_filename is not None:
+                client.save_host_keys(host_key_filename)
                 logger.info('Added host key')
             else:
-                logger.warning('Failed to add host key. No host key file defined!')
+                logger.warning(
+                    'Failed to add host key. No host key file defined!',
+                )
 
         else:
-            raise paramiko.SSHException(f"Server {hostname!r} not found in known_hosts")
+            raise paramiko.SSHException(
+                f"Server {hostname!r} not found in known_hosts",
+            )

aws_ssh_utils-0.2.1/pyproject.toml

@@ -0,0 +1,161 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "aws-ssh-utils"
+description = "Easy AWS SSHing"
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+authors = [
+    { name = "Michiel Vanderlee", email = "jmt.vanderlee@gmail.com" },
+]
+classifiers = [
+    "Intended Audience :: Information Technology",
+    "Intended Audience :: System Administrators",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: Software Development :: Libraries",
+    "Topic :: Software Development",
+    "Typing :: Typed",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+dependencies = [
+  "boto3",
+  "click",
+  "click_spinner",
+  "coloredlogs",
+  "environs",
+  "loguru",
+  "paramiko",
+  "questionary",
+  "typing_extensions",
+]
+dynamic = ["version"]
+
+[project.urls]
+Homepage = "https://github.com/mvanderlee/aws-ssh-utils"
+
+[project.scripts]
+aws_ssh = "aws_ssh_utils.ssh:cli"
+
+[project.optional-dependencies]
+
+dev = [
+  "boto3-stubs[ec2,emr]",
+]
+test = [
+]
+publish = [
+  "hatch >= 1.7.0",
+]
+
+[tool.hatch.build]
+include = [
+  "aws_ssh_utils",
+]
+
+[tool.hatch.version]
+path = "aws_ssh_utils/__init__.py"
+
+
+[tool.coverage.run]
+parallel = true
+source = [
+  "aws_ssh_utils"
+]
+
+[tool.autopep8]
+max_line_length = 130
+
+[tool.ruff]
+exclude = [
+    "*.ipynb",
+    "**/_version.py",
+]
+lint.select = [
+    "E",  # pycodestyle errors
+    "W",  # pycodestyle warnings
+    "F",  # pyflakes
+    "I",  # isort
+    "C",  # flake8-comprehensions
+    "B",  # flake8-bugbear
+    "A",  # flake8-builtins
+    "N",  # pep8-naming
+    "UP",  # pyupgrade
+    "S",  # flake8-bandit
+    "COM",  # flake8-commas
+    "ISC",  # flake8-implicit-str-concat
+    "INP",  # flake8-no-pep420
+    "PIE",  # flake8-pie
+    "SIM",  # flake8-simplify
+    "RUF",  # Ruff-specific rules
+    "T20",  # flake8-print
+    "PT",  # flake8-pytest-style
+]
+lint.ignore = [
+    "B008",  # do not perform function calls in argument defaults
+    "B028",  # No explicit stacklevel argument found.
+    "C408",  # unnecessary-collection-call - Rewrite as a literal
+    "C901",  # too complex
+    "PT003",  # pytest-extraneous-scope-function
+
+    # TODO: remove
+    "A002",  # builtin-argument-shadowing
+    "PT011",  # pytest-raises-too-broad
+    "N818",  # error-suffix-on-exception-name
+
+    # Unsure, need more data
+    "SIM117",  # multiple-with-statements - In general good, but sometimes becomes harder to read, i.e.: aiohttp?
+]
+# The formatter wraps lines at a length of 88.
+line-length = 130  # Enforce 130. But ~100 is recommended.
+
+[tool.ruff.lint.pycodestyle]
+max-line-length = 150  # E501 reports lines that exceed the length of 150.
+
+[tool.ruff.format]
+quote-style = "preserve"  # Don't change quotes - https://docs.astral.sh/ruff/settings/#format_quote-style
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
+
+[tool.ruff.lint.flake8-bugbear]
+# Allow default arguments like, e.g., `data: List[str] = fastapi.Query(None)`.
+extend-immutable-calls = []
+
+[tool.pytest.ini_options]
+# https://pytest-asyncio.readthedocs.io/en/latest/reference/configuration.html
+asyncio_mode = "auto"
+asyncio_default_test_loop_scope = "session"
+asyncio_default_fixture_loop_scope = "session"
+
+[tool.pyright]
+venvPath = "."
+venv = ".venv"
+
+include = [
+    "aws_ssh_utils/",
+]
+
+pythonVersion = "3.11"
+
+# basedpyright exclusions
+reportAny = "none"
+reportUnreachable = "none"
+reportExplicitAny = "none"
+reportUnusedCallResult = "none"
+reportUnknownMemberType = "none"
+reportUnusedParameter = "none"
+reportUntypedFunctionDecorator = "none"
+reportMissingTypeStubs = "none"
+reportUnknownVariableType = "none"
+reportImplicitStringConcatenation = "none"

aws_ssh_utils-0.1.1/aws_ssh_utils/__init__.py

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

aws_ssh_utils-0.1.1/pyproject.toml

@@ -1,102 +0,0 @@
-[build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
-
-[project]
-name = "aws-ssh-utils"
-description = "Easy AWS SSHing"
-readme = "README.md"
-requires-python = ">=3.11"
-license = "MIT"
-authors = [
-    { name = "Michiel Vanderlee", email = "jmt.vanderlee@gmail.com" },
-]
-classifiers = [
-    "Intended Audience :: Information Technology",
-    "Intended Audience :: System Administrators",
-    "Operating System :: OS Independent",
-    "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python",
-    "Topic :: Software Development :: Libraries :: Python Modules",
-    "Topic :: Software Development :: Libraries",
-    "Topic :: Software Development",
-    "Typing :: Typed",
-    "Intended Audience :: Developers",
-    "License :: OSI Approved :: MIT License",
-    "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.11",
-    "Programming Language :: Python :: 3.12",
-    "Programming Language :: Python :: 3.13",
-]
-dependencies = [
-  "boto3",
-  "click",
-  "click_spinner",
-  "coloredlogs",
-  "environs",
-  "loguru",
-  "paramiko",
-  "questionary",
-]
-dynamic = ["version"]
-
-[project.urls]
-Homepage = "https://github.com/mvanderlee/aws_ssh_utils"
-
-[project.scripts]
-aws_ssh = "aws_ssh_utils.ssh:cli"
-
-[project.optional-dependencies]
-
-dev = [
-  "boto3-stubs[ec2,emr]",
-]
-test = [
-]
-publish = [
-  "hatch >= 1.7.0",
-]
-
-[tool.hatch.build]
-include = [
-  "aws_ssh_utils",
-]
-
-[tool.hatch.version]
-path = "aws_ssh_utils/__init__.py"
-
-[tool.isort]
-profile = "hug"
-line_length = 100
-
-[tool.coverage.run]
-parallel = true
-source = [
-  "aws_ssh_utils"
-]
-
-[tool.ruff]
-select = [
-  "E",  # pycodestyle errors
-  "W",  # pycodestyle warnings
-  "F",  # pyflakes
-  # "I",  # isort
-  "C",  # flake8-comprehensions
-  "B",  # flake8-bugbear
-]
-ignore = [
-  "B008",  # do not perform function calls in argument defaults
-  "B028",  # No explicit stacklevel argument found.
-  "C901",  # too complex
-]
-line-length = 130  # Enforce 130. But ~100 is recommended.
-
-
-[tool.ruff.per-file-ignores]
-"__init__.py" = ["F401"]
-"migrations/env.py" = ["E402"]
-
-[tool.pytest.ini_options]
-# https://pytest-asyncio.readthedocs.io/en/latest/reference/configuration.html
-asyncio_mode = "auto"
-asyncio_default_fixture_loop_scope = "session"