colab_tunnel/_ssh.py

"""
Провайдеры туннелей на основе SSH.
"""
import subprocess
from atexit import register as exit_register
from pathlib import Path
from subprocess import DEVNULL, PIPE, Popen, TimeoutExpired

from ._logger import logger
from ._registry import tunnel_provider
from ._utils import kill_process_by_name, terminate_process, drain_process_output, read_until_pattern

# Реестр активных SSH-туннелей: имя процесса → объект Popen
# dict вместо list — предотвращает накопление мёртвых процессов при повторных вызовах
_active_tunnels: dict[str, Popen] = {}

_cmd = [
    'ssh',
    '-o', 'StrictHostKeyChecking=no',
    '-o', 'UserKnownHostsFile=/dev/null',
    '-o', 'ExitOnForwardFailure=yes',
    '-o', 'ServerAliveInterval=10',
    '-o', 'ServerAliveCountMax=3',
]


def _register_active_tunnel(name: str, process: Popen) -> None:
    """Регистрирует туннель, корректно завершая предыдущий с тем же именем."""
    if name in _active_tunnels:
        old = _active_tunnels[name]
        try:
            old.terminate()
            old.wait(timeout=3)
        except TimeoutExpired:
            old.kill()
            old.wait()
        except Exception:
            pass
    _active_tunnels[name] = process


def ensure_ssh_key() -> None:
    """Гарантирует наличие ed25519 SSH-ключа. Генерирует без пароля при отсутствии."""
    ssh_dir = Path.home() / '.ssh'
    ssh_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
    key_file = ssh_dir / 'id_ed25519'

    if not key_file.exists():
        subprocess.run(
            ['ssh-keygen', '-t', 'ed25519', '-N', '', '-f', str(key_file)],
            stdout=DEVNULL,
            stderr=DEVNULL,
        )
        logger.debug('SSH-ключ ed25519 сгенерирован.')


def get_ssh_tunnel_url(
        start_commands: list[str],
        process_name: str,
        url_pattern: str,
        timeout: float = 15.0,
) -> str:
    """
    Запускает SSH-туннельный процесс и ждёт появления публичного URL.

    stdin намеренно оставляется ОТКРЫТЫМ: SSH завершается при получении EOF,
    что прервало бы туннель. Объект процесса сохраняется в _active_tunnels
    для предотвращения сборки мусора.

    Args:
        start_commands: Полная команда SSH с флагами.
        process_name:   Идентификатор процесса (часть адреса хоста).
        url_pattern:    Regex для поиска URL в выводе.
        timeout:        Таймаут ожидания URL в секундах.

    Returns:
        Публичный URL туннеля.
    """
    ensure_ssh_key()
    kill_process_by_name(process_name)

    process = Popen(start_commands, stdout=PIPE, stderr=PIPE, stdin=PIPE)

    try:
        url, _ = read_until_pattern(
            process=process,
            url_pattern=url_pattern,
            timeout=timeout,
            read_both_streams=True,  # SSH может писать URL в stdout или stderr
        )
    except RuntimeError:
        process.terminate()
        try:
            process.wait(timeout=3)
        except TimeoutExpired:
            process.kill()
            process.wait()
        kill_process_by_name(process_name)
        raise

    drain_process_output(process)
    _register_active_tunnel(process_name, process)
    exit_register(terminate_process, process_name, process)

    logger.debug(f'[{process_name}] SSH-туннель: {url}')
    return url


# ---------------------------------------------------------------------------
# Провайдеры
# ---------------------------------------------------------------------------

@tunnel_provider('optimistix')
def get_optimistix_url(port: int) -> str:
    return get_ssh_tunnel_url(
        start_commands=_cmd + [
            '-p', '1122',
            '-R', f'80:127.0.0.1:{port}',
            'ssh.optimistixtunnel.com',
        ],
        process_name='ssh.optimistixtunnel.com',
        url_pattern=r'https://\S+\.otnl\.link',
    )


@tunnel_provider('srvus')
def get_srvus_url(port: int) -> str:
    return get_ssh_tunnel_url(
        start_commands=_cmd + [
            '-R', f'1:127.0.0.1:{port}',
            'srv.us',
        ],
        process_name='srv.us',
        url_pattern=r'https://[a-z0-9]{10,}\.srv\.us/?',
    )


@tunnel_provider('serveo')
def get_serveo_url(port: int) -> str:
    return get_ssh_tunnel_url(
        start_commands=_cmd + [
            '-R', f'80:127.0.0.1:{port}',
            'serveo.net',
        ],
        process_name='serveo.net',
        url_pattern=r'https://\S+\.serveousercontent\.com',
    )


@tunnel_provider('localhostrun')
def get_localhostrun_url(port: int) -> str:
    return get_ssh_tunnel_url(
        start_commands=_cmd + [
            '-R', f'80:127.0.0.1:{port}',
            'nokey@localhost.run',
        ],
        process_name='localhost.run',
        url_pattern=r'https://(?!admin\b)[a-zA-Z0-9-]+\.(?:localhost\.run|lhr\.life)',
    )