Files
mytoolkit/mycli/commands/ssh.py
T

101 lines
2.7 KiB
Python

"""SSH connection utilities."""
import subprocess
from pathlib import Path
import click
from mycli.config import config
from mycli.utils import handle_errors
def _get_hosts() -> dict:
"""Get SSH hosts from config."""
return config.get("ssh_hosts", {})
def _build_ssh_base_cmd(host_config: dict) -> list[str]:
"""Build base SSH command with host and port."""
cmd = ["ssh", "-p", str(host_config["port"])]
if "key" in host_config:
key_path = Path(host_config["key"]).expanduser()
cmd.extend(["-i", str(key_path)])
cmd.append(host_config["host"])
return cmd
@click.group(name="ssh")
def ssh_cmd():
"""SSH connection utilities."""
pass
@ssh_cmd.command("connect")
@click.argument("name")
@handle_errors
def connect(name):
"""Connect to a configured host."""
hosts = _get_hosts()
if name not in hosts:
click.echo(f"Unknown host: {name}", err=True)
click.echo(f"Available: {', '.join(hosts.keys())}", err=True)
raise click.Exit(1)
cmd = _build_ssh_base_cmd(hosts[name])
click.secho(f"Connecting to {name}...", fg="cyan")
subprocess.run(cmd)
@ssh_cmd.command("tunnel")
@click.argument("name")
@click.option("--local-port", "-l", default=11111, help="Local port")
@click.option("--remote-port", "-r", default=11111, help="Remote port")
@handle_errors
def tunnel(name, local_port, remote_port):
"""Create SSH tunnel with port forwarding."""
hosts = _get_hosts()
if name not in hosts:
click.echo(f"Unknown host: {name}", err=True)
raise click.Exit(1)
host_config = hosts[name]
cmd = ["ssh", "-p", str(host_config["port"]), "-L", f"{local_port}:localhost:{remote_port}"]
if "key" in host_config:
key_path = Path(host_config["key"]).expanduser()
cmd.extend(["-i", str(key_path)])
cmd.append(host_config["host"])
click.secho(f"Creating tunnel {local_port} -> {remote_port} on {name}...", fg="cyan")
click.secho(f"Command: {' '.join(cmd)}", fg="dim")
subprocess.run(cmd)
@ssh_cmd.command("list")
def list_hosts():
"""List configured SSH hosts."""
hosts = _get_hosts()
click.secho("Configured hosts:", fg="cyan")
for name, host_config in sorted(hosts.items()):
click.echo(f" {name:12} {host_config['host']}:{host_config['port']}")
@ssh_cmd.command("cmd")
@click.argument("name")
@click.argument("command")
@handle_errors
def run_cmd(name, command):
"""Run a command on remote host."""
hosts = _get_hosts()
if name not in hosts:
click.echo(f"Unknown host: {name}", err=True)
raise click.Exit(1)
cmd = _build_ssh_base_cmd(hosts[name])
cmd.append(command)
subprocess.run(cmd)