""" 连接管理器:保存多个 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)