ssm-cli 1.0.0.dev0__tar.gz → 1.0.0.dev3__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 (37) hide show
  1. {ssm_cli-1.0.0.dev0/ssm_cli.egg-info → ssm_cli-1.0.0.dev3}/PKG-INFO +1 -1
  2. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/cli.py +1 -8
  3. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/cli_args.py +16 -14
  4. ssm_cli-1.0.0.dev3/ssm_cli/commands/ssh_proxy/forwarding.py +290 -0
  5. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/ssh_proxy/server.py +24 -54
  6. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/config.py +2 -1
  7. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/instances.py +17 -27
  8. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/logging.py +2 -2
  9. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/xdg.py +4 -6
  10. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3/ssm_cli.egg-info}/PKG-INFO +1 -1
  11. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli.egg-info/SOURCES.txt +1 -0
  12. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/LICENCE +0 -0
  13. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/README.md +0 -0
  14. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/pyproject.toml +0 -0
  15. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/setup.cfg +0 -0
  16. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/__init__.py +0 -0
  17. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/__main__.py +0 -0
  18. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/aws.py +0 -0
  19. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/__init__.py +0 -0
  20. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/base.py +0 -0
  21. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/list.py +0 -0
  22. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/setup.py +0 -0
  23. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/shell.py +0 -0
  24. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/ssh_proxy/__init__.py +0 -0
  25. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/ssh_proxy/channels.py +0 -0
  26. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/ssh_proxy/shell.py +0 -0
  27. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/commands/ssh_proxy/socket.py +0 -0
  28. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/console.py +0 -0
  29. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/selectors/__init__.py +0 -0
  30. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/selectors/first.py +0 -0
  31. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/selectors/tui/__init__.py +0 -0
  32. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/selectors/tui/posix.py +0 -0
  33. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli/selectors/tui/win.py +0 -0
  34. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli.egg-info/dependency_links.txt +0 -0
  35. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli.egg-info/entry_points.txt +0 -0
  36. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli.egg-info/requires.txt +0 -0
  37. {ssm_cli-1.0.0.dev0 → ssm_cli-1.0.0.dev3}/ssm_cli.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ssm-cli
3
- Version: 1.0.0.dev0
3
+ Version: 1.0.0.dev3
4
4
  Summary: CLI tool to help with SSM functionality, aimed at adminstrators
5
5
  Author-email: Simon Fletcher <simon.fletcher@lexisnexisrisk.com>
6
6
  License: MIT License
@@ -1,7 +1,6 @@
1
1
  import sys
2
2
  import logging
3
3
 
4
- from rich_argparse import ArgumentDefaultsRichHelpFormatter
5
4
 
6
5
  import confclasses
7
6
  from ssm_cli.config import CONFIG
@@ -31,13 +30,7 @@ def cli(argv: list = None) -> int:
31
30
 
32
31
  logger.debug(f"CLI called with {argv}")
33
32
 
34
- # Build the actual parser
35
- parser = CliArgumentParser(
36
- prog="ssm",
37
- description="tool to manage AWS SSM",
38
- formatter_class=ArgumentDefaultsRichHelpFormatter,
39
- )
40
- parser.add_global_argument("--profile", type=str, help="Which AWS profile to use")
33
+ parser = CliArgumentParser()
41
34
 
42
35
  for name, command in COMMANDS.items():
43
36
  command_parser = parser.add_command_parser(name, command.HELP)
@@ -2,12 +2,15 @@ import argparse
2
2
  import sys
3
3
  from confclasses import fields, is_confclass
4
4
  from ssm_cli.config import CONFIG
5
+ from rich_argparse import ArgumentDefaultsRichHelpFormatter
5
6
  import os
6
7
 
7
8
  # The long term aim is to move this module into the config module and "bind" config to cli arguments
8
9
  # This will need a rethink of where some of the arguments/config come from because right now they are in the commands modules
9
10
 
10
11
  class CliNamespace(argparse.Namespace):
12
+ global_args: "CliNamespace"
13
+
11
14
  def update_config(self):
12
15
  self._do_update_config(CONFIG, vars(self.global_args))
