colab_tunnel/_tunnels.py

"""
Провайдеры туннелей на основе бинарных утилит.
"""
from atexit import register as exit_register
from pathlib import Path
from subprocess import PIPE, Popen
from time import sleep

from requests import get as get_url

from ._config import (
    WORK_FOLDER,
    cloudflare_bin, tmole_bin, tunwg_bin,
    go_localt_bin, gradio_bin, mmar_bin,
    tunnelite_bin, beeceptor_bin, bore_bin,
    boredigital_bin, colab_native_url, links_file,
)
from ._logger import logger
from ._registry import tunnel_provider
from ._utils import (
    download, unpack_archive, run,
    read_until_pattern, drain_process_output,
    kill_process_by_name, terminate_process,
    move_path, is_ipv4, _snapshot,
    get_github_latest_release_url,
)


def get_revproxy_url(
    bin_url: str,
    need_unpack: bool,
    bin_path: Path,
    start_commands: list,
    read_from_stderr: bool,
    url_pattern: str = r'https://\S+',
    timeout: float = 20.0,
    write_link: bool = False,
    stdin_input: str | None = None,
    extra_env: dict[str, str] | None = None,
    read_both_streams: bool = False,
) -> str:
    """
    Универсальный запуск бинарного туннельного провайдера.

    При первом вызове скачивает и устанавливает бинарник.
    Использует неблокирующее чтение вывода для поиска URL.

    Args:
        bin_url:          URL скачивания бинарника (или архива).
        need_unpack:      Нужно ли распаковывать архив.
        bin_path:         Целевой путь бинарника на диске.
        start_commands:   Команда запуска со всеми аргументами.
        read_from_stderr: Искать URL в stderr (иначе в stdout).
        url_pattern:      Regex для поиска публичного URL.
        timeout:          Максимальное время ожидания URL (сек).
        write_link:       Сохранять ли ссылку в links.txt.
        stdin_input:      Строка для отправки в stdin процесса.
        extra_env: Дополнительные переменные окружения для процесса.
                   Накладываются поверх os.environ (не заменяют его).
        read_both_streams: Читать ли оба потока.

    Returns:
        Публичный URL туннеля.
    """
    # ── Установка бинарника ──────────────────────────────────────────────
    if not bin_path.exists():
        WORK_FOLDER.mkdir(parents=True, exist_ok=True)
        snapshot_before = _snapshot(WORK_FOLDER)

        if need_unpack:
            archive = download(bin_url, save_path=WORK_FOLDER, progress=False)
            unpack_archive(archive, WORK_FOLDER, rm_archive=True)
        else:
            download(bin_url, save_path=bin_path.parent, progress=False)

        new_files = list(_snapshot(WORK_FOLDER) - snapshot_before)

        if not new_files:
            raise RuntimeError(
                f'В {WORK_FOLDER} не появилось новых файлов после загрузки {bin_url}'
            )

        if len(new_files) == 1:
            move_path(new_files[0], bin_path)
        else:
            # Несколько файлов (например, архив с доп. файлами) — берём наибольший
            largest = max(new_files, key=lambda f: f.stat().st_size)
            move_path(largest, bin_path)
            for f in new_files:
                if f != largest and f.exists():
                    f.unlink(missing_ok=True)

        bin_path.chmod(0o755)
        logger.debug(f'Установлен бинарник: {bin_path}')

    # ── Завершение предыдущего экземпляра ───────────────────────────────
    kill_process_by_name(bin_path.name)
    # ── Подготовка окружения ─────────────────────────────────────────────
    import os as _os
    proc_env = _os.environ.copy()
    if extra_env:
        proc_env.update(extra_env)
    # ── Запуск процесса ──────────────────────────────────────────────────
    stdin_flag = PIPE if stdin_input is not None else None
    process = Popen(
        start_commands,
        stdout=PIPE, stderr=PIPE, stdin=stdin_flag,
        env=proc_env,  # ← передаём окружение
    )
    if stdin_input is not None and process.stdin:
        try:
            process.stdin.write(stdin_input.encode())
            process.stdin.flush()
        finally:
            process.stdin.close()  # Сигнал EOF — многие утилиты ждут его перед стартом

    # ── Чтение URL из вывода ─────────────────────────────────────────────
    try:
        url, full_output = read_until_pattern(
            process=process,
            url_pattern=url_pattern,
            timeout=timeout,
            read_from_stderr=read_from_stderr,
            read_both_streams=read_both_streams,  # ← передаём дальше
        )
    except RuntimeError:
        kill_process_by_name(bin_path.name)
        try:
            process.wait(timeout=3)
        except Exception:
            process.kill()
            process.wait()
        raise

    # Слив вывода в фоне — предотвращает блокировку из-за переполнения пайпа
    drain_process_output(process)

    # Проверка наличия IPv4 в выводе (пароль/IP у некоторых провайдеров)
    ipv4 = next(
        (ln.strip() for ln in full_output.splitlines() if is_ipv4(ln.strip())),
        None,
    )

    exit_register(terminate_process, bin_path.name, process)

    if write_link:
        try:
            links_file.write_text(url)
        except Exception as e:
            logger.warning(f'Не удалось записать ссылку в {links_file}: {e}')

    logger.debug(f'[{bin_path.name}] Туннель: {url}')
    return f'{url}\n пароль(IP): {ipv4}' if ipv4 else url


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

