aws-ssh-utils 0.2.6__tar.gz → 0.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (21) hide show
  1. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/PKG-INFO +13 -7
  2. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/README.md +10 -6
  3. aws_ssh_utils-0.3.0/aws_ssh_utils/__init__.py +1 -0
  4. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/cli_utils.py +53 -7
  5. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/commands/app.py +13 -6
  6. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/commands/ec2.py +1 -1
  7. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/commands/emr.py +5 -3
  8. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/connection.py +112 -65
  9. aws_ssh_utils-0.3.0/aws_ssh_utils/session.py +128 -0
  10. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/ssh.py +5 -1
  11. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/terminal.py +6 -5
  12. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/pyproject.toml +2 -0
  13. aws_ssh_utils-0.2.6/aws_ssh_utils/__init__.py +0 -1
  14. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/.gitignore +0 -0
  15. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/LICENSE +0 -0
  16. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/commands/__init__.py +0 -0
  17. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/commands/emr_all.py +0 -0
  18. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/commands/healthcheck.py +0 -0
  19. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/ec2_utils.py +0 -0
  20. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/emr_utils.py +0 -0
  21. {aws_ssh_utils-0.2.6 → aws_ssh_utils-0.3.0}/aws_ssh_utils/py.typed +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: aws-ssh-utils
3
- Version: 0.2.6
3
+ Version: 0.3.0
4
4
  Summary: Easy AWS SSHing
5
5
  Project-URL: Homepage, https://github.com/mvanderlee/aws-ssh-utils
6
6
  Author-email: Michiel Vanderlee <jmt.vanderlee@gmail.com>
@@ -26,7 +26,9 @@ Requires-Dist: boto3
26
26
  Requires-Dist: click
27
27
  Requires-Dist: loguru
28
28
  Requires-Dist: paramiko
29
+ Requires-Dist: ptyprocess>=0.7.0; sys_platform != 'win32'
29
30
  Requires-Dist: pyte>=0.8.2
31
+ Requires-Dist: pywinpty>=3.0.0; sys_platform == 'win32'
30
32
  Requires-Dist: questionary
31
33
  Requires-Dist: requests>=2.34.2
32
34
  Requires-Dist: rich>=15.0.0