13
16
 
@@ -28,19 +31,20 @@ class CliNamespace(argparse.Namespace):
28
31
  ARGS = CliNamespace()
29
32
 
30
33
  class CliArgumentParser(argparse.ArgumentParser):
31
- def __init__(self, *args, help_as_global=True, **kwargs):
34
+ def __init__(self):
32
35
  self.global_args_parser = argparse.ArgumentParser(add_help=False)
33
36
  self.global_args_parser_group = self.global_args_parser.add_argument_group("Global Options")
34
-
35
- self.help_as_global = help_as_global
36
- if help_as_global:
37
- kwargs['add_help'] = False
38
- # we cannot use help action here because it will just return the global arguments
39
- self.global_args_parser_group.add_argument('--help', '-h', action="store_true", help="show this help message and exit")
40
-
37
+ self.global_args_parser_group.add_argument("--profile", type=str, help="Which AWS profile to use")
38
+ self.global_args_parser_group.add_argument('--help', '-h', action="store_true", help="show this help message and exit")
41
39
  self.global_args_parser_group.add_argument('--version', '-v', action=VersionAction)
42
40
 
43
- super().__init__(*args, **kwargs)
41
+ super().__init__(
42
+ prog="ssm",
43
+ description="tool to manage AWS SSM",
44
+ formatter_class=ArgumentDefaultsRichHelpFormatter,
45
+ add_help=False
46
+ )
47
+
44
48
  self._command_subparsers = self.add_subparsers(title="Commands", dest="command", metavar="<command>", parser_class=argparse.ArgumentParser)
45
49
  self._command_subparsers_map = {}
46
50
 
@@ -57,12 +61,13 @@ class CliArgumentParser(argparse.ArgumentParser):
57
61
 
58
62
  if argv is None:
59
63
  argv = sys.argv[1:]
64
+
60
65
  global_args, unknown = self.global_args_parser.parse_known_args(argv, CliNamespace())
61
66
 
62
67
  super().parse_args(unknown, ARGS)
63
68
  ARGS.global_args = global_args
64
69
 
65
- if self.help_as_global and global_args.help:
70
+ if global_args.help:
66
71
  if ARGS.command and ARGS.command in self._command_subparsers_map:
67
72
  self._command_subparsers_map[ARGS.command].print_help()
68
73
  self.exit()
@@ -77,12 +82,9 @@ class CliArgumentParser(argparse.ArgumentParser):
77
82
  delattr(global_args, 'help')
78
83
 
79
84
  return ARGS
80
-
81
- def add_global_argument(self, *args, **kwargs):
82
- self.global_args_parser_group.add_argument(*args, **kwargs)
83
85
 
84
86
  def add_command_parser(self, name, help):
85
- parser = self._command_subparsers.add_parser(name, help=help, formatter_class=self.formatter_class, parents=[self.global_args_parser], add_help=not self.help_as_global)
87
+ parser = self._command_subparsers.add_parser(name, help=help, formatter_class=self.formatter_class, parents=[self.global_args_parser], add_help=False)
86
88
  self._command_subparsers_map[name] = parser
87
89
  return parser
88
90
 