@tunnel_provider('tmole')
def get_tmole_url(port: int) -> str:
    return get_revproxy_url(
        bin_url='https://tunnelmole.com/downloads/tmole-linux.gz',
        need_unpack=True,
        bin_path=tmole_bin,
        start_commands=[str(tmole_bin), str(port)],
        read_from_stderr=False,
        timeout=20.0,
    )


@tunnel_provider('tunwg', enabled=False)  # Сервис недоступен с 2026-06-16, включить: enable_provider('tunwg')
def get_tunwg_url(port: int) -> str:
    return get_revproxy_url(
        bin_url='https://github.com/ntnj/tunwg/releases/latest/download/tunwg',
        need_unpack=False,
        bin_path=tunwg_bin,
        start_commands=[str(tunwg_bin), f'--forward=http://127.0.0.1:{port}'],
        read_from_stderr=True,
        timeout=20.0,
    )


@tunnel_provider('cloudflared')
def get_cloudflared_url(port: int) -> str:
    return get_revproxy_url(
        bin_url='https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64',
        need_unpack=False,
        bin_path=cloudflare_bin,
        start_commands=[str(cloudflare_bin), 'tunnel', '--url', f'http://127.0.0.1:{port}'],
        read_from_stderr=True,
        url_pattern=r'(?P<url>https?://\S+\.trycloudflare\.com)',
        timeout=25.0,
    )


@tunnel_provider('localt')
def get_localt_url(port: int) -> str:
    return get_revproxy_url(
        bin_url='https://huggingface.co/prolapse/go_localt/resolve/main/go_localt',
        need_unpack=False,
        bin_path=go_localt_bin,
        start_commands=[str(go_localt_bin), str(port)],
        read_from_stderr=False,
        timeout=20.0,
    )


@tunnel_provider('gradio')
def get_gradio_url(port: int) -> str:
    max_attempts = 3
    last_error: Exception | None = None

    for attempt in range(max_attempts):
        try:
            resp = get_url('https://api.gradio.app/v2/tunnel-request', timeout=10)
            resp.raise_for_status()
            info = resp.json()[0]
            remote_host, remote_port = info['host'], int(info['port'])
            cmd = [
                str(gradio_bin), 'http', '-n', 'random',
                '-l', str(port), '-i', '127.0.0.1',
                '--uc', '--sd', 'random', '--ue',
                '--server_addr', f'{remote_host}:{remote_port}',
                '--disable_log_color',
            ]
            return get_revproxy_url(
                bin_url='https://cdn-media.huggingface.co/frpc-gradio-0.1/frpc_linux_amd64',
                need_unpack=False,
                bin_path=gradio_bin,
                start_commands=cmd,
                read_from_stderr=False,
                timeout=20.0,
            )
        except Exception as e:
            last_error = e
            if attempt < max_attempts - 1:
                logger.warning(f'Попытка {attempt + 1} Gradio провалилась: {e}. Повтор через 5с...')
                sleep(5)

    raise RuntimeError(f'После {max_attempts} попыток Gradio-туннель не запущен: {last_error}')