@@ -58,17 +60,21 @@ It utilizes [questionary](https://pypi.org/project/questionary/) to ask you whic
58
60
 
59
61
  ## App
60
62
 
61
- Terminal UI to browse EC2 instances and EMR clusters, with SSH shells in tabs.
63
+ Terminal UI to browse EC2 instances and EMR clusters, with SSM or SSH shells in tabs.
62
64
  It's the default command, so `aws_ssh` and `aws_ssh -p my-profile` launch it too.
63
- Each connection tries, in order: `~/.ssh/config`, SSH over AWS SSM (`session-manager-plugin`), opkssh (via instance/cluster tags),
65
+ Each connection tries, in order: `~/.ssh/config`, a keyless AWS SSM shell, opkssh (via instance/cluster tags),
64
66
  the instance's EC2 key found in `~/.ssh`, then your default keys/ssh-agent.
65
- Unknown host keys are accepted and appended to `~/.ssh/known_hosts`; changed host keys are rejected.
67
+ SSM uses `aws ssm start-session --target INSTANCE_ID` with the selected profile and region. It requires the AWS CLI,
68
+ `session-manager-plugin`, and Session Manager access to the instance, but no SSH key, SSH server, or inbound port 22.
69
+ The shell runs as the Session Manager configured user (usually `ssm-user`); SSH usernames and key options only apply
70
+ to SSH connections. The `ec2` and `emr` commands use the same connection order; `emr-all` still uses SSH in tmux.
71
+ For SSH connections, unknown host keys are accepted and appended to `~/.ssh/known_hosts`; changed host keys are rejected.
66
72
 
67
73
  ```shell
68
74
  $ aws_ssh app --help
69
75
  Usage: aws_ssh app [OPTIONS]
70
76
 
71
- Browse EC2 instances and EMR clusters, and open SSH shells in tabs.
77
+ Browse EC2 instances and EMR clusters, and open SSM or SSH shells in tabs.
72
78
 
73
79
  Options:
74
80
  -p, --profile TEXT Which AWS profile to use
@@ -92,7 +98,7 @@ $ aws_auth ec2 --help
92
98
  Usage: aws_ssh ec2 [OPTIONS]
93
99
 
94
100
  Asks user which EC2 instance they want to connect to, then opens an
95
- interactive SSH session to the instance
101
+ interactive SSM or SSH session to the instance
96
102
 
97
103
  Options:
98
104
  -p, --profile TEXT Which AWS profile to use
@@ -114,7 +120,7 @@ Select the EMR cluster, instance group, and instance to connect to.
114
120
  $ aws_ssh emr --help
115
121
  Usage: aws_ssh emr [OPTIONS]
116
122
 
117
- Asks user which Cluster and EC2 instance they want to connect to, then opens an interactive SSH session to the instance
123
+ Asks user which Cluster and EC2 instance they want to connect to, then opens an interactive SSM or SSH session to the instance
118
124
 
119
125
  Options:
120
126
  -p, --profile TEXT Which AWS profile to use
@@ -22,17 +22,21 @@ It utilizes [questionary](https://pypi.org/project/questionary/) to ask you whic
22
22
 
23
23
  ## App
24
24
 
25
- Terminal UI to browse EC2 instances and EMR clusters, with SSH shells in tabs.
25
+ Terminal UI to browse EC2 instances and EMR clusters, with SSM or SSH shells in tabs.
26
26
  It's the default command, so `aws_ssh` and `aws_ssh -p my-profile` launch it too.
27
- Each connection tries, in order: `~/.ssh/config`, SSH over AWS SSM (`session-manager-plugin`), opkssh (via instance/cluster tags),
27
+ Each connection tries, in order: `~/.ssh/config`, a keyless AWS SSM shell, opkssh (via instance/cluster tags),
28
28
  the instance's EC2 key found in `~/.ssh`, then your default keys/ssh-agent.
29
- Unknown host keys are accepted and appended to `~/.ssh/known_hosts`; changed host keys are rejected.
29
+ SSM uses `aws ssm start-session --target INSTANCE_ID` with the selected profile and region. It requires the AWS CLI,
30
+ `session-manager-plugin`, and Session Manager access to the instance, but no SSH key, SSH server, or inbound port 22.
31
+ The shell runs as the Session Manager configured user (usually `ssm-user`); SSH usernames and key options only apply
32
+ to SSH connections. The `ec2` and `emr` commands use the same connection order; `emr-all` still uses SSH in tmux.
33
+ For SSH connections, unknown host keys are accepted and appended to `~/.ssh/known_hosts`; changed host keys are rejected.
30
34
 
31
35
  ```shell
32
36
  $ aws_ssh app --help
33
37
  Usage: aws_ssh app [OPTIONS]
34
38
 
35
- Browse EC2 instances and EMR clusters, and open SSH shells in tabs.
39
+ Browse EC2 instances and EMR clusters, and open SSM or SSH shells in tabs.
36
40
 
37
41
  Options:
38
42
  -p, --profile TEXT Which AWS profile to use
@@ -56,7 +60,7 @@ $ aws_auth ec2 --help
56
60
  Usage: aws_ssh ec2 [OPTIONS]
57
61
 
58
62
  Asks user which EC2 instance they want to connect to, then opens an
59
- interactive SSH session to the instance
63
+ interactive SSM or SSH session to the instance
60
64
 
61
65
  Options:
62
66
  -p, --profile TEXT Which AWS profile to use
@@ -78,7 +82,7 @@ Select the EMR cluster, instance group, and instance to connect to.
78
82
  $ aws_ssh emr --help
79
83
  Usage: aws_ssh emr [OPTIONS]
80
84
 
81
- Asks user which Cluster and EC2 instance they want to connect to, then opens an interactive SSH session to the instance
85
+ Asks user which Cluster and EC2 instance they want to connect to, then opens an interactive SSM or SSH session to the instance
82
86
 
83
87
  Options:
84
88
  -p, --profile TEXT Which AWS profile to use
@@ -0,0 +1 @@
1
+ __version__ = "0.3.0"
@@ -6,10 +6,11 @@ import json
6
6
  import os
7
7
  import subprocess
8
8
  import sys
9
- from collections.abc import Callable
9
+ import time
10
+ from collections.abc import Callable, Sequence
10
11
  from pathlib import Path
11
12
  from shutil import which
12
- from typing import ParamSpec, TypeVar
13
+ from typing import ParamSpec, TypeVar, cast
13
14
 
14
15
  import click
15
16
  import paramiko
@@ -17,11 +18,14 @@ import questionary
17
18
  from botocore.exceptions import ClientError
18
19
  from loguru import logger
19
20
  from rich.console import Console
21
+ from rich.live import Live
20
22
  from rich.markup import escape
23
+ from rich.spinner import Spinner
21
24
  from rich.status import Status
25
+ from rich.text import Text
22
26
 
23
27
  from aws_ssh_utils.commands.healthcheck import is_aws_ssm_installed, is_opkssh_installed
24
- from aws_ssh_utils.connection import Environment, Target, connect, open_shell_channel
28
+ from aws_ssh_utils.connection import Attempt, Environment, Target, connect, open_shell_channel
25
29
  from aws_ssh_utils.terminal import ShellApp
26
30
 
27
31
  P = ParamSpec('P')
@@ -37,6 +41,7 @@ WINDOWS_TERMINAL_SETTINGS = (
37
41
  """Relative to %LOCALAPPDATA%: Store, Store preview, and unpackaged installs."""
38
42
 
39
43
  console = Console(stderr=True)
44
+ connecting_spinner = Spinner('dots')
40
45
 
41
46
 
42
47
  class ShellError(Exception):
@@ -50,6 +55,29 @@ def spinner(message: str = '') -> Status:
50
55
  return console.status(message, spinner='dots')
51
56
 
52
57
 
58
+ def render_attempts(attempts: Sequence[Attempt]) -> Text:
59
+ """The connection checklist: done attempts, the current one with a spinner, and the remaining ones muted."""
60
+
61
+ def line(attempt: Attempt) -> Text:
62
+ match attempt.state:
63
+ case 'connecting':
64
+ frame = cast(Text, connecting_spinner.render(time.monotonic()))
65
+ return Text.assemble(frame, f" Connecting via {attempt.via} ...")
66
+ case 'connected':
67
+ return Text(f"✓ Connected via {attempt.via}", 'green')
68
+ case 'failed':
69
+ return Text.assemble(
70
+ f"❌ Failed to connect via {attempt.via} - {attempt.detail}\n",
71
+ (f" {attempt.command}", 'dim'),
72
+ )
73
+ case 'skipped':
74
+ return Text(f" {attempt.detail}, skipping.", 'dim')
75
+ case 'pending':
76
+ return Text(f" {attempt.via}", 'dim')
77
+
78
+ return Text('\n').join(line(a) for a in attempts)
79
+
80
+
53
81
  def handle_errors(func: Callable[P, R]) -> Callable[P, R]:
54
82
  """Log ShellErrors and AWS errors, and exit, instead of printing a traceback."""
55
83
 
@@ -145,10 +173,28 @@ def open_shell(target: Target, env: Environment, title: str, scrollback: int | N
145
173
 
146
174
  `scrollback` defaults to the terminal's own setting, see `detect_scrollback`.
147
175
  """
148
- result = connect(target, env, on_output=logger.info, on_waiting=lambda: None, confirm_host_key=confirm_host_key)
149
- if isinstance(result, list):
150
- for line in result:
151
- console.print(escape(line))
176
+ attempts: list[Attempt] = []
177
+ with Live(console=console, get_renderable=lambda: render_attempts(attempts)) as live:
178
+
179
+ def on_progress(snapshot: list[Attempt]):
180
+ attempts[:] = snapshot
181
+
182
+ def confirm(hostname: str, key: paramiko.PKey) -> bool:
183
+ live.stop()
184
+ try:
185
+ return confirm_host_key(hostname, key)
186
+ finally:
187
+ live.start()
188
+
189
+ result = connect(
190
+ target,
191
+ env,
192
+ on_output=lambda line: live.console.print(escape(line)),
193
+ on_waiting=lambda: None,
194
+ on_progress=on_progress,
195
+ confirm_host_key=confirm,
196
+ )
197
+ if result is None:
152
198
  raise ShellError(f'Failed to connect to {target.user}@{target.ip}')
153
199
 
154
200
  channel = open_shell_channel(result.client)
@@ -22,8 +22,8 @@ from textual.widgets import Footer, Input, OptionList, RichLog, Static, TabbedCo
22
22
  from textual.widgets.option_list import Option
23
23
  from textual.widgets.tree import TreeNode
24
24
 
25
- from aws_ssh_utils.cli_utils import detect_environment, detect_scrollback, scrollback_option
26
- from aws_ssh_utils.connection import Environment, Target, connect, open_shell_channel
25
+ from aws_ssh_utils.cli_utils import detect_environment, detect_scrollback, render_attempts, scrollback_option
26
+ from aws_ssh_utils.connection import Attempt, Environment, Target, connect, open_shell_channel
27
27
  from aws_ssh_utils.ec2_utils import ec2_name, ec2_target, get_running_ec2_instances
28
28
  from aws_ssh_utils.emr_utils import emr_target, get_emr_clusters, get_emr_instances, group_role, group_sort_key
29
29
  from aws_ssh_utils.terminal import ShellStatus, Terminal
@@ -48,7 +48,7 @@ ANIMATION_INTERVAL = 1 / 12
48
48
  @click.option('-r', '--region', default=None, help='Which AWS region to use')
49
49
  @scrollback_option
50
50
  def app(profile: str | None = None, region: str | None = None, scrollback: int | None = None, **kwargs: Any):
51
- """Browse EC2 instances and EMR clusters, and open SSH shells in tabs."""
51
+ """Browse EC2 instances and EMR clusters, and open SSM or SSH shells in tabs."""
52
52
  try:
53
53
  session = boto3.Session(profile_name=profile, region_name=region)
54
54
  session.client('sts').get_caller_identity()
@@ -99,8 +99,10 @@ class ShellPane(TabPane):
99
99
  self.scrollback = scrollback
100
100
  self.status: Status = 'connecting'
101
101
  self.spinner = Spinner('dots')
102
+ self.attempts: list[Attempt] = []
102
103
 
103
104
  def compose(self) -> ComposeResult:
105
+ yield Static(classes='attempts')
104
106
  yield RichLog(wrap=True, markup=True)
105
107
 
106
108
  def on_mount(self):
@@ -110,6 +112,11 @@ class ShellPane(TabPane):
110
112
  def update_label(self):
111
113
  icon = cast(Text, self.spinner.render(time.monotonic())) if self.status == 'connecting' else STATUS_ICONS[self.status]
112
114
  self.query_ancestor(TabbedContent).get_tab(self).label = Text.assemble(icon, ' ', self.title_)
115
+ self.query_one('.attempts', Static).update(render_attempts(self.attempts))
116
+
117
+ def set_attempts(self, attempts: list[Attempt]):
118
+ self.attempts = attempts
119
+ self.update_label()
113
120
 
114
121
  def set_status(self, status: Status):
115
122
  self.status = status
@@ -131,10 +138,9 @@ class ShellPane(TabPane):
131
138
  self.env,
132
139
  on_output=write,
133
140
  on_waiting=lambda: self.app.call_from_thread(self.set_status, 'waiting'),
141
+ on_progress=lambda attempts: self.app.call_from_thread(self.set_attempts, attempts),
134
142
  )
135
- if isinstance(result, list):
136
- for line in result:
137
- write(line)
143
+ if result is None:
138
144
  self.app.call_from_thread(self.set_status, 'failed')
139
145
  return
140
146
 
@@ -149,6 +155,7 @@ class ShellPane(TabPane):
149
155
 
150
156
  async def attach(self, via: str, terminal: Terminal):
151
157
  await self.query_one(RichLog).remove()
158
+ await self.query_one('.attempts').remove()
152
159
  await self.mount(ShellStatus(via), terminal)
153
160
  self.set_status('connected')
154
161
 
@@ -41,7 +41,7 @@ def ec2(
41
41
  ):
42
42
  """
43
43
  Asks user which EC2 instance they want to connect to,
44
- then opens an interactive SSH session to the instance
44
+ then opens an interactive SSM or SSH session to the instance
45
45
  """
46
46
  session = boto3.Session(profile_name=profile, region_name=region)
47
47
  ec2_client = session.client('ec2')
@@ -45,7 +45,7 @@ def emr(
45
45
  ):
46
46
  """
47
47
  Asks user which Cluster and EC2 instance they want to connect to,
48
- then opens an interactive SSH session to the instance
48
+ then opens an interactive SSM or SSH session to the instance
49
49
  """
50
50
  session = boto3.Session(profile_name=profile, region_name=region)
51
51
  emr_client = session.client('emr')
@@ -65,9 +65,11 @@ def emr(
65
65
  with spinner():
66
66
  target = emr_target(emr_client, cluster_id, instance_ip.instance_id or '', ip)
67
67
  target = replace(target, user=user or target.user, key_file=key_file)
68
+ env = detect_environment(profile, session.region_name)
68
69
 
69
70
  if (
70
- key_file is None
71
+ not env.ssm_installed
72
+ and key_file is None
71
73
  and target.opkssh_provider is None
72
74
  and find_ssh_key_file(target.key_name) is None
73
75
  and not questionary.confirm(f'Could not find the ssh key {target.key_name}, would you like to continue?').unsafe_ask()
@@ -78,7 +80,7 @@ def emr(
78
80
  postfix = f'[{group_idx}]' if group_idx else ''
79
81
  open_shell(
80
82
  target,
81
- detect_environment(profile, session.region_name),
83
+ env,
82
84
  title=f'{cluster_name} - {group}{postfix}',
83
85
  scrollback=scrollback,
84
86
  initial_input=CORE_NODE_INPUT if group.lower().startswith('core') else b'',
@@ -1,4 +1,4 @@
1
- """SSH connection strategies, tried in order until one succeeds."""
1
+ """SSM and SSH connection strategies, tried in order until one succeeds."""
2
2
 
3
3
  import os
4
4
  import shlex
@@ -7,17 +7,19 @@ import subprocess
7
7
  import threading
8
8
  import time
9
9
  from collections.abc import Callable, Sequence
10
- from dataclasses import dataclass
10
+ from dataclasses import dataclass, field, replace
11
11
  from fnmatch import fnmatch
12
12
  from functools import partial
13
13
  from hashlib import sha256
14
14
  from pathlib import Path
15
- from typing import IO, cast
15
+ from typing import IO, Literal, cast
16
16
 
17
17
  import paramiko
18
18
  from loguru import logger
19
19
  from typing_extensions import override
20
20
 
21
+ from aws_ssh_utils.session import ShellChannel, SSMSession
22
+
21
23
  # CSV of issuer,client
22
24
  OPKSSH_PROVIDER_TAG = 'opkssh_provider'
23
25
  # EMR doesn't support comma's in tags.
@@ -54,97 +56,121 @@ class Environment:
54
56
 
55
57
  @dataclass(frozen=True)
56
58
  class Connection:
57
- client: paramiko.SSHClient
59
+ client: paramiko.SSHClient | SSMSession
58
60
  via: str
59
61
  """How the connection was made: AWS SSM, opkssh, a key name, ..."""
60
62
 
61
63
 
64
+ @dataclass(frozen=True)
65
+ class Attempt:
66
+ via: str
67
+ command: str
68
+ """The equivalent command line, to retry by hand."""
69
+ open_client: Callable[[], paramiko.SSHClient | SSMSession] = field(repr=False, compare=False)
70
+ state: Literal['pending', 'connecting', 'connected', 'failed', 'skipped'] = 'pending'
71
+ detail: str = ''
72
+ """The error when failed, the reason when skipped."""
73
+
74
+
75
+ OnProgress = Callable[[list[Attempt]], None]
76
+
77
+
62
78
  def connect(
63
79
  target: Target,
64
80
  env: Environment,
65
81
  on_output: OnOutput,
66
82
  on_waiting: Callable[[], None],
83
+ on_progress: OnProgress,
67
84
  confirm_host_key: ConfirmHostKey | None = None,
68
- ) -> Connection | list[str]:
69
- """Return the first successful connection, or the failure checklist.
85
+ ) -> Connection | None:
86
+ """Return the first successful connection, or None when all fail.
70
87
 
71
88
  Tried in order:
72
89
  1. ~/.ssh/config, when a Host pattern (other than `*`) matches the instance id, an alias or the IP.
73
- 2. SSH over AWS SSM to the instance id, with any existing opkssh/EC2 key plus the default keys.
90
+ 2. A keyless AWS SSM shell, authenticated with the AWS profile.
74
91
  3. opkssh, when the instance/cluster has opkssh tags. May open a browser login.
75
92
  4. The EC2 key pair (or `target.key_file`), when a matching file is found in ~/.ssh.
76
93
  5. The ssh-agent and default ~/.ssh/id_* keys.
77
94
 
95
+ `on_progress` gets a snapshot of all attempts whenever one changes state.
78
96
  Unknown host keys are trusted, unless `confirm_host_key` declines them.
79
97
  """
80
- failures: list[str] = []
81
98
  open_ssh_ = partial(open_ssh, confirm_host_key=confirm_host_key)
99
+ attempts: list[Attempt] = []
82
100
 
83
- def attempt(via: str, failure: str, open_client: Callable[[], paramiko.SSHClient]) -> Connection | None:
84
- logger.info(f"Connecting via {via}")
101
+ def try_attempt(i: int) -> Connection | None:
102
+ attempt = attempts[i]
103
+ if attempt.state == 'skipped':
104
+ on_progress(list(attempts))
105
+ return None
106
+ attempts[i] = replace(attempt, state='connecting')
107
+ on_progress(list(attempts))
85
108
  try:
86
- return Connection(open_client(), via)
109
+ client = attempt.open_client()
87
110
  except Exception as e:
88
- logger.debug(f"{failure}: {e!r}")
89
- failures.append(f"{failure}\n {e or type(e).__name__}")
111
+ logger.debug(f"Failed to connect via {attempt.via}: {e!r}")
112
+ attempts[i] = replace(attempt, state='failed', detail=str(e) or type(e).__name__)
113
+ on_progress(list(attempts))
90
114
  return None
115
+ attempts[i] = replace(attempt, state='connected')
116
+ on_progress(list(attempts))
117
+ return Connection(client, attempt.via)
91
118
 
92
119
  config = find_ssh_config(target)
93
120
  if config is not None:
94
121
  alias, options = config
95
- if connection := attempt(
96
- f"ssh config ({alias})",
97
- f":x: Failed with ssh config - ssh {alias}",
98
- lambda: open_with_config(alias, options, target, confirm_host_key),
99
- ):
100
- return connection
122
+ attempts.append(
123
+ Attempt(
124
+ f"ssh config ({alias})",
125
+ shlex.join(['ssh', alias]),
126
+ lambda: open_with_config(alias, options, target, confirm_host_key),
127
+ ),
128
+ )
129
+
130
+ command = ssm_command(target.instance_id, env)
131
+ ssm = Attempt('AWS SSM', shlex.join(command), lambda: SSMSession.open(command))
132
+ attempts.append(
133
+ ssm if env.ssm_installed else replace(ssm, state='skipped', detail="AWS CLI or session-manager-plugin not installed"),
134
+ )
135
+ for i in range(len(attempts)):
136
+ if result := try_attempt(i):
137
+ return result
101
138
 
139
+ # Key discovery is only needed for the remaining SSH fallbacks.
140
+ fallback_start = len(attempts)
102
141
  ec2_key = target.key_file or find_ssh_key_file(target.key_name)
103
- known_keys = [k for k in (valid_opkssh_key_file(target.opkssh_provider), ec2_key) if k]
104
-
105
- if not env.ssm_installed:
106
- failures.append(":x: AWS session-manager-plugin not installed")
107
- elif connection := attempt(
108
- "AWS SSM",
109
- ":x: Failed to start session with AWS session-manager",
110
- lambda: open_ssh_(
111
- target.instance_id,
112
- target.user,
113
- known_keys,
114
- use_defaults=True,
115
- sock=proxy_socket(ssm_command(target.instance_id, env), on_output),
116
- ),
117
- ):
118
- return connection
119
-
120
- if target.opkssh_provider is not None:
121
- provider = target.opkssh_provider
122
- if not env.opkssh_installed:
123
- failures.append(":x: Failed to authenticate with opkssh\n opkssh not installed")
124
- elif connection := attempt(
142
+
143
+ if (provider := target.opkssh_provider) is not None:
144
+ opkssh_key = str(opkssh_key_path(provider))
145
+ opkssh = Attempt(
125
146
  "opkssh",
126
- ":x: Failed to authenticate with opkssh",
147
+ f"{shlex.join(opkssh_login_command(provider, opkssh_key))} && "
148
+ + ssh_command(target.user, target.ip, [opkssh_key], identities_only=True),
127
149
  lambda: open_ssh_(target.ip, target.user, [opkssh_login(provider, on_output, on_waiting)], use_defaults=False),
128
- ):
129
- return connection
130
-
131
- if ec2_key is not None and (
132
- connection := attempt(
133
- Path(ec2_key).stem,
134
- f":x: Key found, but failed - ssh -o IdentitiesOnly=yes -i {ec2_key} {target.user}@{target.ip}",
135
- lambda: open_ssh_(target.ip, target.user, [ec2_key], use_defaults=False),
136
150
  )
137
- ):
138
- return connection
151
+ attempts.append(opkssh if env.opkssh_installed else replace(opkssh, state='skipped', detail="opkssh not installed"))
152
+
153
+ if ec2_key is not None:
154
+ attempts.append(
155
+ Attempt(
156
+ Path(ec2_key).stem,
157
+ ssh_command(target.user, target.ip, [ec2_key], identities_only=True),
158
+ lambda: open_ssh_(target.ip, target.user, [ec2_key], use_defaults=False),
159
+ ),
160
+ )
139
161
 
140
- if connection := attempt(
141
- "default key",
142
- f":x: Failed with default key - ssh {target.user}@{target.ip}",
143
- lambda: open_ssh_(target.ip, target.user, [], use_defaults=True),
144
- ):
145
- return connection
162
+ attempts.append(
163
+ Attempt(
164
+ "default key",
165
+ ssh_command(target.user, target.ip),
166
+ lambda: open_ssh_(target.ip, target.user, [], use_defaults=True),
167
+ ),
168
+ )
146
169
 
147
- return failures
170
+ for i in range(fallback_start, len(attempts)):
171
+ if result := try_attempt(i):
172
+ return result
173
+ return None
148
174
 
149
175
 
150
176
  def open_ssh(
@@ -186,7 +212,28 @@ def open_ssh(
186
212
  return client
187
213
 
188
214
 
189
- def open_shell_channel(client: paramiko.SSHClient) -> paramiko.Channel:
215
+ def ssh_command(
216
+ user: str,
217
+ host: str,
218
+ key_files: Sequence[str] = (),
219
+ *,
220
+ identities_only: bool = False,
221
+ proxy: list[str] | None = None,
222
+ ) -> str:
223
+ """The OpenSSH command line equivalent of an `open_ssh` call."""
224
+ cmd = ['ssh']
225
+ if proxy:
226
+ cmd += ['-o', f'ProxyCommand={shlex.join(proxy)}']
227
+ if identities_only:
228
+ cmd += ['-o', 'IdentitiesOnly=yes']
229
+ for key_file in key_files:
230
+ cmd += ['-i', key_file]
231
+ return shlex.join([*cmd, f'{user}@{host}'])
232
+
233
+
234
+ def open_shell_channel(client: paramiko.SSHClient | SSMSession) -> ShellChannel:
235
+ if isinstance(client, SSMSession):
236
+ return client
190
237
  channel = client.get_transport().open_session() # pyright: ignore[reportOptionalMemberAccess]
191
238
  channel.get_pty(term='xterm-256color')
192
239
  channel.invoke_shell()
@@ -261,10 +308,6 @@ def ssm_command(instance_id: str, env: Environment) -> list[str]:
261
308
  'start-session',
262
309
  '--target',
263
310
  instance_id,
264
- '--document-name',
265
- 'AWS-StartSSHSession',
266
- '--parameters',
267
- 'portNumber=22',
268
311
  ]
269
312
  if env.profile:
270
313
  cmd += ['--profile', env.profile]
@@ -357,6 +400,10 @@ def valid_opkssh_key_file(provider: str | None) -> str | None:
357
400
  return str(key) if key.is_file() and key.stat().st_mtime > time.time() - OPKSSH_KEY_MAX_AGE else None
358
401
 
359
402
 
403
+ def opkssh_login_command(provider: str, key_file: str) -> list[str]:
404
+ return ['opkssh', 'login', '--provider', provider, '-i', key_file]
405
+
406
+
360
407
  def opkssh_login(provider: str, on_output: OnOutput, on_waiting: Callable[[], None]) -> str:
361
408
  """Return a valid opkssh key file, logging in when missing or expired."""
362
409
  if key_file := valid_opkssh_key_file(provider):
@@ -368,7 +415,7 @@ def opkssh_login(provider: str, on_output: OnOutput, on_waiting: Callable[[], No
368
415
  path.unlink(missing_ok=True)
369
416
 
370
417
  proc = subprocess.Popen( # noqa: S603
371
- ['opkssh', 'login', '--provider', provider, '-i', str(key)], # noqa: S607
418
+ opkssh_login_command(provider, str(key)),
372
419
  stdout=subprocess.PIPE,
373
420
  stderr=subprocess.STDOUT,
374
421
  text=True,
@@ -0,0 +1,128 @@
1
+ """Terminal channels for native, keyless Session Manager shells."""
2
+
3
+ import os
4
+ import queue
5
+ import re
6
+ import threading
7
+ import time
8
+ from typing import Protocol, cast
9
+
10
+
11
+ class ShellChannel(Protocol):
12
+ def recv(self, nbytes: int, /) -> bytes: ...
13
+ def send(self, data: bytes, /) -> int: ...
14
+ def resize_pty(self, width: int, height: int) -> None: ...
15
+
16
+
17
+ class ShellClient(Protocol):
18
+ def close(self) -> None: ...
19
+
20
+
21
+ class Pty(Protocol):
22
+ def read(self, size: int) -> str: ...
23
+ def write(self, data: str) -> int: ...
24
+ def setwinsize(self, rows: int, cols: int) -> None: ...
25
+ def close(self, force: bool = False) -> None: ...
26
+
27
+
28
+ def spawn_pty(command: list[str]) -> Pty:
29
+ """SSM needs a real terminal for raw input, control keys, and window sizes."""
30
+ env = {**os.environ, 'TERM': 'xterm-256color', 'AWS_PAGER': '', 'AWS_CLI_AUTO_PROMPT': 'off'}
31
+ if os.name == 'nt':
32
+ from winpty import PtyProcess
33
+
34
+ return cast(Pty, PtyProcess.spawn(command, env=env, dimensions=(24, 80)))
35
+
36
+ from ptyprocess import PtyProcessUnicode
37
+
38
+ return cast(Pty, PtyProcessUnicode.spawn(command, env=env, dimensions=(24, 80), echo=False))
39
+
40
+
41
+ class SSMSession:
42
+ """Adapt an AWS CLI pseudo-terminal to the terminal widget's channel interface."""
43
+
44
+ def __init__(self, process: Pty):
45
+ self.process = process
46
+ self.closed = False
47
+ self._output: queue.Queue[bytes | Exception | None] = queue.Queue()
48
+ self._pending = b''
49
+ self._eof = False
50
+ threading.Thread(target=self._read, daemon=True).start()
51
+
52
+ @classmethod
53
+ def open(cls, command: list[str], timeout: float = 30) -> 'SSMSession':
54
+ session = cls(spawn_pty(command))
55
+ try:
56
+ session._wait_started(timeout)
57
+ except BaseException:
58
+ session.close()
59
+ raise
60
+ return session
61
+
62
+ def _read(self):
63
+ try:
64
+ while not self.closed:
65
+ data = self.process.read(65536)
66
+ if data:
67
+ self._output.put(data.encode('utf-8'))
68
+ except EOFError:
69
+ pass
70
+ except Exception as e:
71
+ if not self.closed:
72
+ self._output.put(e)
73
+ finally:
74
+ self._output.put(None)
75
+
76
+ def _wait_started(self, timeout: float):
77
+ """Keep startup output for the terminal, but report CLI failures to the fallback list."""
78
+ deadline = time.monotonic() + timeout
79
+ output = bytearray()
80
+ try:
81
+ while True:
82
+ data = self._output.get(timeout=max(0, deadline - time.monotonic()))
83
+ if isinstance(data, Exception):
84
+ raise OSError(str(data)) from data
85
+ if data is None:
86
+ raise OSError(output.decode('utf-8', errors='replace').strip() or 'SSM exited before starting a session')
87
+ output.extend(data)
88
+ # ConPTY adds terminal escape sequences to console output on Windows.
89
+ text = re.sub(rb'\x1b\[[0-?]*[ -/]*[@-~]', b'', output)
90
+ if re.search(rb'Starting session with SessionId:\s*\S+', text):
91
+ self._pending = bytes(output)
92
+ return
93
+ except queue.Empty:
94
+ detail = output.decode('utf-8', errors='replace').strip()
95
+ raise TimeoutError(f'Timed out starting SSM session{": " + detail if detail else ""}') from None
96
+
97
+ def recv(self, nbytes: int) -> bytes:
98
+ if nbytes <= 0:
99
+ return b''
100
+ while not self._pending and not self._eof:
101
+ data = self._output.get()
102
+ if isinstance(data, Exception):
103
+ raise OSError(str(data)) from data
104
+ if data is None:
105
+ self._eof = True
106
+ else:
107
+ self._pending = data
108
+ data, self._pending = self._pending[:nbytes], self._pending[nbytes:]
109
+ return data
110
+
111
+ def send(self, data: bytes) -> int:
112
+ try:
113
+ self.process.write(data.decode('utf-8'))
114
+ except EOFError as e:
115
+ raise OSError('SSM session is closed') from e
116
+ return len(data)
117
+
118
+ def resize_pty(self, width: int, height: int):
119
+ if not self.closed:
120
+ self.process.setwinsize(height, width)
121
+
122
+ def close(self):
123
+ if not self.closed:
124
+ self.closed = True
125
+ try:
126
+ self.process.close(force=True)
127
+ finally:
128
+ self._output.put(None)
@@ -53,8 +53,12 @@ def cli(
53
53
  )
54
54
  level = "DEBUG" if verbose else 100 if quiet else "INFO"
55
55
 
56
+ def to_stdout(message: str):
57
+ # Look up sys.stdout per message, so rich.live can redirect it above its display.
58
+ sys.stdout.write(message)
59
+
56
60
  logger.remove()
57
- logger.add(sys.stdout, level=level, format=log_format, colorize=True)
61
+ logger.add(to_stdout, level=level, format=log_format, colorize=True)
58
62
 
59
63
 
60
64
  for command in (app, ec2, emr, emr_all, healthcheck):
@@ -1,4 +1,4 @@
1
- """Textual widgets rendering a paramiko shell channel through the pyte VT100 emulator."""
1
+ """Textual widgets rendering SSH and SSM shell channels through the pyte VT100 emulator."""
2
2
 
3
3
  import datetime as dt
4
4
  import time
@@ -7,7 +7,6 @@ from collections.abc import Sequence
7
7
  from functools import lru_cache
8
8
  from typing import Any
9
9
 
10
- import paramiko
11
10
  import pyte
12
11
  from pyte.screens import Char, Margins
13
12
  from rich.segment import Segment
@@ -22,6 +21,8 @@ from textual.widget import Widget
22
21
  from textual.widgets import Static
23
22
  from typing_extensions import override
24
23
 
24
+ from aws_ssh_utils.session import ShellChannel, ShellClient
25
+
25
26
  KEYS = {
26
27
  'enter': '\r',
27
28
  'tab': '\t',
@@ -108,7 +109,7 @@ def cell_style(
108
109
  class ChannelScreen(pyte.Screen):
109
110
  """Answers terminal queries over the channel, and keeps the lines scrolled off the top as history."""
110
111
 
111
- def __init__(self, channel: paramiko.Channel, columns: int, lines: int, scrollback: int = 0):
112
+ def __init__(self, channel: ShellChannel, columns: int, lines: int, scrollback: int = 0):
112
113
  super().__init__(columns, lines)
113
114
  self.channel = channel
114
115
  self.history: deque[tuple[Char, ...]] = deque(maxlen=scrollback)
@@ -145,7 +146,7 @@ class Terminal(Widget, can_focus=True):
145
146
  class Closed(Message):
146
147
  pass
147
148
 
148
- def __init__(self, client: paramiko.SSHClient, channel: paramiko.Channel, scrollback: int = 0, **kwargs: object):
149
+ def __init__(self, client: ShellClient, channel: ShellChannel, scrollback: int = 0, **kwargs: object):
149
150
  super().__init__(**kwargs) # pyright: ignore[reportArgumentType]
150
151
  self.client = client
151
152
  self.channel = channel
@@ -287,7 +288,7 @@ class ShellApp(App[None]):
287
288
  # Its priority ctrl+p binding would steal shell history navigation from the terminal.
288
289
  ENABLE_COMMAND_PALETTE = False
289
290
 
290
- def __init__(self, via: str, client: paramiko.SSHClient, channel: paramiko.Channel, scrollback: int = 0):
291
+ def __init__(self, via: str, client: ShellClient, channel: ShellChannel, scrollback: int = 0):
291
292
  super().__init__()
292
293
  self.via = via
293
294
  self.client = client
@@ -34,6 +34,8 @@ dependencies = [
34
34
  "loguru",
35
35
  "paramiko",
36
36
  "pyte>=0.8.2",
37
+ "ptyprocess>=0.7.0; sys_platform != 'win32'",
38
+ "pywinpty>=3.0.0; sys_platform == 'win32'",
37
39
  "questionary",
38
40
  "requests>=2.34.2",
39
41
  "rich>=15.0.0",
@@ -1 +0,0 @@
1
- __version__ = "0.2.6"
File without changes
File without changes