@@ -0,0 +1,290 @@
1
+ r"""
2
+ This module is split into 2 parts, the proxy side and the manager side.
3
+
4
+ The proxy side is the code that runs in the main process along side the ssh server, it is responsible for handling all the requests for port forwarding and starting the manager process.
5
+
6
+ The manager side is the code that runs in a separate process and is responsible for starting and stopping the session-manager-plugin processes, including the clean up if things end unexpectedly.
7
+
8
+ [ssh server] <--> [PortForwardingSession] <--> [PortForwardingSessionProcess] <--> [session-manager-plugin]
9
+ [ssh server] --> [PortForwardingSession] will send commands to us across many threads
10
+ [PortForwardingSession] --> [manager] is a pipe for ipc
11
+ [proxy] --> [session-manager-plugin] is a local socket on internal_port
12
+ [manager] <-/-> [session-manager-plugin] doesnt connect to session-manager-plugin, it just starts and stops them
13
+ [manager] --> [proxy] monitors the pipe and the process to see if cleanup is needed
14
+
15
+ """
16
+
17
+ import socket
18
+ from typing import Dict, Any, List
19
+ import time
20
+ import subprocess
21
+ import multiprocessing
22
+ import threading
23
+ from multiprocessing.connection import PipeConnection
24
+
25
+ from ssm_cli.aws import aws_client
26
+ from ssm_cli.cli_args import ARGS
27
+ from ssm_cli.config import CONFIG
28
+ from ssm_cli.instances import Instance, SessionManagerPluginError, SessionManagerPluginPortError
29
+
30
+ import logging
31
+
32
+ logger = logging.getLogger(__name__)
33
+
34
+ #######
35
+ # Proxy Side
36
+ # Names should reflect both sides
37
+ #######
38
+
39
+ class PortForwardingSession():
40
+ """
41
+ This is the object that the user will interact with
42
+ """
43
+ def __init__(self, manager: "PortForwarding", session_id: str, internal_port: int):
44
+ self.manager = manager
45
+ self.session_id = session_id
46
+ self.internal_port = internal_port
47
+
48
+ def is_open(self):
49
+ return self.manager.is_open(self)
50
+
51
+ def close(self):
52
+ self.manager.close_session(self)
53
+
54
+ class PortForwarding():
55
+ """
56
+ Starts the manager process and creates the session objects.
57
+ """
58
+ session_cache: Dict[tuple, PortForwardingSession] # (host, port) -> session
59
+ instance: Instance
60
+ proxy_pipe: PipeConnection
61
+ pipe_lock: threading.Lock
62
+ manager_process: multiprocessing.Process
63
+
64
+ def __init__(self, instance):
65
+ self.session_cache = {}
66
+ self.instance = instance
67
+
68
+ def start(self):
69
+ logger.debug("starting manager process")
70
+ self.proxy_pipe, manager_pipe = multiprocessing.Pipe()
71
+ self.manager_process = PortForwardingManagerProcess(manager_pipe, self.instance)
72
+ self.manager_process.start()
73
+
74
+ logger.debug("waiting for manager process to start")
75
+ self.pipe_lock = threading.Lock()
76
+ ready = self.recv()
77
+ if ready != "ready":
78
+ logger.error(f"manager process failed to start: 'ready' != '{ready}'")
79
+ self.manager_process.terminate()
80
+ self.proxy_pipe.close()
81
+ raise Exception("Manager process failed to start")
82
+
83
+ def open_session(self, host: str, remote_port: int) -> PortForwardingSession:
84
+ session = self.session_cache.get((host, remote_port))
85
+ if session is not None:
86
+ if self.is_open(session):
87
+ logger.debug(f"{session.session_id} still running, reusing")
88
+ return session
89
+ else:
90
+ logger.debug(f"{session.session_id} closed, opening new one")
91
+ del self.session_cache[(host, remote_port)]
92
+
93
+ session_id, internal_port = self.send_recv(("open_session", (host, remote_port)))
94
+ session = PortForwardingSession(self, session_id, internal_port)
95
+ self.session_cache[(host, remote_port)] = session
96
+ return session
97
+
98
+ def close_session(self, session: PortForwardingSession):
99
+ self.proxy_pipe.send(("close_session", session.session_id))
100
+ del self.session_cache[(session.host, session.remote_port)]
101
+
102
+ def is_open(self, session: PortForwardingSession):
103
+ return self.send_recv(("is_open", session.session_id))
104
+
105
+ def is_alive(self):
106
+ return self.manager_process.is_alive()
107
+
108
+ def send(self, obj: Any):
109
+ if not self.is_alive():
110
+ raise Exception("Manager process has ended")
111
+ with self.pipe_lock:
112
+ logger.debug(f"sending {obj} to manager")
113
+ self.proxy_pipe.send(obj)
114
+
115
+ def recv(self) -> Any:
116
+ if not self.is_alive():
117
+ raise Exception("Manager process has ended")
118
+ with self.pipe_lock:
119
+ result = self.proxy_pipe.recv()
120
+ logger.debug(f"received {result} from manager")
121
+ return result
122
+
123
+ def send_recv(self, obj: Any) -> Any:
124
+ if not self.is_alive():
125
+ raise Exception("Manager process has ended")
126
+ with self.pipe_lock:
127
+ logger.debug(f"sending {obj} to manager")
128
+ self.proxy_pipe.send(obj)
129
+ result = self.proxy_pipe.recv()
130
+ logger.debug(f"received {result} from manager")
131
+ return result
132
+
133
+ #######
134
+ # Manager Side
135
+ # class suffix should be Process
136
+ #######
137
+
138
+ class PortForwardingSessionProcess():
139
+ """
140
+ Exposes open/close methods for the sessions, this will send messages to the manager process to do the actual work.
141
+ """
142
+ _open: bool = False
143
+ host: str
144
+ remote_port: int
145
+ internal_port: int
146
+ session_id: str
147
+ proc: subprocess.Popen
148
+ instance: Instance
149
+
150
+ def __init__(self, host: str, remote_port: int, instance: Instance):
151
+ self.host = host
152
+ self.remote_port = remote_port
153
+ self.instance = instance
154
+
155
+ def open(self):
156
+ # Retry because of rare race condition from get_free_port
157
+ for attempt in range(3):
158
+ try:
159
+ self.internal_port = get_free_port()
160
+ logger.debug(f"getting session for localhost:{self.internal_port} to {self.host}:{self.remote_port} over {self.instance.id} from aws")
161
+ session_id, proc = self.instance.start_port_forwarding_session_to_remote_host(self.host, self.remote_port, self.internal_port)
162
+ self.session_id = session_id
163
+ self.proc = proc
164
+ self._open = True
165
+ logger.info(f"{self.session_id} opened to {self.host}:{self.remote_port} on 127.0.0.1:{self.internal_port}, pid {self.proc.pid}")
166
+ return
167
+ except SessionManagerPluginPortError as e:
168
+ logger.warning(f"session-manager-plugin attempt {attempt} failed due to port clash, retrying: {e}")
169
+ time.sleep(0.1)
170
+
171
+ logger.error(f"session-manager-plugin failed to open session to {self.host}:{self.remote_port} after {attempt} attempts")
172
+ raise SessionManagerPluginError("Max retries hit") from e
173
+
174
+ def close(self):
175
+ logger.info(f"{self.session_id} closing")
176
+ self._open = False
177
+ try:
178
+ logger.debug(f"{self.session_id} killing process {self.proc.pid}")
179
+ self.proc.terminate()
180
+ self.proc.wait()
181
+ except Exception as e:
182
+ logger.error(f"{self.session_id} failed to kill process: {e}")
183
+
184
+ try:
185
+ logger.debug(f"{self.session_id} terminating session")
186
+ with aws_client('ssm') as client:
187
+ client.terminate_session(SessionId=self.session_id)
188
+ except Exception as e:
189
+ logger.error(f"{self.session_id} to terminate session: {e}")
190
+
191
+ def is_open(self):
192
+ if self.proc.poll() is not None:
193
+ logger.debug(f"process for {self.host}:{self.remote_port} has exited, restarting")
194
+ self._open = False
195
+ # the process ends when aws terminates the session, no need to check that
196
+ return self._open
197
+
198
+
199
+ class PortForwardingManagerProcess(multiprocessing.Process):
200
+ """
201
+ This is the process that will run the manager code, it is responsible for starting and stopping the session-manager-plugin processes.
202
+ As well as monitoring the proxy side and cleaning up.
203
+
204
+ This runs in a single thread, no need for all the locking and threading like the other class.
205
+ """
206
+
207
+ sessions: List[PortForwardingSessionProcess]
208
+ def __init__(self, pipe: PipeConnection, instance: Instance):
209
+ self.pipe = pipe
210
+ self.instance = instance
211
+ self.log_level = CONFIG.log.level
212
+ self.log_loggers = CONFIG.log.loggers
213
+ self.profile = ARGS.global_args.profile
214
+ super().__init__()
215
+
216
+ def run(self):
217
+ # Rebuild some global state that ssm_cli.cli normally deals with, the manager process is a clean slate,
218
+ # We can be selective about what we rebuild
219
+ # when config/args are refactored, this should be taken into account.
220
+ from ssm_cli.logging import setup_logging, configure_log_level
221
+ setup_logging("manager")
222
+ configure_log_level(self.log_level)
223
+ for logger_name, level in self.log_loggers.items():
224
+ configure_log_level(level, name=logger_name)
225
+
226
+ from ssm_cli.cli_args import ARGS, CliNamespace
227
+ ARGS.global_args = CliNamespace(profile=self.profile)
228
+ self.sessions = []
229
+ self.pipe.send("ready")
230
+ while True:
231
+ try:
232
+ msg, data = self.pipe.recv()
233
+ logger.debug(f"received {msg} {data} from proxy")
234
+ reply = None
235
+ if msg == "open_session":
236
+ reply = self.open_session(*data)
237
+ elif msg == "close_session":
238
+ self.close_session(data)
239
+ elif msg == "is_open":
240
+ reply = self.is_open(data)
241
+
242
+ if reply is not None:
243
+ logger.debug(f"sending {reply} to proxy")
244
+ self.pipe.send(reply)
245
+ except EOFError:
246
+ break
247
+
248
+ def open_session(self, host: str, remote_port: int):
249
+ session = None
250
+ for s in self.sessions:
251
+ if s.host == host and s.remote_port == remote_port:
252
+ logger.debug(f"{s.session_id} in cache for {host}:{remote_port}")
253
+ if s.is_open():
254
+ logger.debug(f"{s.session_id} still running, reusing")
255
+ return s
256
+ else:
257
+ logger.debug(f"{s.session_id} closed, removing from cache")
258
+ self.sessions.remove(s)
259
+
260
+ session = PortForwardingSessionProcess(host, remote_port, self.instance)
261
+ session.open()
262
+ self.sessions.append(session)
263
+ return session.session_id, session.internal_port
264
+
265
+ def close_session(self, session_id: str):
266
+ for session in self.sessions:
267
+ if session.session_id == session_id:
268
+ session.close()
269
+ self.sessions.remove(session)
270
+
271
+ def is_open(self, session_id: str) -> bool:
272
+ for session in self.sessions:
273
+ if session.session_id == session_id:
274
+ if session.is_open():
275
+ return True
276
+ else:
277
+ self.sessions.remove(session)
278
+ break
279
+ return False
280
+
281
+ def get_free_port(bind_host="127.0.0.1"):
282
+ """
283
+ Ask OS for an ephemeral free port. Returns the port number, however it is not guaranteed that the port will remain free. A retry should be used.
284
+ """
285
+ s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
286
+ s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
287
+ s.bind((bind_host, 0))
288
+ port = s.getsockname()[1]
289
+ s.close()
290
+ return port
@@ -8,8 +8,9 @@ import select
8
8
  from ssm_cli.commands.ssh_proxy.socket import StdIoSocket