@tunnel_provider('native')
def get_native_url(port: int) -> str:
    """
    Встроенный прокси Google Colab. Результат кэшируется по номеру порта.
    """
    try:
        from google.colab.output import eval_js  # type: ignore[import]
    except ImportError:
        raise RuntimeError('get_native_url доступен только в среде Google Colab.')

    cache_url_file  = colab_native_url
    cache_port_file = colab_native_url.with_suffix('.port')

    if cache_url_file.exists() and cache_port_file.exists():
        try:
            if int(cache_port_file.read_text().strip()) == port:
                cached = cache_url_file.read_text().strip()
                if cached:
                    logger.debug(f'native URL из кэша (порт {port}): {cached}')
                    return cached
        except Exception:
            pass

    url = eval_js(f'google.colab.kernel.proxyPort({port})')

    try:
        WORK_FOLDER.mkdir(parents=True, exist_ok=True)
        cache_url_file.write_text(url)
        cache_port_file.write_text(str(port))
    except Exception as e:
        logger.warning(f'Не удалось кэшировать native URL: {e}')

    return url


@tunnel_provider('mmar')
def get_mmar_url(port: int) -> str:
    return get_revproxy_url(
        bin_url='https://github.com/yusuf-musleh/mmar/releases/latest/download/mmar_Linux_x86_64.tar.gz',
        need_unpack=True,
        bin_path=mmar_bin,
        start_commands=[str(mmar_bin), 'client', '--local-port', str(port)],
        read_from_stderr=True,
        url_pattern=r'(?P<url>https?://\S+\.mmar\.dev)',
        timeout=20.0,
    )


@tunnel_provider('tunnelite')
def get_tunnelite_url(port: int) -> str:
    """
    tunnelite — интерактивное TUI-приложение, которое проверяет isatty(stdout).
    Без PTY оно выводит URL и сразу завершается, не поддерживая туннель → 404.
    `script -q -c "..." /dev/null` выделяет pseudo-TTY для дочернего процесса,
    tunnelite думает что работает в терминале и держит соединение.
    script при этом транслирует весь вывод в свой stdout — читаем через обычный PIPE.
    """
    kill_process_by_name(tunnelite_bin.name)
    return get_revproxy_url(
        bin_url='https://github.com/hehepiska/tunnellite-linux/releases/download/latest/tunnelite',
        need_unpack=False,
        bin_path=tunnelite_bin,
        start_commands=[
            'script', '-q', '-c',
            f'{tunnelite_bin} http://127.0.0.1:{port}',
            '/dev/null',
        ],
        read_from_stderr=False,
        url_pattern=r'(?P<url>https?://\S+\.tunnelite\.com)',
        timeout=45.0,
    )


@tunnel_provider('beeceptor')
def get_beeceptor_url(port: int) -> str:
    return get_revproxy_url(
        bin_url='https://cdn.beeceptor.com/downloads/cli/latest/beeceptor-cli-linux-x64.tar.gz',
        need_unpack=True,
        bin_path=beeceptor_bin,
        start_commands=[str(beeceptor_bin), '-p', str(port), '-y'],
        read_from_stderr=False,
        url_pattern=r'(?P<url>https?://\S+\.beeceptor\.com)',
        timeout=20.0,
        stdin_input='\n',
    )


@tunnel_provider('bore')
def get_bore_url(port: int) -> str:
    bin_url = 'https://github.com/ekzhang/bore/releases/download/v0.6.0/bore-v0.6.0-x86_64-unknown-linux-musl.tar.gz'
    if not bore_bin.exists():
        try:
            bin_url = get_github_latest_release_url(
                'ekzhang/bore',
                r'bore-v[\d.]+-x86_64-unknown-linux-musl\.tar\.gz',
            )
            logger.debug(f'Актуальный URL bore: {bin_url}')
        except Exception as e:
            logger.warning(f'GitHub API недоступен, используется закреплённая версия bore: {e}')

    raw = get_revproxy_url(
        bin_url=bin_url,
        need_unpack=True,
        bin_path=bore_bin,
        start_commands=[str(bore_bin), 'local', str(port), '--to', 'bore.pub'],
        # read_both_streams: bore пишет в stderr через tracing-subscriber,
        # но читаем оба потока — на случай если конкретная сборка пишет в stdout
        read_from_stderr=False,
        read_both_streams=True,
        url_pattern=r'bore\.pub:\d+',
        timeout=20.0,
    )
    return f'http://{raw}/'



@tunnel_provider('boredigital')
def get_boredigital_url(port: int) -> str:
    return get_revproxy_url(
        bin_url='https://github.com/jkuri/bore/releases/latest/download/bore_linux_amd64',
        need_unpack=False,
        bin_path=boredigital_bin,
        start_commands=[str(boredigital_bin), '-lp', str(port)],
        read_from_stderr=False,
        url_pattern=r'(?P<url>https://\S+\.bore\.digital)',
        timeout=20.0,
    )