PangolinSentinel/backend/remote.py

import asyncio
import os
import shutil
import time
from pathlib import Path
from typing import Dict, Any, List, Optional, Tuple
from backend.config import load_settings, DEFAULT_KEY_FILE, DEFAULT_PUB_KEY_FILE, DATA_DIR

def get_public_key() -> str:
    if DEFAULT_PUB_KEY_FILE.exists():
        return DEFAULT_PUB_KEY_FILE.read_text(encoding="utf-8").strip()
    return ""

def ensure_ssh_key() -> str:
    if not DEFAULT_KEY_FILE.exists():
        DEFAULT_KEY_FILE.parent.mkdir(parents=True, exist_ok=True)
        os.system(f'ssh-keygen -t ed25519 -f "{DEFAULT_KEY_FILE}" -N "" -C "pangolin-sentinel@unraid" -q')
    if DEFAULT_KEY_FILE.exists():
        try:
            os.chmod(DEFAULT_KEY_FILE, 0o600)
        except Exception:
            pass
    return get_public_key()

def resolve_key_path(configured_path: Optional[str] = None) -> str:
    if configured_path and os.path.exists(configured_path):
        return configured_path
    if DEFAULT_KEY_FILE.exists():
        return str(DEFAULT_KEY_FILE)
    return configured_path or str(DEFAULT_KEY_FILE)

def build_ssh_prefix(settings: Optional[Dict[str, Any]] = None) -> List[str]:
    if not settings:
        settings = load_settings()

    host = settings.get("vps_host", "").strip()
    port = str(settings.get("vps_port", 22))
    user = settings.get("vps_user", "root").strip()
    auth_method = settings.get("vps_auth_method", "key")
    raw_key_path = settings.get("vps_ssh_key_path", str(DEFAULT_KEY_FILE))
    key_path = resolve_key_path(raw_key_path)
    password = settings.get("vps_ssh_password", "")
    known_hosts = str(DATA_DIR / "known_hosts")

    ssh_opts = [
        "-T",
        "-o", "StrictHostKeyChecking=accept-new",
        "-o", f"UserKnownHostsFile={known_hosts}",
        "-o", "ConnectTimeout=8",
        "-o", "ServerAliveInterval=15",
        "-p", port
    ]

    if auth_method == "password" and password:
        return ["sshpass", "-p", password, "ssh"] + ssh_opts + [f"{user}@{host}"]
    else:
        # Use Key Auth
        if os.path.exists(key_path):
            try:
                os.chmod(key_path, 0o600)
            except Exception:
                pass
            return ["ssh", "-i", key_path, "-o", "BatchMode=yes"] + ssh_opts + [f"{user}@{host}"]
        return ["ssh", "-o", "BatchMode=yes"] + ssh_opts + [f"{user}@{host}"]

def build_scp_prefix(settings: Optional[Dict[str, Any]] = None) -> List[str]:
    if not settings:
        settings = load_settings()

    port = str(settings.get("vps_port", 22))
    auth_method = settings.get("vps_auth_method", "key")
    raw_key_path = settings.get("vps_ssh_key_path", str(DEFAULT_KEY_FILE))
    key_path = resolve_key_path(raw_key_path)
    password = settings.get("vps_ssh_password", "")
    known_hosts = str(DATA_DIR / "known_hosts")

    scp_opts = [
        "-P", port,
        "-o", "StrictHostKeyChecking=accept-new",
        "-o", f"UserKnownHostsFile={known_hosts}",
        "-o", "ConnectTimeout=10"
    ]

    if auth_method == "password" and password:
        return ["sshpass", "-p", password, "scp"] + scp_opts
    else:
        if os.path.exists(key_path):
            try:
                os.chmod(key_path, 0o600)
            except Exception:
                pass
            return ["scp", "-i", key_path, "-o", "BatchMode=yes"] + scp_opts
        return ["scp", "-o", "BatchMode=yes"] + scp_opts

async def run_remote_quick(cmd: str, settings: Optional[Dict[str, Any]] = None) -> Tuple[int, str, str]:
    prefix = build_ssh_prefix(settings)
    full_cmd = prefix + [cmd]
    try:
        proc = await asyncio.create_subprocess_exec(
            *full_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        stdout, stderr = await proc.communicate()
        return proc.returncode, stdout.decode('utf-8', errors='replace'), stderr.decode('utf-8', errors='replace')
    except Exception as e:
        return 1, "", str(e)

async def test_vps_connection(settings: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    if not settings:
        settings = load_settings()

    host = settings.get("vps_host", "").strip()
    if not host:
        return {"success": False, "error": "No VPS host configured. Please provide IP or hostname."}

    # 1. Test basic SSH connectivity & gather OS summary
    cmd_ping = 'echo "PANGOLIN_PONG" && uname -srm && cat /etc/os-release 2>/dev/null | grep PRETTY_NAME | cut -d= -f2 | tr -d \'"\''
    code, stdout, stderr = await run_remote_quick(cmd_ping, settings)

    if code != 0:
        err_msg = stderr.strip() or "Connection timed out or failed."
        if "Permission denied" in err_msg:
            err_msg += " (SSH Authentication failed. Check key in ~/.ssh/authorized_keys or password)"
        return {"success": False, "error": err_msg}

    lines = stdout.strip().split("\n")
    os_name = lines[2] if len(lines) > 2 else (lines[1] if len(lines) > 1 else "Linux")
    kernel = lines[1] if len(lines) > 1 else ""

    # 2. Check Docker availability on VPS
    code_docker, stdout_docker, _ = await run_remote_quick('docker --version', settings)
    docker_installed = (code_docker == 0)
    docker_version = stdout_docker.strip() if docker_installed else "Not detected"

    # 3. Check Pangolin Directory
    pangolin_dir = settings.get("vps_pangolin_dir", "/opt/pangolin")
    chk_dir_cmd = f'[ -d "{pangolin_dir}" ] && echo "DIR_EXISTS" || echo "DIR_MISSING"'
    _, stdout_dir, _ = await run_remote_quick(chk_dir_cmd, settings)
    pangolin_dir_found = "DIR_EXISTS" in stdout_dir

    # 4. Check SQLite database status in Pangolin directory
    chk_db_cmd = f'find "{pangolin_dir}" -name "*.db" -o -name "*.sqlite" 2>/dev/null'
    _, stdout_db, _ = await run_remote_quick(chk_db_cmd, settings)
    dbs_found = [p.strip() for p in stdout_db.strip().split("\n") if p.strip()]

    return {
        "success": True,
        "os_name": os_name,
        "kernel": kernel,
        "docker_installed": docker_installed,
        "docker_version": docker_version,
        "pangolin_dir_found": pangolin_dir_found,
        "sqlite_databases": dbs_found,
        "message": "SSH Connection successful! Remote host is ready for backups and updates."
    }