9
9
  from ssm_cli.commands.ssh_proxy.shell import ShellThread
10
10
  from ssm_cli.commands.ssh_proxy.channels import Channels
11
+ from ssm_cli.commands.ssh_proxy.forwarding import PortForwarding, PortForwardingSession
11
12
  from ssm_cli.xdg import get_ssh_hostkey
12
- from ssm_cli.instances import Instance, SessionManagerPluginError, SessionManagerPluginPortError
13
+ from ssm_cli.instances import Instance
13
14
 
14
15
  import logging
15
16
  logger = logging.getLogger(__name__)
@@ -20,16 +21,16 @@ class SshServer(paramiko.ServerInterface):
20
21
  """
21
22
  event: threading.Event
22
23
  instance: Instance
23
- connections: Dict[tuple, tuple] # (host, port) -> (internal_port, session)
24
24
  channels: Channels
25
25
  transport: paramiko.Transport
26
+ port_forwarding_manager: PortForwarding
26
27
 
27
28
  def __init__(self, instance: Instance):
28
29
  logger.debug("creating server")
29
30
  self.event = threading.Event()
30
31
  self.instance = instance
31
- self.connections = {}
32
-
32
+ self.port_forwarding_manager = PortForwarding(instance)
33
+
33
34
  def start(self):
34
35
  logger.info("starting server")
35
36
 
@@ -40,8 +41,9 @@ class SshServer(paramiko.ServerInterface):
40
41
  key_path = get_ssh_hostkey()
41
42
  host_key = paramiko.RSAKey(filename=key_path)
42
43
  logger.info("Loaded existing host key")
43
-
44
44
  self.transport.add_server_key(host_key)
45
+
46
+ self.port_forwarding_manager.start()
45
47
  self.transport.start_server(server=self)
46
48
 
47
49
  self.event.wait()
@@ -83,53 +85,33 @@ class SshServer(paramiko.ServerInterface):
83
85
  logger.info(f"direct TCP/IP request: chan={chanid} origin={origin} destination={destination}")
84
86
  host = destination[0]
85
87
  remote_port = destination[1]
86
-
87
- internal_port, session = None, None
88
-
89
- if (host, remote_port) in self.connections:
90
- internal_port, session = self.connections[(host, remote_port)]
91
- if session.proc.poll() is not None:
92
- logger.debug(f"process for {host}:{remote_port} has exited, restarting")
93
- del self.connections[(host, remote_port)]
94
- internal_port, session = None, None
95
- else:
96
- logger.debug(f"reusing existing connection to {host}:{remote_port} with internal port {internal_port} and process {session.proc.pid}")
97
88
 
98
- if internal_port is None:
99
- # Retry because of rare race condition from get_free_port
100
- for attempt in range(3):
101
- try:
102
- internal_port = get_free_port()
103
- session = self.instance.start_port_forwarding_session_to_remote_host(host, remote_port, internal_port)
104
- self.connections[(host, remote_port)] = (internal_port, session)
105
- break
106
- except SessionManagerPluginPortError as e:
107
- logger.warning(f"session-manager-plugin attempt {attempt} failed due to port clash, retrying: {e}")
108
- time.sleep(0.1)
109
- except SessionManagerPluginError as e:
110
- logger.error(f"session-manager-plugin failed: {e}")
111
- return paramiko.OPEN_FAILED_CONNECT_FAILED
112
-
113
- logger.debug(f"connecting to session manager plugin on 127.0.0.1:{internal_port}")
89
+ try:
90
+ session = self.port_forwarding_manager.open_session(host, remote_port)
91
+ except Exception as e:
92
+ logger.error(f"failed to connect: {e}")
93
+ return paramiko.OPEN_FAILED_CONNECT_FAILED
94
+
95
+ logger.debug(f"connecting to session manager plugin on 127.0.0.1:{session.internal_port}")
114
96
  # Even though we wait for the process to say its connected, we STILL need to wait for it
115
97
  for attempt in range(10):
116
98
  try:
117
- if session.proc.poll() is not None:
118
- logger.error(f"session-manager-plugin has exited")
119
- raise RuntimeError("session-manager-plugin has exited")
120
- sock = socket.create_connection(('127.0.0.1', internal_port))
121
- logger.info(f"connected to 127.0.0.1:{internal_port}")
99
+ sock = socket.create_connection(('127.0.0.1', session.internal_port))
100
+ logger.info(f"connected to 127.0.0.1:{session.internal_port}")
122
101
  break
123
102
  except Exception as e:
124
103
  logger.warning(f"connection attempt {attempt} failed: {e}")
125
104
  time.sleep(0.1)
126
-
127
- if not sock:
105
+ if not session.is_open():
106
+ logger.error(f"session is closed")
107
+ return paramiko.OPEN_FAILED_CONNECT_FAILED
108
+ else:
128
109
  logger.error("max retries reached, giving up")
129
110
  return paramiko.OPEN_FAILED_CONNECT_FAILED
130
111
 
131
112
 
132
113
  chunk_size = 1024
114
+
133
115
  # Start thread to open the channel and forward data
134
116
  def forwarding_thread():
135
117
  logger.info(f"starting forward thread chan={chanid}")
@@ -148,9 +130,9 @@ class SshServer(paramiko.ServerInterface):
148
130
  if len(data) == 0:
149
131
  break
150
132
  sock.send(data)
151
-
152
- del session
153
- logger.info(f"forward thread chan={self.chanid} exiting")
133
+
134
+ session.close()
135
+ logger.info(f"forward thread chan={chanid} exiting")
154
136
 
155
137
  t = threading.Thread(target=forwarding_thread)
156
138
  t.start()
@@ -160,15 +142,3 @@ class SshServer(paramiko.ServerInterface):
160
142
 
161
143
  def get_banner(self):
162
144
  return ("SSM CLI - SSH server\r\n", "en-US")
163
-
164
-
165
- def get_free_port(bind_host="127.0.0.1"):
166
- """
167
- Ask OS for an ephemeral free port. Returns the port number, however it is not guaranteed that the port will remain free. A retry should be used.
168
- """
169
- s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
170
- s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
171
- s.bind((bind_host, 0))
172
- port = s.getsockname()[1]
173
- s.close()
174
- return port
@@ -15,7 +15,8 @@ class ActionsConfig:
15
15
  class LoggingConfig:
16
16
  level: str = "info"
17
17
  loggers: Dict[str, str] = {
18
- "botocore": "warn"
18
+ "botocore": "warn",
19
+ "paramiko": "warn"
19
20
  }
20
21
  """key value dictionary to override log level on, some modules make a lot of noise, botocore for example"""
21
22
 
@@ -4,7 +4,7 @@ import re
4
4
  import signal
5
5
  import subprocess
6
6
  import sys
7
- from typing import List
7
+ from typing import List, Tuple
8
8
 
9
9
  from ssm_cli.selectors import SELECTORS
10
10
  from ssm_cli.config import CONFIG
@@ -17,7 +17,7 @@ class SessionManagerPluginError(Exception):
17
17
  """ A generic exception for any AWS errors """
18
18
  stdout = []
19
19
  returncode = 0
20
- def __init__(self, message, stdout, returncode):
20
+ def __init__(self, message, stdout = None, returncode = None):
21
21
  super().__init__(message)
22
22
  self.stdout = stdout
23
23
  self.returncode = returncode
@@ -29,27 +29,6 @@ class SessionManagerPluginPortError(SessionManagerPluginError):
29
29
  """ A specific exception for a timeout error """
30
30
  pass
31
31
 
32
-
33
- class SessionManagerSession:
34
- """ A class to manage an SSM session, helps with cleanup """
35
- def __init__(self, session_id: str, proc: subprocess.Popen):
36
- self.session_id = session_id
37
- self.proc = proc
38
-
39
- def __del__(self):
40
- logger.info(f"cleaning up session {self.session_id}")
41
- try:
42
- self.proc.terminate()
43
- self.proc.wait()
44
- except Exception as e:
45
- logger.error(f"Failed to cleanup process {self.session_id} but continuing anyway: {e}")
46
- try:
47
- with aws_client('ssm') as client:
48
- client.terminate_session(SessionId=self.session_id)
49
- except Exception as e:
50
- logger.error(f"Failed to cleanup session {self.session_id} but continuing anyway: {e}")
51
-
52
-
53
32
  class Instance:
54
33
  """
