sshclient/core/manager.py

"""
连接管理器：保存多个 SSH 会话的配置和活跃连接。
配置持久化到 ~/.sshclient/hosts.json。
"""
import json
import os
import threading
from pathlib import Path
from typing import Dict, List, Optional

from .ssh_client import SSHConnection


CONFIG_DIR = Path.home() / ".sshclient"
CONFIG_FILE = CONFIG_DIR / "hosts.json"
AI_CONFIG_FILE = CONFIG_DIR / "ai.json"


class ConnectionManager:
    """多主机连接管理 + 配置持久化"""

    def __init__(self):
        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
        self._lock = threading.RLock()
        self.hosts: List[dict] = []  # 主机配置
        self.connections: Dict[str, SSHConnection] = {}  # host_id -> SSHConnection
        self._load_hosts()

    def _load_hosts(self):
        if CONFIG_FILE.exists():
            try:
                with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                    self.hosts = json.load(f)
            except Exception:
                self.hosts = []
        if not self.hosts:
            # 给一个示例条目，让 UI 不为空
            self.hosts = [{
                "id": "demo", "name": "示例主机", "host": "127.0.0.1",
                "port": 22, "username": "root", "password": "", "key_path": "",
            }]

    def save_hosts(self):
        with open(CONFIG_FILE, "w", encoding="utf-8") as f:
            json.dump(self.hosts, f, ensure_ascii=False, indent=2)

    def list_hosts(self) -> List[dict]:
        return list(self.hosts)

    def get_host(self, host_id: str) -> Optional[dict]:
        for h in self.hosts:
            if h.get("id") == host_id:
                return dict(h)
        return None

    def add_host(self, host_info: dict) -> str:
        """新增主机；返回 id"""
        with self._lock:
            new_id = host_info.get("id") or f"host-{int(__import__('time').time()*1000)}"
            host_info["id"] = new_id
            self.hosts.append(host_info)
            self.save_hosts()
            return new_id

    def update_host(self, host_id: str, host_info: dict):
        with self._lock:
            for i, h in enumerate(self.hosts):
                if h.get("id") == host_id:
                    host_info["id"] = host_id
                    self.hosts[i] = host_info
                    self.save_hosts()
                    # 断开旧连接
                    if host_id in self.connections:
                        self.connections[host_id].disconnect()
                        del self.connections[host_id]
                    return

    def remove_host(self, host_id: str):
        with self._lock:
            self.hosts = [h for h in self.hosts if h.get("id") != host_id]
            if host_id in self.connections:
                self.connections[host_id].disconnect()
                del self.connections[host_id]
            self.save_hosts()

    def connect(self, host_id: str) -> tuple:
        """连接指定主机；返回 (conn, 成功, 消息)"""
        info = self.get_host(host_id)
        if not info:
            return None, False, "主机不存在"
        with self._lock:
            conn = self.connections.get(host_id)
            if conn and conn.connected:
                return conn, True, "已连接"
            conn = SSHConnection(
                host=info["host"], port=info.get("port", 22),
                username=info.get("username", ""),
                password=info.get("password", ""),
                key_path=info.get("key_path", ""),
            )
            ok, msg = conn.connect()
            if ok:
                self.connections[host_id] = conn
            return conn, ok, msg

    def disconnect(self, host_id: str):
        with self._lock:
            if host_id in self.connections:
                self.connections[host_id].disconnect()
                del self.connections[host_id]

    def get_connection(self, host_id: str) -> Optional[SSHConnection]:
        return self.connections.get(host_id)

    def close_all(self):
        with self._lock:
            for c in self.connections.values():
                c.disconnect()
            self.connections.clear()


def load_ai_config() -> dict:
    if AI_CONFIG_FILE.exists():
        try:
            with open(AI_CONFIG_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            pass
    return {
        "api_key": "",
        "base_url": "https://api.openai.com/v1",
        "model": "gpt-4o-mini",
        "system_prompt": "",
    }


def save_ai_config(cfg: dict):
    with open(AI_CONFIG_FILE, "w", encoding="utf-8") as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)