Files
autossh-mgr/src/autossh_mgr/cli.py
2026-05-21 12:47:32 +08:00

244 lines
7.6 KiB
Python

import click
from dataclasses import replace
from autossh_mgr.config import (
TunnelConfig, get_config_dir, ensure_dirs,
load_tunnels, save_tunnels, get_tunnel,
)
from autossh_mgr.check import check_connectivity
from autossh_mgr.process import (
get_status, start_tunnel, stop_tunnel, delete_pid_file,
)
from autossh_mgr.display import (
print_tunnel_table, print_tunnel_detail, print_status_table,
)
_FIELD_TYPES: dict[str, type] = {
"host": str,
"port": int,
"user": str,
"identity_file": str,
"local_host": str,
"local_port": int,
"remote_host": str,
"remote_port": int,
"ssh_options": str,
}
@click.group()
@click.pass_context
def cli(ctx: click.Context) -> None:
ctx.ensure_object(dict)
config_dir = get_config_dir()
ensure_dirs(config_dir)
ctx.obj["config_dir"] = config_dir
@cli.command()
@click.pass_context
def init(ctx: click.Context) -> None:
"""Initialize config directory."""
config_dir = ctx.obj["config_dir"]
ensure_dirs(config_dir)
click.echo(f"Initialized {config_dir}")
@cli.command()
@click.argument("name")
@click.option("--host", default=None)
@click.option("--port", type=int, default=None)
@click.option("--user", default=None)
@click.option("--identity-file", default=None)
@click.option("--local-host", default=None)
@click.option("--local-port", type=int, default=None)
@click.option("--remote-host", default=None)
@click.option("--remote-port", type=int, default=None)
@click.option("--ssh-options", default=None)
@click.pass_context
def add(
ctx, name, host, port, user, identity_file,
local_host, local_port, remote_host, remote_port, ssh_options,
):
"""Add a new tunnel configuration."""
config_dir = ctx.obj["config_dir"]
tunnels = load_tunnels(config_dir)
if any(t.name == name for t in tunnels):
raise click.ClickException(f"Tunnel '{name}' already exists")
if host is None:
host = click.prompt("Public server host")
if user is None:
user = click.prompt("SSH user")
if local_port is None:
local_port = click.prompt("Local port", type=int)
if remote_port is None:
remote_port = click.prompt("Remote port", type=int)
if port is None:
port = 22
if identity_file is None:
identity_file = "~/.ssh/id_ed25519"
if local_host is None:
local_host = "127.0.0.1"
if remote_host is None:
remote_host = "0.0.0.0"
if ssh_options is None:
ssh_options = ""
tunnel = TunnelConfig(
name=name, host=host, port=port, user=user,
identity_file=identity_file, local_host=local_host,
local_port=local_port, remote_host=remote_host,
remote_port=remote_port, ssh_options=ssh_options,
)
tunnels.append(tunnel)
save_tunnels(config_dir, tunnels)
click.echo(f"Added tunnel '{name}'")
@cli.command()
@click.argument("name")
@click.pass_context
def remove(ctx, name):
"""Remove a tunnel configuration."""
config_dir = ctx.obj["config_dir"]
tunnels = load_tunnels(config_dir)
if not any(t.name == name for t in tunnels):
raise click.ClickException(f"No tunnel named '{name}'")
status = get_status(config_dir, name)
if status.state == "running":
raise click.ClickException(f"Stop '{name}' before removing it")
if not click.confirm(f"Remove tunnel '{name}'?"):
return
save_tunnels(config_dir, [t for t in tunnels if t.name != name])
click.echo(f"Removed tunnel '{name}'")
@cli.command(name="list")
@click.pass_context
def list_cmd(ctx):
"""List all tunnel configurations."""
config_dir = ctx.obj["config_dir"]
tunnels = load_tunnels(config_dir)
statuses = {t.name: get_status(config_dir, t.name) for t in tunnels}
print_tunnel_table(tunnels, statuses)
@cli.command()
@click.argument("name")
@click.pass_context
def show(ctx, name):
"""Show full details for a tunnel."""
config_dir = ctx.obj["config_dir"]
tunnel = get_tunnel(config_dir, name)
status = get_status(config_dir, name)
print_tunnel_detail(tunnel, status)
@cli.command(name="config")
@click.argument("name")
@click.argument("key")
@click.argument("value")
@click.pass_context
def config_cmd(ctx, name, key, value):
"""Update a single tunnel configuration field."""
config_dir = ctx.obj["config_dir"]
if key == "name":
raise click.ClickException("Cannot rename a tunnel via config. Remove and re-add.")
if key not in _FIELD_TYPES:
raise click.ClickException(f"Unknown field '{key}'")
try:
typed_value = _FIELD_TYPES[key](value)
except ValueError:
raise click.ClickException(
f"Invalid value for '{key}': expected {_FIELD_TYPES[key].__name__}"
)
tunnels = load_tunnels(config_dir)
if not any(t.name == name for t in tunnels):
raise click.ClickException(f"No tunnel named '{name}'")
tunnels = [replace(t, **{key: typed_value}) if t.name == name else t for t in tunnels]
save_tunnels(config_dir, tunnels)
click.echo(f"Updated {name}.{key} = {typed_value}")
@cli.command()
@click.argument("name", required=False)
@click.pass_context
def start(ctx, name):
"""Start a tunnel (or all tunnels if no name given)."""
config_dir = ctx.obj["config_dir"]
tunnels = load_tunnels(config_dir)
targets = (
[get_tunnel(config_dir, name)] if name
else sorted(tunnels, key=lambda t: t.name)
)
for tunnel in targets:
status = get_status(config_dir, tunnel.name)
if status.state == "running":
click.echo(f"{tunnel.name} already running (pid: {status.pid})")
continue
if status.state == "stale":
click.echo(f"{tunnel.name}: stale pid detected, cleaning up")
delete_pid_file(config_dir, tunnel.name)
pid = start_tunnel(config_dir, tunnel)
click.echo(f"Started {tunnel.name} (pid: {pid})")
@cli.command()
@click.argument("name", required=False)
@click.pass_context
def stop(ctx, name):
"""Stop a tunnel (or all running tunnels if no name given)."""
config_dir = ctx.obj["config_dir"]
tunnels = load_tunnels(config_dir)
targets = (
[get_tunnel(config_dir, name)] if name
else sorted(tunnels, key=lambda t: t.name)
)
for tunnel in targets:
status = get_status(config_dir, tunnel.name)
if status.state == "stopped":
click.echo(f"{tunnel.name} is not running")
continue
stop_tunnel(config_dir, tunnel.name)
click.echo(f"Stopped {tunnel.name}")
@cli.command()
@click.argument("name", required=False)
@click.pass_context
def restart(ctx, name):
"""Restart a tunnel (or all tunnels if no name given)."""
ctx.invoke(stop, name=name)
ctx.invoke(start, name=name)
@cli.command()
@click.argument("name", required=False)
@click.pass_context
def status(ctx, name):
"""Show tunnel status (all tunnels if no name given)."""
config_dir = ctx.obj["config_dir"]
if name:
tunnel = get_tunnel(config_dir, name)
s = get_status(config_dir, name)
print_tunnel_detail(tunnel, s)
else:
tunnels = load_tunnels(config_dir)
statuses = {t.name: get_status(config_dir, t.name) for t in tunnels}
print_status_table(tunnels, statuses)
@cli.command(name="check")
@click.argument("name")
@click.pass_context
def check_cmd(ctx, name):
"""Check SSH connectivity to a tunnel's server."""
config_dir = ctx.obj["config_dir"]
tunnel = get_tunnel(config_dir, name)
success, msg = check_connectivity(tunnel)
if success:
click.echo(f"OK: {tunnel.user}@{tunnel.host}:{tunnel.port} is reachable")
else:
raise click.ClickException(f"Connection failed: {msg}")