diff --git a/changelog.md b/changelog.md index a1a14bdf..d7618da0 100644 --- a/changelog.md +++ b/changelog.md @@ -6,6 +6,11 @@ Breaking Changes * Remove undocumented support for `tcp://` and `socket://` DSNs. +Features +--------- +* Add `--ssh-jump` CLI argument, restoring SSH jump functionality. + + Documentation --------- * Fix a typo in myclirc commentary. diff --git a/mycli/cli_runner.py b/mycli/cli_runner.py index 5d71f5e2..585860fc 100644 --- a/mycli/cli_runner.py +++ b/mycli/cli_runner.py @@ -311,45 +311,48 @@ def run_from_cli_args(cli_args: 'CliArgs', client_factory: ClientFactory) -> Non use_keyring = str_to_bool(cli_args.use_keyring) reset_keyring = False - mycli.connect( - database=database, - user=cli_args.user, - passwd=cli_args.password, - host=cli_args.host, - port=cli_args.port, - socket=cli_args.socket, - local_infile=cli_args.local_infile, - ssl=ssl, - init_command=combined_init_cmd, - unbuffered=cli_args.unbuffered, - character_set=cli_args.character_set, - use_keyring=use_keyring, - reset_keyring=reset_keyring, - keepalive_ticks=keepalive_ticks, - ) - - if combined_init_cmd: - click.echo(f"Executing init-command: {combined_init_cmd}", err=True) - - mycli.logger.debug( - "Launch Params: \n\tdatabase: %r\tuser: %r\thost: %r\tport: %r", - database, - cli_args.user, - cli_args.host, - cli_args.port, - ) + try: + mycli.connect( + database=database, + user=cli_args.user, + passwd=cli_args.password, + host=cli_args.host, + port=cli_args.port, + socket=cli_args.socket, + local_infile=cli_args.local_infile, + ssl=ssl, + init_command=combined_init_cmd, + unbuffered=cli_args.unbuffered, + character_set=cli_args.character_set, + use_keyring=use_keyring, + reset_keyring=reset_keyring, + keepalive_ticks=keepalive_ticks, + ssh_jump=cli_args.ssh_jump, + ) + + if combined_init_cmd: + click.echo(f"Executing init-command: {combined_init_cmd}", err=True) + + mycli.logger.debug( + "Launch Params: \n\tdatabase: %r\tuser: %r\thost: %r\tport: %r", + database, + cli_args.user, + cli_args.host, + cli_args.port, + ) - if cli_args.execute is not None: - sys.exit(main_execute_from_cli(mycli, cli_args)) + if cli_args.execute is not None: + sys.exit(main_execute_from_cli(mycli, cli_args)) - if cli_args.batch is not None and cli_args.batch != '-' and cli_args.progress and sys.stderr.isatty(): - sys.exit(main_batch_with_progress_bar(mycli, cli_args)) + if cli_args.batch is not None and cli_args.batch != '-' and cli_args.progress and sys.stderr.isatty(): + sys.exit(main_batch_with_progress_bar(mycli, cli_args)) - if cli_args.batch is not None: - sys.exit(main_batch_without_progress_bar(mycli, cli_args)) + if cli_args.batch is not None: + sys.exit(main_batch_without_progress_bar(mycli, cli_args)) - if not sys.stdin.isatty(): - sys.exit(main_batch_from_stdin(mycli, cli_args)) + if not sys.stdin.isatty(): + sys.exit(main_batch_from_stdin(mycli, cli_args)) - mycli.run_cli() - mycli.close() + mycli.run_cli() + finally: + mycli.close() diff --git a/mycli/client.py b/mycli/client.py index 320d2fa8..9fcadca4 100644 --- a/mycli/client.py +++ b/mycli/client.py @@ -40,6 +40,7 @@ from mycli.schema_prefetcher import SchemaPrefetcher from mycli.sqlcompleter import SQLCompleter from mycli.sqlexecute import SQLExecute +from mycli.ssh_tunnel import SshTunnel from mycli.types import Query sqlparse.engine.grouping.MAX_GROUPING_DEPTH = None # type: ignore[assignment] @@ -75,6 +76,7 @@ def __init__( cli_verbosity: int = 0, ) -> None: self.sqlexecute = sqlexecute + self.ssh_tunnel: SshTunnel | None = None self.logfile = logfile self.login_path = login_path self.toolbar_error_message: str | None = None @@ -198,10 +200,20 @@ def __init__( special.set_destructive_keywords(self.destructive_keywords) def close(self) -> None: - if hasattr(self, 'schema_prefetcher'): + try: self.schema_prefetcher.stop() + except Exception: + pass if self.sqlexecute is not None: - self.sqlexecute.close() + try: + self.sqlexecute.close() + except Exception: + pass + if self.ssh_tunnel is not None: + try: + self.ssh_tunnel.close() + except Exception: + pass def run_cli(self) -> None: repl_package.main_repl(self) diff --git a/mycli/client_connection.py b/mycli/client_connection.py index 4625ad55..96309ff7 100644 --- a/mycli/client_connection.py +++ b/mycli/client_connection.py @@ -4,6 +4,7 @@ import sys import traceback from typing import TYPE_CHECKING, Any +from urllib.parse import quote as urlquote import click import keyring @@ -22,6 +23,7 @@ ) from mycli.packages.filepaths import guess_socket_location from mycli.sqlexecute import SQLExecute +from mycli.ssh_tunnel import SshTunnel, SshTunnelError try: from pwd import getpwuid @@ -58,6 +60,7 @@ def connect( use_keyring: bool | None = None, reset_keyring: bool | None = None, keepalive_ticks: int | None = None, + ssh_jump: str | None = None, ) -> None: mylogin_cnf: dict[str, Any] = self.read_mylogin_cnf(self.mylogin_cnf) # Fall back to .mylogin.cnf values only if user did not specify a value. @@ -67,6 +70,7 @@ def connect( ssl_config: dict[str, Any] = ssl or {} user_connection_config = self.config_without_package_defaults.get('connection', {}) self.keepalive_ticks = keepalive_ticks + self.ssh_tunnel = None int_port = port and int(port) if not int_port: @@ -74,6 +78,31 @@ def connect( if not host or host == DEFAULT_HOST: socket = socket or user_connection_config.get("default_socket") or mylogin_cnf["socket"] or guess_socket_location() + if ssh_jump: + remote_host = host or DEFAULT_HOST + remote_port = int_port or DEFAULT_PORT + remote_socket = socket or None + ssh_executable = self.config.get('ssh', {}).get('ssh_executable', 'ssh') or 'ssh' + ssh_options = self.config.get('ssh', {}).get('ssh_options') + try: + self.ssh_tunnel = SshTunnel.from_target( + ssh_jump, + remote_host=remote_host, + remote_port=int(remote_port), + remote_socket=remote_socket, + ssh_executable=ssh_executable, + ssh_options=ssh_options, + ) + self.ssh_tunnel.start() + except (OSError, ValueError, SshTunnelError) as exc: + click.secho(f'Error: Unable to start SSH tunnel: {exc}', err=True, fg='red') + try: + if self.ssh_tunnel: + self.ssh_tunnel.close() + except Exception: + pass + sys.exit(1) + passwd = passwd if isinstance(passwd, (str, int)) else mylogin_cnf["password"] if not character_set: @@ -135,7 +164,10 @@ def connect( # 5. .mylogin.cnf # 6. keyring - keyring_identifier = f'{user}@{host}:{"" if socket else int_port}:{socket or ""}' + ssh_tunnel_field = urlquote(ssh_jump or '') + if ssh_tunnel_field: + ssh_tunnel_field = ':' + ssh_tunnel_field + keyring_identifier = f'{user}@{host}:{"" if socket else int_port}:{socket or ""}{ssh_tunnel_field}' keyring_domain = 'mycli.net' keyring_retrieved_cleanly = False @@ -152,19 +184,24 @@ def connect( # should not fail, but will help the typechecker assert not isinstance(passwd, int) - connection_info: dict[Any, Any] = { - "database": database, - "user": user, - "password": passwd, - "host": host, - "port": int_port, - "socket": socket, - "character_set": character_set, - "local_infile": use_local_infile, - "ssl": ssl_config, - "init_command": init_command, - "unbuffered": unbuffered, + connection_info: dict[str, Any] = { + 'database': database, + 'user': user, + 'password': passwd, + 'character_set': character_set, + 'local_infile': use_local_infile, + 'ssl': ssl_config, + 'init_command': init_command, + 'unbuffered': unbuffered, } + if self.ssh_tunnel: + connection_info['host'] = self.ssh_tunnel.local_host + connection_info['port'] = self.ssh_tunnel.local_port + connection_info['socket'] = None + else: + connection_info['host'] = host + connection_info['port'] = int_port + connection_info['socket'] = socket def _update_keyring(password: str | None, keyring_retrieved_cleanly: bool): if not password: @@ -239,6 +276,8 @@ def _connect( socket_owner = getpwuid(os.stat(socket).st_uid).pw_name except KeyError: socket_owner = '' + except FileNotFoundError: + socket_owner = '' self.echo(f"Connecting to socket {socket}, owned by user {socket_owner}", err=True) try: _connect(keyring_retrieved_cleanly=keyring_retrieved_cleanly) diff --git a/mycli/main.py b/mycli/main.py index e8f615ea..66ac5a05 100755 --- a/mycli/main.py +++ b/mycli/main.py @@ -271,6 +271,10 @@ class CliArgs: type=int, help='Send regular keepalive pings to the connection, roughly every seconds.', ) + ssh_jump: str | None = clickdc.option( + type=str, + help='Open an SSH tunnel via [user@]host[:port] and connect to MySQL through it.', + ) checkup: bool = clickdc.option( is_flag=True, help='Run a checkup on your configuration.', @@ -287,47 +291,47 @@ class CliArgs: ssh_user: str | None = clickdc.option( type=str, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) ssh_host: str | None = clickdc.option( type=str, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) ssh_port: int = clickdc.option( type=int, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) ssh_password: str | None = clickdc.option( type=str, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) ssh_key_filename: str | None = clickdc.option( type=str, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) ssh_config_path: str = clickdc.option( type=str, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) ssh_config_host: str | None = clickdc.option( type=str, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) list_ssh_config: bool = clickdc.option( is_flag=True, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) ssh_warning_off: bool = clickdc.option( is_flag=True, hidden=True, - deprecated='No effect. See https://github.com/dbcli/mycli/issues/1960 .', + deprecated='No effect. See --ssh-jump.', ) diff --git a/mycli/myclirc b/mycli/myclirc index ff21f6ae..99ff0244 100644 --- a/mycli/myclirc +++ b/mycli/myclirc @@ -403,3 +403,10 @@ matching-bracket.other = '#000000 bg:#aacccc' [alias_dsn.init-commands] # Define one or more SQL statements per alias (semicolon-separated). # example_dsn = "SET sql_select_limit=1000; SET time_zone='+00:00'" + +[ssh] +# Path to the ssh executable used by --ssh-jump. +ssh_executable = ssh + +# options to pass to the ssh executable when making a tunnel +ssh_options = -a -o ServerAliveInterval=60 -o ExitOnForwardFailure=yes -o IPQoS=af13 -o LogLevel=FATAL diff --git a/mycli/ssh_tunnel.py b/mycli/ssh_tunnel.py new file mode 100644 index 00000000..cd9b3b8c --- /dev/null +++ b/mycli/ssh_tunnel.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from dataclasses import dataclass +import shlex +import socket +import subprocess +import threading +import time + + +class SshTunnelError(RuntimeError): + pass + + +def _find_free_local_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(('127.0.0.1', 0)) + return int(sock.getsockname()[1]) + + +@dataclass(slots=True) +class SshTunnelTarget: + ssh_target: str + ssh_port: int | None = None + + @classmethod + def parse(cls, value: str) -> SshTunnelTarget: + target, separator, port = value.rpartition(':') + if separator and port.isdigit() and ']' not in port: + return cls(target, int(port)) + return cls(value, None) + + +class SshTunnel: + def __init__( + self, + *, + ssh_target: str, + remote_host: str, + remote_port: int, + remote_socket: str | None = None, + ssh_executable: str = 'ssh', + ssh_options: str | None = None, + ssh_port: int | None = None, + local_port: int | None = None, + ready_timeout: float = 30.0, + ) -> None: + self.ssh_executable = ssh_executable + self.ssh_options = ssh_options + self.ssh_target = ssh_target + self.ssh_port = ssh_port + self.remote_host = remote_host + self.remote_port = remote_port + self.remote_socket = remote_socket + self.local_host = '127.0.0.1' + self.local_port = local_port or _find_free_local_port() + self.ready_timeout = ready_timeout + self.process: subprocess.Popen | None = None + self._startup_error: OSError | None = None + self._ready = threading.Event() + self._failed = threading.Event() + self._thread: threading.Thread | None = None + + @classmethod + def from_target( + cls, + ssh_jump_spec: str, + *, + remote_host: str, + remote_port: int, + remote_socket: str | None = None, + ssh_executable: str = 'ssh', + ssh_options: str | None = None, + ) -> SshTunnel: + target = SshTunnelTarget.parse(ssh_jump_spec) + return cls( + ssh_target=target.ssh_target, + ssh_port=target.ssh_port, + remote_host=remote_host, + remote_port=remote_port, + remote_socket=remote_socket, + ssh_executable=ssh_executable, + ssh_options=ssh_options, + ) + + def _forward_spec(self) -> str: + if self.remote_socket: + return f'{self.local_host}:{self.local_port}:{self.remote_socket}' + return f'{self.local_host}:{self.local_port}:{self.remote_host}:{self.remote_port}' + + def command(self) -> list[str]: + opts = shlex.split(self.ssh_options or '') + command = [ + self.ssh_executable, + *opts, + '-N', + '-L', + self._forward_spec(), + ] + if self.ssh_port is not None: + command.extend(['-p', str(self.ssh_port)]) + command.append(self.ssh_target) + return command + + def start(self) -> None: + self._thread = threading.Thread(target=self._run, name='mycli-ssh-tunnel', daemon=True) + self._thread.start() + deadline = time.monotonic() + self.ready_timeout + while time.monotonic() < deadline: + if self._failed.is_set(): + self.close() + if self._startup_error is not None: + raise SshTunnelError(f'Unable to start SSH tunnel process: {self._startup_error}') from self._startup_error + raise SshTunnelError('SSH tunnel process exited before it was ready.') + if self._is_listening(): + self._ready.set() + return + time.sleep(0.05) + self.close() + raise SshTunnelError('Timed out waiting for SSH tunnel to become ready.') + + def close(self) -> None: + process = self.process + if process is not None and process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + if self._thread is not None and self._thread.is_alive(): + self._thread.join(timeout=5) + + def _run(self) -> None: + try: + self.process = subprocess.Popen( + self.command(), + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + except OSError as exc: + self._startup_error = exc + self._failed.set() + return + return_code = self.process.wait() + if return_code != 0 and not self._ready.is_set(): + self._failed.set() + + def _is_listening(self) -> bool: + try: + with socket.create_connection((self.local_host, self.local_port), timeout=0.05): + return True + except OSError: + return False diff --git a/test/myclirc b/test/myclirc index 45dafe32..9795ba64 100644 --- a/test/myclirc +++ b/test/myclirc @@ -406,3 +406,10 @@ global_limit = set sql_select_limit=9999 [alias_dsn.init-commands] # Define one or more SQL statements per alias (semicolon-separated). # example_dsn = "SET sql_select_limit=1000; SET time_zone='+00:00'" + +[ssh] +# Path to the ssh executable used by --ssh-jump. +ssh_executable = ssh + +# options to pass to the ssh executable when making a tunnel +ssh_options = -a -o ServerAliveInterval=60 -o ExitOnForwardFailure=yes -o IPQoS=af13 -o LogLevel=FATAL diff --git a/test/pytests/test_cli_runner.py b/test/pytests/test_cli_runner.py index 4adba139..262c80c3 100644 --- a/test/pytests/test_cli_runner.py +++ b/test/pytests/test_cli_runner.py @@ -31,6 +31,7 @@ def __init__( self.ssl_mode: str | None = None self.logger = DummyLogger() self.dsn_alias: str | None = None + self.ssh_tunnel: Any = None self.connect_calls: list[dict[str, Any]] = [] self.run_cli_called = False self.close_called = False @@ -42,6 +43,8 @@ def run_cli(self) -> None: self.run_cli_called = True def close(self) -> None: + if getattr(self, 'ssh_tunnel', None) is not None: + self.ssh_tunnel.close() self.close_called = True diff --git a/test/pytests/test_client.py b/test/pytests/test_client.py index 537cccf8..eec4b225 100644 --- a/test/pytests/test_client.py +++ b/test/pytests/test_client.py @@ -3,7 +3,7 @@ from io import TextIOWrapper from pathlib import Path from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest @@ -133,13 +133,29 @@ def test_close_stops_schema_prefetcher_and_closes_sqlexecute() -> None: cli = MyCli.__new__(MyCli) stopped: list[bool] = [] closed: list[bool] = [] + tunnel_closed: list[bool] = [] cli.schema_prefetcher = SimpleNamespace(stop=lambda: stopped.append(True)) cli.sqlexecute = SimpleNamespace(close=lambda: closed.append(True)) # type: ignore[assignment] + cast(Any, cli).ssh_tunnel = SimpleNamespace(close=lambda: tunnel_closed.append(True)) MyCli.close(cli) assert stopped == [True] assert closed == [True] + assert tunnel_closed == [True] + + +def test_close_swallows_cleanup_errors() -> None: + cli = MyCli.__new__(MyCli) + + def fail() -> None: + raise RuntimeError('cleanup failed') + + cli.schema_prefetcher = SimpleNamespace(stop=fail) + cli.sqlexecute = SimpleNamespace(close=fail) # type: ignore[assignment] + cast(Any, cli).ssh_tunnel = SimpleNamespace(close=fail) + + MyCli.close(cli) def test_run_cli_delegates_to_main_repl(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/test/pytests/test_client_connection.py b/test/pytests/test_client_connection.py index 765b69c1..eddf09ca 100644 --- a/test/pytests/test_client_connection.py +++ b/test/pytests/test_client_connection.py @@ -307,6 +307,142 @@ def fail_set_password(domain: str, identifier: str, password: str) -> None: assert secho_calls == [('Password not saved to the system keyring: locked', {'err': True, 'fg': 'red'})] +def test_connect_uses_ssh_jump_with_remote_socket(monkeypatch: pytest.MonkeyPatch) -> None: + tunnel_calls: list[dict[str, Any]] = [] + + class FakeTunnel: + local_host = '127.0.0.1' + local_port = 4406 + remote_socket = '/var/run/mysqld/mysqld.sock' + + @classmethod + def from_target( + cls, + ssh_jump: str, + *, + remote_host: str, + remote_port: int, + remote_socket: str | None = None, + ssh_executable: str = 'ssh', + ssh_options: str | None = None, + ) -> 'FakeTunnel': + tunnel_calls.append({ + 'ssh_jump': ssh_jump, + 'remote_host': remote_host, + 'remote_port': remote_port, + 'remote_socket': remote_socket, + 'ssh_executable': ssh_executable, + 'ssh_options': None, + }) + return cls() + + def start(self) -> None: + pass + + def close(self) -> None: + pass + + monkeypatch.setattr(client_connection, 'SshTunnel', FakeTunnel) + client = DummyClient(config={'main': {}, 'ssh': {'ssh_executable': '/opt/bin/ssh'}, 'connection': {}}) + + client.connect( + user='alice', + host='db.internal', + port=None, + socket='/var/run/mysqld/mysqld.sock', + ssh_jump='bastion', + ) + + assert tunnel_calls == [ + { + 'ssh_jump': 'bastion', + 'remote_host': 'db.internal', + 'remote_port': 3306, + 'remote_socket': '/var/run/mysqld/mysqld.sock', + 'ssh_executable': '/opt/bin/ssh', + 'ssh_options': None, + } + ] + assert FakeSQLExecute.calls[-1]['host'] == '127.0.0.1' + assert FakeSQLExecute.calls[-1]['port'] == 4406 + assert FakeSQLExecute.calls[-1]['socket'] is None + + +def test_connect_reports_ssh_jump_start_error_and_closes_tunnel(monkeypatch: pytest.MonkeyPatch) -> None: + close_calls: list[bool] = [] + secho_calls: list[tuple[str, dict[str, Any]]] = [] + + class FakeTunnel: + local_host = '127.0.0.1' + local_port = 4406 + remote_socket = None + + @classmethod + def from_target( + cls, + _ssh_jump: str, + *, + remote_host: str, + remote_port: int, + remote_socket: str | None = None, + ssh_executable: str = 'ssh', + ssh_options: str | None = None, + ) -> 'FakeTunnel': + return cls() + + def start(self) -> None: + raise client_connection.SshTunnelError('no tunnel') + + def close(self) -> None: + close_calls.append(True) + + monkeypatch.setattr(client_connection, 'SshTunnel', FakeTunnel) + monkeypatch.setattr(client_connection.click, 'secho', lambda message, **kwargs: secho_calls.append((message, kwargs))) + client = DummyClient() + + with pytest.raises(SystemExit) as excinfo: + client.connect(host='db.internal', ssh_jump='bastion') + + assert excinfo.value.code == 1 + assert close_calls == [True] + assert secho_calls == [('Error: Unable to start SSH tunnel: no tunnel', {'err': True, 'fg': 'red'})] + assert FakeSQLExecute.calls == [] + + +def test_connect_swallows_ssh_jump_cleanup_error(monkeypatch: pytest.MonkeyPatch) -> None: + class FakeTunnel: + local_host = '127.0.0.1' + local_port = 4406 + remote_socket = None + + @classmethod + def from_target( + cls, + _ssh_jump: str, + *, + remote_host: str, + remote_port: int, + remote_socket: str | None = None, + ssh_executable: str = 'ssh', + ssh_options: str | None = None, + ) -> 'FakeTunnel': + return cls() + + def start(self) -> None: + raise OSError('no process') + + def close(self) -> None: + raise RuntimeError('close failed') + + monkeypatch.setattr(client_connection, 'SshTunnel', FakeTunnel) + client = DummyClient() + + with pytest.raises(SystemExit) as excinfo: + client.connect(host='db.internal', ssh_jump='bastion') + + assert excinfo.value.code == 1 + + def test_connect_retries_without_ssl_for_auto_handshake_error() -> None: client = DummyClient() FakeSQLExecute.effects = [op_error(client_connection.HANDSHAKE_ERROR), None] diff --git a/test/pytests/test_main.py b/test/pytests/test_main.py index 34138293..1d2aa1fc 100644 --- a/test/pytests/test_main.py +++ b/test/pytests/test_main.py @@ -922,6 +922,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, "MyCli", MockMyCli) @@ -1181,6 +1184,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1233,6 +1239,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1291,6 +1300,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1359,6 +1371,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1440,6 +1455,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1512,6 +1530,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1580,6 +1601,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1634,6 +1658,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1702,6 +1729,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) @@ -1768,6 +1798,9 @@ def connect(self, **args): def run_query(self, query, new_line=True): pass + def close(self): + pass + import mycli.main monkeypatch.setattr(mycli.main, 'MyCli', MockMyCli) diff --git a/test/pytests/test_ssh_tunnel.py b/test/pytests/test_ssh_tunnel.py new file mode 100644 index 00000000..eae8c465 --- /dev/null +++ b/test/pytests/test_ssh_tunnel.py @@ -0,0 +1,307 @@ +from __future__ import annotations + +import socket +import subprocess +from typing import Any, cast + +import pytest + +from mycli import ssh_tunnel +from mycli.ssh_tunnel import SshTunnel, SshTunnelError, SshTunnelTarget + + +def test_ssh_tunnel_target_parse_handles_target_with_port() -> None: + target = SshTunnelTarget.parse('alice@bastion:2222') + + assert target.ssh_target == 'alice@bastion' + assert target.ssh_port == 2222 + + +def test_ssh_tunnel_target_parse_handles_target_without_port() -> None: + target = SshTunnelTarget.parse('alice@bastion') + + assert target.ssh_target == 'alice@bastion' + assert target.ssh_port is None + + +def test_ssh_tunnel_command_includes_forward_and_ssh_port() -> None: + tunnel = SshTunnel( + ssh_target='alice@bastion', + ssh_port=2222, + remote_host='db.internal', + remote_port=3307, + local_port=4406, + ) + + assert tunnel.command()[0] == 'ssh' + assert '127.0.0.1:4406:db.internal:3307' in tunnel.command() + assert 'alice@bastion' in tunnel.command() + assert '2222' in tunnel.command() + + +def test_ssh_tunnel_command_uses_configured_ssh_executable() -> None: + tunnel = SshTunnel( + ssh_target='alice@bastion', + remote_host='db.internal', + remote_port=3307, + local_port=4406, + ssh_executable='/opt/bin/ssh', + ) + + assert tunnel.command()[0] == '/opt/bin/ssh' + + +def test_ssh_tunnel_command_forwards_remote_socket() -> None: + tunnel = SshTunnel( + ssh_target='alice@bastion', + remote_host='db.internal', + remote_port=3307, + remote_socket='/var/run/mysqld/mysqld.sock', + local_port=4406, + ) + + assert '127.0.0.1:4406:/var/run/mysqld/mysqld.sock' in tunnel.command() + assert '127.0.0.1:4406:db.internal:3307' not in tunnel.command() + + +def test_ssh_tunnel_from_target_allocates_local_port(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(ssh_tunnel, '_find_free_local_port', lambda: 4406) + + tunnel = SshTunnel.from_target('alice@bastion:2222', remote_host='db.internal', remote_port=3307) + + assert tunnel.ssh_target == 'alice@bastion' + assert tunnel.ssh_port == 2222 + assert tunnel.local_port == 4406 + assert tunnel.remote_host == 'db.internal' + assert tunnel.remote_port == 3307 + + +def test_ssh_tunnel_from_target_uses_configured_ssh_executable() -> None: + tunnel = SshTunnel.from_target( + 'alice@bastion', + remote_host='db.internal', + remote_port=3307, + ssh_executable='/opt/bin/ssh', + ) + + assert tunnel.ssh_executable == '/opt/bin/ssh' + + +def test_ssh_tunnel_start_waits_until_local_port_listens(monkeypatch: pytest.MonkeyPatch) -> None: + started_commands: list[list[str]] = [] + + class FakeProcess: + def __init__(self, command: list[str], **_kwargs: Any) -> None: + started_commands.append(command) + + def wait(self, timeout: float | None = None) -> int: + return 0 + + def poll(self) -> None: + return None + + def terminate(self) -> None: + pass + + def kill(self) -> None: + pass + + monkeypatch.setattr(ssh_tunnel.subprocess, 'Popen', FakeProcess) + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ) + checks = iter([False, True]) + monkeypatch.setattr(tunnel, '_is_listening', lambda: next(checks)) + + tunnel.start() + tunnel.close() + + assert started_commands == [tunnel.command()] + + +def test_ssh_tunnel_start_reports_process_exit_before_ready(monkeypatch: pytest.MonkeyPatch) -> None: + class FakeProcess: + def __init__(self, *_args: Any, **_kwargs: Any) -> None: + pass + + def wait(self, timeout: float | None = None) -> int: + return 255 + + def poll(self) -> int: + return 255 + + monkeypatch.setattr(ssh_tunnel.subprocess, 'Popen', FakeProcess) + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ) + monkeypatch.setattr(tunnel, '_is_listening', lambda: False) + + with pytest.raises(SshTunnelError, match='exited before it was ready'): + tunnel.start() + + +def test_ssh_tunnel_start_reports_process_start_error(monkeypatch: pytest.MonkeyPatch) -> None: + def fail_popen(*_args: Any, **_kwargs: Any) -> None: + raise FileNotFoundError('missing ssh') + + monkeypatch.setattr(ssh_tunnel.subprocess, 'Popen', fail_popen) + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ) + monkeypatch.setattr(tunnel, '_is_listening', lambda: False) + + with pytest.raises(SshTunnelError, match='Unable to start SSH tunnel process: missing ssh') as excinfo: + tunnel.start() + + assert isinstance(excinfo.value.__cause__, FileNotFoundError) + + +def test_ssh_tunnel_start_reports_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ready_timeout=0, + ) + monkeypatch.setattr(tunnel, '_run', lambda: None) + + with pytest.raises(SshTunnelError, match='Timed out waiting for SSH tunnel'): + tunnel.start() + + +def test_ssh_tunnel_close_terminates_running_process() -> None: + calls: list[str] = [] + + class FakeProcess: + def poll(self) -> None: + return None + + def terminate(self) -> None: + calls.append('terminate') + + def wait(self, timeout: float | None = None) -> int: + calls.append(f'wait:{timeout}') + return 0 + + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ) + tunnel.process = cast(Any, FakeProcess()) + + tunnel.close() + + assert calls == ['terminate', 'wait:5'] + + +def test_ssh_tunnel_close_kills_process_after_terminate_timeout() -> None: + calls: list[str] = [] + + class FakeProcess: + def __init__(self) -> None: + self.wait_calls = 0 + + def poll(self) -> None: + return None + + def terminate(self) -> None: + calls.append('terminate') + + def wait(self, timeout: float | None = None) -> int: + calls.append(f'wait:{timeout}') + self.wait_calls += 1 + if self.wait_calls == 1: + assert timeout is not None + raise subprocess.TimeoutExpired('ssh', timeout) + return 0 + + def kill(self) -> None: + calls.append('kill') + + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ) + tunnel.process = cast(Any, FakeProcess()) + + tunnel.close() + + assert calls == ['terminate', 'wait:5', 'kill', 'wait:None'] + + +def test_ssh_tunnel_close_joins_running_thread() -> None: + calls: list[str] = [] + + class FakeThread: + def is_alive(self) -> bool: + return True + + def join(self, timeout: float | None = None) -> None: + calls.append(f'join:{timeout}') + + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ) + tunnel._thread = cast(Any, FakeThread()) + + tunnel.close() + + assert calls == ['join:5'] + + +def test_ssh_tunnel_is_listening_returns_true(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[tuple[tuple[str, int], float]] = [] + + class FakeConnection: + def __enter__(self) -> 'FakeConnection': + return self + + def __exit__(self, *_args: Any) -> None: + pass + + def fake_create_connection(address: tuple[str, int], timeout: float) -> FakeConnection: + calls.append((address, timeout)) + return FakeConnection() + + monkeypatch.setattr(ssh_tunnel.socket, 'create_connection', fake_create_connection) + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ) + + assert tunnel._is_listening() is True + assert calls == [(('127.0.0.1', 4406), 0.05)] + + +def test_ssh_tunnel_is_listening_returns_false(monkeypatch: pytest.MonkeyPatch) -> None: + def fail_create_connection(*_args: Any, **_kwargs: Any) -> None: + raise socket.timeout + + monkeypatch.setattr(ssh_tunnel.socket, 'create_connection', fail_create_connection) + tunnel = SshTunnel( + ssh_target='bastion', + remote_host='db.internal', + remote_port=3306, + local_port=4406, + ) + + assert tunnel._is_listening() is False