55
34
  Contains information about an EC2 instance and methods to handle sessions with them.
@@ -89,7 +68,18 @@ class Instance:
89
68
  logger.error(f"Failed to connect to session: {result.stderr.decode()}")
90
69
  raise RuntimeError(f"Failed to connect to session: {result.stderr.decode()}")
91
70
 
92
- def start_port_forwarding_session_to_remote_host(self, host: str, remote_port: int, internal_port: int) -> SessionManagerSession:
71
+ def start_port_forwarding_session_to_remote_host(self, host: str, remote_port: int, internal_port: int) -> Tuple[str, subprocess.Popen]:
72
+ """
73
+ Start a port forwarding session to a remote host.
74
+
75
+ Args:
76
+ host: The remote host to forward to.
77
+ remote_port: The remote port to forward to.
78
+ internal_port: The local port use for forwarding.
79
+
80
+ Returns:
81
+ A tuple containing the session id and the subprocess.
82
+ """
93
83
  logger.debug(f"start port forwarding between localhost:{internal_port} and {host}:{remote_port} via {self.id}")
94
84
  with aws_session(False) as session:
95
85
  client = session.client('ssm') # we need a fresh connection where session is the same as the client
@@ -124,18 +114,18 @@ class Instance:
124
114
  stdout=subprocess.PIPE,
