aws-ssh-utils 0.1.0__py3-none-any.whl

aws_ssh_utils/__init__.py

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

aws_ssh_utils/emr_utils.py

@@ -0,0 +1,173 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, MutableMapping
+
+import click_spinner
+import questionary
+
+if TYPE_CHECKING:
+    from mypy_boto3_emr import EMRClient
+    from mypy_boto3_emr.type_defs import ClusterSummaryTypeDef, InstanceTypeDef
+
+
+@dataclass
+class IP:
+    private: str
+    public: str = None
+
+
+def prompt_for_emr_cluster(
+    emr: "EMRClient",
+    prompt="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,
+) -> tuple[str, str]:
+    """Discovers the available EMR clusters and asks the user which to use.
+
+        Returns the cluster id of the cluster that the user selected.
+    """
+    # Discover available clusters
+    with click_spinner.spinner():
+        clusters = get_emr_clusters(emr, applications=applications)
+
+    # Prompt the user
+    cluster_name = questionary.select(
+        prompt,
+        choices=sorted(clusters),
+    ).unsafe_ask()
+    cluster_id, cluster_name = clusters[cluster_name]
+
+    return cluster_id, cluster_name
+
+
+def prompt_for_emr_instance_group(
+    emr: "EMRClient",
+    cluster_id: str, prompt="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.
+
+        Returns the list of instances in the selected group, as well as the group name
+    """
+    # Discover the available instances grouped by groupname
+    with click_spinner.spinner():
+        grouped_instances = get_emr_instance_ips(emr, cluster_id)
+
+    # Prompt the user
+    group_name = questionary.select(
+        prompt,
+        choices=sorted(grouped_instances),
+    ).unsafe_ask()
+
+    instances = grouped_instances[group_name]
+    return instances, group_name
+
+
+def does_cluster_have_applications(
+    emr: "EMRClient",
+    cluster_id: str,
+    applications: dict[str, str] | list[str] = None,
+) -> bool:
+    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"]}
+
+    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())
+    else:
+        return all(k in cluster_applications for k in applications)
+
+
+def get_emr_clusters(
+    emr: "EMRClient",
+    states: list[str] = None,
+    # Either a list of application names, or a dict[name, version]
+    applications: dict[str, str] | list[str] = None,
+) -> dict[str, tuple[str, str]]:
+    """Discover the available EMR cluster
+
+        Returns a dict where the key is the cluster "Name - ID - CreatedDT" and the value is the ID.
+    """
+    if states is None:
+        states = ["RUNNING", "WAITING"]
+
+    clusters = emr.list_clusters(
+        ClusterStates=states,
+    ).get("Clusters", [])
+
+    if applications:
+        clusters = [c for c in clusters if does_cluster_have_applications(emr, c["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}"
+
+    return {f"{get_display_name(c)}": (c["Id"], c["Name"]) for c in clusters}
+
+
+def get_emr_groups(
+    emr: "EMRClient",
+    cluster_id: str,
+):
+    """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")
+
+    if instance_collection_type == "INSTANCE_FLEET":
+        id_name = "InstanceFleetId"
+        groups = emr.list_instance_fleets(ClusterId=cluster_id)
+        groups = groups["InstanceFleets"]
+
+    elif instance_collection_type == "INSTANCE_GROUP":
+        id_name = "InstanceGroupId"
+        groups = emr.list_instance_groups(ClusterId=cluster_id)
+        groups = groups["InstanceGroups"]
+
+    return groups, id_name
+
+
+def get_emr_instances(
+    emr: "EMRClient",
+    cluster_id: str,
+) -> dict[str, list["InstanceTypeDef"]]:
+    """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}
+    grouped_instances = {g: [] for g in groups.values()}
+
+    instances = emr.list_instances(ClusterId=cluster_id, InstanceStates=["RUNNING"])
+    for instance in instances["Instances"]:
+        grouped_instances[groups[instance[id_name]]].append(instance)
+
+    return grouped_instances
+
+
+def get_emr_instance_ips(
+    emr: "EMRClient",
+    cluster_id: str,
+) -> dict[str, list[IP]]:
+    '''Returns a dict of IPs per group'''
+    return {
+        k: [
+            IP(x["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/interactive_ssh.py

@@ -0,0 +1,280 @@
+# Source: https://github.com/paramiko/paramiko/blob/main/demos/interactive.py
+
+# Copyright (C) 2003-2007  Robey Pointer <robeypointer@gmail.com>
+#
+# This file is part of paramiko.
+#
+# Paramiko is free software; you can redistribute it and/or modify it under the
+# terms of the GNU Lesser General Public License as published by the Free
+# Software Foundation; either version 2.1 of the License, or (at your option)
+# any later version.
+#
+# Paramiko is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+# A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
+# details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Paramiko; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA.
+import os
+import socket
+import sys
+
+import paramiko
+
+# https://invisible-island.net/xterm/ctlseqs/ctlseqs.html#h2-Bracketed-Paste-Mode
+START_PASTE = "\x1B\x5B\x32\x30\x30\x7E"  # ESC[200~
+END_PASTE = "\x1B\x5B\x32\x30\x31\x7E"  # ESC[201~
+
+
+ALL_CODECS = [
+    'ascii', 'big5', 'big5hkscs', 'cp037', 'cp273', 'cp424', 'cp437',
+    'cp500', 'cp720', 'cp737', 'cp775', 'cp850', 'cp852', 'cp855', 'cp856', 'cp857',
+    'cp858', 'cp860', 'cp861', 'cp862', 'cp863', 'cp864', 'cp865', 'cp866', 'cp869',
+    'cp874', 'cp875', 'cp932', 'cp949', 'cp950', 'cp1006', 'cp1026', 'cp1125',
+    'cp1140', 'cp1250', 'cp1251', 'cp1252', 'cp1253', 'cp1254', 'cp1255', 'cp1256',
+    'cp1257', 'cp1258', 'euc_jp', 'euc_jis_2004', 'euc_jisx0213', 'euc_kr',
+    'gb2312', 'gbk', 'gb18030', 'hz', 'iso2022_jp', 'iso2022_jp_1', 'iso2022_jp_2',
+    'iso2022_jp_2004', 'iso2022_jp_3', 'iso2022_jp_ext', 'iso2022_kr', 'latin_1',
+    'iso8859_2', 'iso8859_3', 'iso8859_4', 'iso8859_5', 'iso8859_6', 'iso8859_7',
+    'iso8859_8', 'iso8859_9', 'iso8859_10', 'iso8859_11', 'iso8859_13',
+    'iso8859_14', 'iso8859_15', 'iso8859_16', 'johab', 'koi8_r', 'koi8_t', 'koi8_u',
+    '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'
+]
+
+# windows does not have termios...
+try:
+    import termios
+    import tty
+
+    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'
+
+
+def is_int(val: str) -> bool:
+    try:
+        int(val)
+        return True
+    except Exception:
+        return False
+
+
+def remove_title_change(data: str) -> str:
+    '''
+        Title change is a string starting with '\x1b]0;', and ending with '\x07'.
+        Remove all such substrings
+    '''
+    while True:
+        a = data.find(_TITLE_START)
+        b = data.find(_TITLE_END)
+
+        if a == -1 or b == -1:
+            break
+
+        data = data[:a] + data[b + len(_TITLE_END):]
+
+    return data
+
+
+def interactive_shell(chan: paramiko.Channel, allow_title_changes: bool = True):
+    if has_termios:
+        posix_shell(chan, allow_title_changes=allow_title_changes)
+    else:
+        windows_shell(chan, allow_title_changes=allow_title_changes)
+
+
+def decode(chars: bytes) -> str:
+    '''
+        Decodes the bytes and handles encoding errors
+        `htop` scrolling for example uses cp037
+    '''
+    try:
+        # Attempt to decode the entire input as UTF-8
+        return chars.decode('utf-8')
+    except Exception:
+        # Attempt to decode each individual character
+        def decode_char(byte: bytes) -> str:
+            for codec in ALL_CODECS:
+                try:
+                    c = byte.decode(codec)
+                    return c
+                except Exception:
+                    pass
+            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))
+
+
+def posix_readkey() -> str:
+    """Get a keypress. If an escaped key is pressed, the full sequence is
+    read and returned.
+
+        Copied from readchar:
+        https://github.com/magmax/python-readchar/blob/master/readchar/_posix_read.py#L30
+    """
+
+    def read():
+        '''
+            Reads one character and handles encoding errors
+            `htop` scrolling for example uses cp037
+        '''
+        return decode(sys.stdin.buffer.raw.read(1))
+
+    c1 = read()
+
+    if c1 != "\x1B":  # ESC
+        return c1
+
+    c2 = read()
+    if c2 not in "\x4F\x5B":  # O[
+        return c1 + c2
+
+    c3 = read()
+    if c3 not in "\x31\x32\x33\x35\x36":  # 12356
+        return c1 + c2 + c3
+
+    c4 = read()
+    if c4 not in "\x30\x31\x33\x34\x35\x37\x38\x39":  # 01345789
+        return c1 + c2 + c3 + c4
+
+    c5 = read()
+    key = c1 + c2 + c3 + c4 + c5
+
+    # Bracketed Paste Mode: # https://invisible-island.net/xterm/ctlseqs/ctlseqs.html#h2-Bracketed-Paste-Mode
+    if key == START_PASTE[:-1] or key == END_PASTE[:-1]:
+        c6 = read()
+        return key + c6
+
+    return key
+
+
+def windows_readkey() -> str:
+    """Reads the next keypress. If an escaped key is pressed, the full
+    sequence is read and returned.
+
+        Copied from readchar:
+        https://github.com/magmax/python-readchar/blob/master/readchar/_win_read.py#LL14C1-L30C24
+    """
+
+    ch = sys.stdin.read(1)
+
+    # if it is a normal character:
+    if ch not in "\x00\xe0":
+        return ch
+
+    # if it is a scpeal key, read second half:
+    ch2 = sys.stdin.read(1)
+
+    return "\x00" + ch2
+
+
+def posix_shell(chan: paramiko.Channel, allow_title_changes: bool = True):  # noqa: C901
+    import select
+
+    oldtty = termios.tcgetattr(sys.stdin)
+
+    # input_history = []
+    # output_history = []
+
+    try:
+        tty.setraw(sys.stdin.fileno())
+        tty.setcbreak(sys.stdin.fileno())
+        chan.settimeout(0.0)
+        while True:
+            r, w, e = select.select([chan, sys.stdin], [], [])
+            if chan in r:
+                try:
+                    data = decode(chan.recv(1024))
+                    if len(data) == 0:
+                        sys.stdout.write("\r\n")
+                        break
+
+                    if not allow_title_changes:
+                        data = remove_title_change(data)
+
+                    # output_history.append(data)
+                    sys.stdout.write(data)
+                    sys.stdout.flush()
+                except socket.timeout:
+                    pass
+            if sys.stdin in r:
+                key = posix_readkey()
+                # When pasting something, we need to read the entire pasted blob at once
+                # Otherwise it'll hang until the next key press.
+                # This has to do with how 'select.select' detects changes.
+                # A paste is a single event of many characters, so we must handle them all as one event
+                if key == START_PASTE:
+                    # Start reading the pasted text
+                    key = posix_readkey()
+                    # Until we reach the end of the pasted text
+                    while key != END_PASTE:
+                        chan.send(key)
+                        # input_history.append(key)
+                        key = posix_readkey()
+                    # We've exhausted the paste event, wait for next event
+                    continue
+
+                if len(key) == 0:
+                    break
+                chan.send(key)
+                # input_history.append(key)
+
+    finally:
+        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
+
+    # Useful in debugging how control characters were send
+    # from pprint import pprint
+    # pprint(input_history)
+    # pprint(output_history)
+
+
+# thanks to Mike Looijmans for this code
+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"
+    )
+
+    def writeall(sock):
+        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.flush()
+                break
+
+            if not allow_title_changes:
+                data = remove_title_change(data)
+
+            sys.stdout.write(data)
+            sys.stdout.flush()
+
+    writer = threading.Thread(target=writeall, args=(chan,))
+    writer.start()
+
+    try:
+        while True:
+            d = windows_readkey()
+            if not d or chan.closed:
+                break
+            chan.send(d)
+    except EOFError:
+        # user hit ^Z or F6
+        pass

aws_ssh_utils/ssh.py

@@ -0,0 +1,493 @@
+import logging
+import os
+import sys
+import textwrap
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import boto3
+import click
+import click_spinner
+import paramiko
+import questionary
+from botocore.exceptions import ClientError
+from environs import Env
+from loguru import logger
+
+from .emr_utils import (
+    IP,
+    get_emr_instance_ips,
+    get_instance_key_name,
+    prompt_for_emr_cluster,
+    prompt_for_emr_instance_group,
+)
+from .interactive_ssh import interactive_shell
+
+if TYPE_CHECKING:
+    from mypy_boto3_ec2.service_resource import Instance
+    from mypy_boto3_emr import EMRClient
+
+Env().read_env()  # Load .env file
+
+
+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
+
+
+def set_terminal_title(title: str = ''):
+    if os.name == 'nt':
+        # Windows - CMD
+        os.system(f'title "{title}"')
+
+        # Windows - Powershell - But it seems that the scripts are always ran inside CMD anyway.
+        # os.system(f'$host.UI.RawUI.WindowTitle = {title}')
+    else:
+        # Unix
+        if os.getenv('TMUX'):
+            # Set the TMUX Window title
+            sys.stdout.write(f"\33k{title}\33")
+        else:
+            # Set the Gnome terminal title
+            sys.stdout.write(f"\33]0;{title}\a")
+        sys.stdout.flush()
+
+
+@dataclass
+class SelectedEMRInstance:
+    cluster_id: str
+    cluster_name: str
+    group_name: str
+    group_idx: str
+    ip: IP
+
+
+@click.group()
+@click.option('-ll/', '--long-log/--no-long-log', default=False, help='Enable long logging')
+@click.option('--verbose/--no-verbose', default=False, help='Enable debug logging')
+@click.option('--quiet/--no-quiet', default=False, help='Disable logging')
+def cli(
+    long_log: bool = False,
+    verbose: bool = False,
+    quiet: bool = False,
+    **kwargs,
+):
+    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>"
+    ) if long_log else "<level>{message}</level>"
+    level = "DEBUG" if verbose else 100 if quiet else "INFO"
+
+    logger.remove()
+    logger.add(sys.stdout, level=level, format=format, colorize=True)
+
+
+@cli.command('ec2')
+@click.option('-p', '--profile', default=None, help='Which AWS profile to use')
+@click.option('-r', '--region', default=None, help='Which AWS region to use')
+@click.option('-u', '--user', default=None, help='Which user to connect as')
+@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,
+    private: bool = True,
+    key_file: str = None,
+    **kwargs,
+):
+    '''
+        Asks user which EC2 instance they want to connect to,
+        then opens an interactive SSH session to the instance
+
+    '''
+    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)
+        terminal_title = f'{user}@{instance_name}'
+
+        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.')
+            exit(1)
+
+
+@cli.command('emr')
+@click.option('-p', '--profile', default=None, help='Which AWS profile to use')
+@click.option('-r', '--region', default=None, help='Which AWS region to use')
+@click.option('-u', '--user', default=None, help='Which user to connect as')
+@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,
+    private: bool = True,
+    key_file: str = None,
+    **kwargs
+):
+    '''
+        Asks user which Cluster and EC2 instance they want to connect to,
+        then opens an interactive SSH session to the instance
+    '''
+    try:
+        b3s = boto3.Session(profile_name=profile, region_name=region)
+        emr = b3s.client("emr")
+
+        emr_instance = get_emr_ssh_options(emr)
+
+        if user is None:
+            user = 'hadoop'
+
+        if key_file is None:
+            key_file = get_emr_ssh_key_file_for_cluster(emr, emr_instance.cluster_id)
+
+        if private:
+            ip = emr_instance.ip.private
+        else:
+            ip = emr_instance.ip.public
+            if ip is None:
+                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 ''
+        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')
+
+        set_terminal_title()
+
+    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.')
+            exit(1)
+
+
+@cli.command('emr-all')
+@click.option('-p', '--profile', default=None, help='Which AWS profile to use')
+@click.option('-r', '--region', default=None, help='Which AWS region to use')
+@click.option('-u', '--user', default=None, help='Which user to connect as')
+@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,
+    private: bool = True,
+    key_file: str = None,
+    **kwargs,
+):
+    '''
+        Asks user which Cluster and EC2 instance they want to connect to,
+        Then prints a tmux cli statement that will open a new session
+        with a window per ec2 instance with ssh shell already opened.
+    '''
+    try:
+        b3s = boto3.Session(profile_name=profile, region_name=region)
+        emr = b3s.client('emr')
+        cluster_id, cluster_name = prompt_for_emr_cluster(emr)
+        grouped_instances = {
+            k.split(" ")[0]: v
+            for k, v in get_emr_instance_ips(emr, cluster_id).items()
+        }
+
+        if user is None:
+            user = 'hadoop'
+
+        if key_file is None:
+            key_file = get_emr_ssh_key_file_for_cluster(emr, cluster_id)
+
+        window_order = ['Master', 'Primary', 'Core', 'Task']
+
+        def get_ip(ip: IP) -> str:
+            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')
+                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')
+            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)
+        ]
+
+        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}"')
+        # 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]}"')
+        # Now switch to the session's first window
+        tmux(f'switch-client -t "{session_name}:{first_window_name}"')
+
+    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.')
+            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
+
+    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}')
+
+        terminal_size = os.get_terminal_size()
+
+        ssh_client = paramiko.SSHClient()
+        # Set hosts key path so we can save to it
+        host_key_path = os.path.expanduser('~/.ssh/known_hosts')
+        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)
+        except paramiko.BadHostKeyException as e:
+            error_message = f'''
+                @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
+                @    WARNING: REMOTE HOST IDENTIFICATION HAS CHANGED!     @
+                @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
+                IT IS POSSIBLE THAT SOMEONE IS DOING SOMETHING NASTY!
+                Someone could be eavesdropping on you right now (man-in-the-middle attack)!
+                It is also possible that a host key has just been changed.
+                The fingerprint for the key sent by the remote host is
+                {e.key.fingerprint}
+                Expected key is
+                {e.expected_key.fingerprint}
+                Please contact your system administrator.
+                Add correct host key in {host_key_path} to get rid of this message.
+                Offending key in {host_key_path}
+                  remove with:
+                  ssh-keygen -f "{host_key_path}" -R "{self.hostname}"
+                Host key for {self.hostname} has changed and you have requested strict checking.
+                Host key verification failed.
+            '''
+            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.invoke_shell()
+
+        self._ssh_client = ssh_client
+        self._channel = channel
+        if self.terminal_title:
+            set_terminal_title(self.terminal_title)
+
+        return channel
+
+    def _launch(self):
+        interactive_shell(self._channel, allow_title_changes=not self.terminal_title)
+
+    def _close(self):
+        self._ssh_client.close()
+        # Reset on connection close
+        if self.terminal_title:
+            set_terminal_title()
+
+    def __enter__(self):
+        return self._open()
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self._launch()
+        self._close()
+
+
+def tmux(command: str):
+    os.system(f'tmux {command}')
+
+
+def try_to_find_ssh_key_file(key_name: str) -> str:
+    '''Recursively iterate over the `~/.ssh/` folder to find the matching key'''
+    if key_name:
+        for dirpath, _, filenames in os.walk(os.path.expanduser('~/.ssh/')):
+            if filenames:
+                for f in filenames:
+                    if f == key_name or os.path.splitext(f)[0] == key_name:
+                        return os.path.join(dirpath, f)
+
+    return None
+
+
+def get_ec2_image_name(instance) -> str:
+    try:
+        return instance.image.name
+    except AttributeError:
+        return None
+
+
+def get_ec2_ssh_options(
+    b3s: boto3.Session,
+    user: str = None,
+    use_private_ip: bool = True,
+    key_file: str = None
+) -> tuple[str, str, str, str]:
+    instance = prompt_for_ec2_instance(b3s)
+    if key_file is None:
+        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'
+
+    if use_private_ip:
+        ip = instance.private_ip_address
+    elif instance.public_ip_address:
+        ip = instance.public_ip_address
+    else:
+        raise ShellError('Public IP was requested, but none was found!')
+
+    return ip, user, key_file, get_ec2_name(instance)
+
+
+def get_emr_ssh_options(
+    emr: "EMRClient",
+) -> SelectedEMRInstance:
+    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_options = [
+        f'{ip.private} ({ip.public})' if ip.public else ip.private
+        for ip in instance_ips
+    ]
+    instance_ip = questionary.select(
+        'Which instance do you want to connect to?',
+        choices=instance_options,
+    ).unsafe_ask()
+    group_idx = instance_options.index(instance_ip)
+
+    return SelectedEMRInstance(cluster_id, cluster_name, group_name, group_idx, instance_ips[group_idx])
+
+
+def get_emr_ssh_key_file_for_cluster(
+    emr: "EMRClient",
+    cluster_id: str,
+) -> str:
+    with click_spinner.spinner():
+        key_name = get_instance_key_name(emr, cluster_id)
+        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
+
+
+def get_running_ec2_instances(b3s: boto3.Session):
+    '''Discover running instances'''
+    with click_spinner.spinner():
+        ec2 = b3s.resource('ec2')
+        running_instances = ec2.instances.filter(
+            Filters=[{'Name': 'instance-state-name', 'Values': ['running']}],
+        )
+
+    return running_instances
+
+
+def prompt_for_ec2_instance(
+    b3s: boto3.Session,
+    prompt: str = 'Which EC2 instance do you want to connect to?',
+) -> "Instance":
+    '''
+        Discovers the runnin EC2 clusters and asks the user which to use.
+
+        Returns the EC2 object that the user selected.
+    '''
+    running_instances = get_running_ec2_instances(b3s)
+
+    name_contains = questionary.text(
+        'Provide a name filter or leave blank to show all',
+    ).unsafe_ask().lower()
+
+    with click_spinner.spinner():
+        grouped_by_name = {
+            get_ec2_name(instance): instance
+            for instance in running_instances
+            if not name_contains or name_contains in get_ec2_name(instance).lower()
+        }
+
+        if not grouped_by_name:
+            raise ShellError('No matching EC2 instances found!')
+
+    # Prompt the user
+    ec2_name = questionary.select(
+        prompt,
+        choices=sorted(grouped_by_name),
+    ).unsafe_ask()
+    instance = grouped_by_name[ec2_name]
+
+    return 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']
+
+    return ec2_instance.instance_id
+
+
+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()
+        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)
+                logger.info('Added host key')
+            else:
+                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")

aws_ssh_utils-0.1.0.dist-info/METADATA

@@ -0,0 +1,113 @@
+Metadata-Version: 2.4
+Name: aws-ssh-utils
+Version: 0.1.0
+Summary: Easy AWS SSHing
+Project-URL: Homepage, https://github.com/mvanderlee/aws_ssh_utils
+Author-email: Michiel Vanderlee <jmt.vanderlee@gmail.com>
+License-Expression: MIT
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Information Technology
+Classifier: Intended Audience :: System Administrators
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: boto3
+Requires-Dist: click
+Requires-Dist: click-spinner
+Requires-Dist: coloredlogs
+Requires-Dist: environs
+Requires-Dist: loguru
+Requires-Dist: paramiko
+Requires-Dist: questionary
+Provides-Extra: dev
+Requires-Dist: boto3-stubs[ec2,emr]; extra == 'dev'
+Provides-Extra: publish
+Requires-Dist: hatch>=1.7.0; extra == 'publish'
+Provides-Extra: test
+Description-Content-Type: text/markdown
+
+# AWS Auth
+
+```shell
+pip install aws-auth-utils
+
+aws configure --profile mfa-source
+
+aws_auth mfa
+```
+
+The commands use [click](https://click.palletsprojects.com/en/stable/) for argument parsing and if required arguments are missing it will prompt you.
+
+To authenticate using your MFA token you will need to have a profile configured using regular an AWS Access Key.
+
+We will use that and your MFA token to generate an authorized session profile.
+By default we will try to use the `mfa-source` and create the `default` profile.
+
+If you only have a single MFA device set up, it will use that automatically. If you have multiple, it will the first one.
+
+## MFA
+
+```shell
+$ aws_auth mfa --help
+Usage: aws_auth mfa [OPTIONS]
+
+Options:
+  -a, --mfa-arn TEXT          The identification number of the MFA device that
+                              is associated with the IAM user. i.e.:
+                              "arn:aws:iam::123456789012:mfa/tony.stark". You
+                              can find this on the IAM page.
+  -c, --code TEXT             The code generated by your MFA device.
+  -d, --duration INTEGER      The duration, in seconds, of the session.
+  -sp, --source-profile TEXT  What AWS profile to get the session token with.
+  -tp, --target-profile TEXT  What AWS profile to store the credentials under.
+  -v, --verbose BOOLEAN
+  --help                      Show this message and exit.
+```
+
+## Assume Role
+
+The assume role is useful for multi-org environments where you want to impersonate a role in a child organization.
+If you access multiple organizations I recommend you set up aliases.
+
+```shell
+aws_auth assume \
+  --role-arn arn:aws:iam::123456789012:role/OrganizationAccountAccessRole \
+  --session-name child_org \
+  --target-profile child_session
+```
+
+```shell
+$ aws_auth assume --help
+Usage: aws_auth assume [OPTIONS]
+
+  Get MFA authenticated and assumed role session credentials and save them to
+  the aws credentials file
+
+  If you have multiple accounts you'd like to switch between, I recommend
+  setting up aliases that call this script with predefined arguments.
+
+Options:
+  -r, --role-arn TEXT         The Arn of the Role to assume.
+  -n, --session-name TEXT     The identifier for the assumed role session.
+  -a, --mfa-arn TEXT          The identification number of the MFA device that
+                              is associated with the IAM user. i.e.:
+                              "arn:aws:iam::123456789012:mfa/tony.stark". You
+                              can find this on the IAM page.
+  -c, --code TEXT             The code generated by your MFA device.
+  -d, --duration INTEGER      The duration, in seconds, of the session.
+                              (defaults to 4 hours)
+  -sp, --source-profile TEXT  What AWS profile to get the session token with.
+  -tp, --target-profile TEXT  What AWS profile to store the credentials under.
+  -v, --verbose BOOLEAN
+  --help                      Show this message and exit.
+```
+

aws_ssh_utils-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+aws_ssh_utils/__init__.py,sha256=kUR5RAFc7HCeiqdlX36dZOHkUI5wI6V_43RpEcD8b-0,22
+aws_ssh_utils/emr_utils.py,sha256=RyAVP2IGXIMfo1pqBmfupFUJruLyxiwEcKYNhA7LK7w,5410
+aws_ssh_utils/interactive_ssh.py,sha256=gcBgDOnIExQmg5Fh6JpAfcA5KeezGjw-Qr2b6rE0LNM,8975
+aws_ssh_utils/ssh.py,sha256=zsqPHMJU3WxhqqhQMTXCzrxmi6IfUliTzgOLyQjTc9A,17375
+aws_ssh_utils-0.1.0.dist-info/METADATA,sha256=Ug2CW31EHg_iU2p0h3LcNsMff3Mq5XenzWo-X2-uLys,4321
+aws_ssh_utils-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+aws_ssh_utils-0.1.0.dist-info/entry_points.txt,sha256=H3d3Sn9YhxItIEzmMmnrOIqfvwGuLHEYbhiI64itUV4,50
+aws_ssh_utils-0.1.0.dist-info/RECORD,,

aws_ssh_utils-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

aws_ssh_utils-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+aws_ssh = aws_ssh_utils.ssh:cli