125
115
  stderr=subprocess.STDOUT
126
116
  )
127
- session = SessionManagerSession(response["SessionId"], proc)
128
117
 
129
118
  # changes needed here:
130
- # SessionManagerPluginPortError needs to be raised correctly
119
+ # SessionManagerPluginPortError needs to be raised when it applies
131
120
  # add a timeout for waiting on stdout
132
121
  # maybe split up stdout/stderr?
122
+ # validate the port in stdout matches the internal port
133
123
  output = b''
134
124
  for line in proc.stdout:
135
125
  line = line.strip()
136
126
  output += line + b'\n'
137
127
  if line == b'Waiting for connections...':
138
- return session
128
+ return (response["SessionId"], proc)
139
129
  else:
140
130
  logger.debug(f"Recieved from session-manager-plugin: {line}")
141
131
 
@@ -5,11 +5,11 @@ from typing import Optional
5
5
  from ssm_cli.xdg import get_log_file, get_all_log_files
6
6
 
7
7
 
8
- def setup_logging():
8
+ def setup_logging(name:str = "cli"):
9
9
  """Set up basic logging configuration with date-based rotation."""
10
10
  logging.basicConfig(
11
11
  level=logging.WARNING,
12
- filename=get_log_file(),
12
+ filename=get_log_file(name),
13
13
  filemode='a',
14
14
  format='%(asctime)s - %(process)d [%(threadName)s] - %(name)s - %(levelname)s - %(message)s',
15
15
  datefmt='%Y-%m-%d %H:%M:%S'
@@ -26,15 +26,13 @@ def get_conf_file(check=True) -> Path:
26
26
  raise EnvironmentError(f"{path} missing, run `ssm setup` to create")
27
27
  return path
28
28
 
29
- def get_log_file(date: Optional[datetime] = None) -> Path:
30
- if date is None:
31
- date = datetime.now()
32
- date_str = date.strftime('%Y-%m-%d')
33
- path = get_state_root() / f"run.{date_str}.log"
29
+ def get_log_file(name="cli") -> Path:
30
+ date_str = datetime.now().strftime('%Y-%m-%d')
31
+ path = get_state_root() / f"ssm-{name}.{date_str}.log"
34
32
  return path
35
33
 
36
34
  def get_all_log_files() -> str:
37
- return get_state_root().glob('run.*.log')
35
+ return get_state_root().glob('ssm-*.*.log')
38
36
 
39
37
  def get_ssh_hostkey(check=True) -> Path:
40
38
  path = get_conf_root(check) / 'hostkey.pem'
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ssm-cli
3
- Version: 1.0.0.dev0
3
+ Version: 1.0.0.dev3
4
4
  Summary: CLI tool to help with SSM functionality, aimed at adminstrators
5
5
  Author-email: Simon Fletcher <simon.fletcher@lexisnexisrisk.com>
6
6
  License: MIT License
@@ -24,6 +24,7 @@ ssm_cli/commands/setup.py
24
24
  ssm_cli/commands/shell.py
25
25
  ssm_cli/commands/ssh_proxy/__init__.py
26
26
  ssm_cli/commands/ssh_proxy/channels.py
27
+ ssm_cli/commands/ssh_proxy/forwarding.py
27
28
  ssm_cli/commands/ssh_proxy/server.py
28
29
  ssm_cli/commands/ssh_proxy/shell.py
29
30
  ssm_cli/commands/ssh_proxy/socket.py
File without changes
File without changes
File without changes