127 lines
3.7 KiB
Python
127 lines
3.7 KiB
Python
"""Secure API key storage for backend switching.
|
|
|
|
Reads keys from two sources, in priority order:
|
|
|
|
1. ``~/.xiaohe/agent/config.json`` under ``keys.*`` — the primary store for
|
|
myagents/xiaohe secrets.
|
|
2. ``~/.mytoolkit/config.json`` under ``keys.*`` — for backward compatibility
|
|
with keys already managed by mytoolkit.
|
|
|
|
Writes always go to ``~/.xiaohe/agent/config.json`` so that xiaohe-managed
|
|
keys shadow mytoolkit keys without modifying them.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import stat
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
XIAOHE_CONFIG_DIR = Path.home() / ".xiaohe" / "agent"
|
|
XIAOHE_CONFIG_PATH = XIAOHE_CONFIG_DIR / "config.json"
|
|
|
|
MYTOOLKIT_CONFIG_PATH = Path.home() / ".mytoolkit" / "config.json"
|
|
|
|
|
|
def _load_json(path: Path) -> dict[str, Any]:
|
|
"""Load JSON from path; return empty dict on missing/corrupt."""
|
|
try:
|
|
text = path.read_text(encoding="utf-8")
|
|
except (OSError, UnicodeDecodeError):
|
|
return {}
|
|
try:
|
|
data = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
def _atomic_write(path: Path, data: dict[str, Any]) -> None:
|
|
"""Write JSON atomically and restrict permissions on Unix."""
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
serialized = json.dumps(data, indent=2, ensure_ascii=False) + "\n"
|
|
|
|
fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.")
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
|
fh.write(serialized)
|
|
if os.name != "nt":
|
|
os.chmod(tmp, stat.S_IRUSR | stat.S_IWUSR)
|
|
os.replace(tmp, path)
|
|
except Exception:
|
|
try:
|
|
os.unlink(tmp)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def _read_key_from_config(path: Path, key_name: str) -> str | None:
|
|
"""Read a single key from a config file's ``keys`` section."""
|
|
data = _load_json(path)
|
|
value = data.get("keys", {}).get(key_name)
|
|
return value if isinstance(value, str) and value else None
|
|
|
|
|
|
def load_xiaohe_config() -> dict[str, Any]:
|
|
"""Load the full ``~/.xiaohe/agent/config.json``."""
|
|
return _load_json(XIAOHE_CONFIG_PATH)
|
|
|
|
|
|
def save_xiaohe_config(data: dict[str, Any]) -> None:
|
|
"""Persist the full ``~/.xiaohe/agent/config.json``."""
|
|
_atomic_write(XIAOHE_CONFIG_PATH, data)
|
|
|
|
|
|
def get_key(key_name: str) -> str | None:
|
|
"""Return a key, preferring the xiaohe store over mytoolkit.
|
|
|
|
Returns ``None`` when no key is stored or the stored value is empty.
|
|
"""
|
|
value = _read_key_from_config(XIAOHE_CONFIG_PATH, key_name)
|
|
if value:
|
|
return value
|
|
return _read_key_from_config(MYTOOLKIT_CONFIG_PATH, key_name)
|
|
|
|
|
|
def set_key(key_name: str, value: str) -> None:
|
|
"""Store a key in ``~/.xiaohe/agent/config.json``.
|
|
|
|
Empty values are rejected.
|
|
"""
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise ValueError("API key cannot be empty")
|
|
|
|
data = load_xiaohe_config()
|
|
if "keys" not in data or not isinstance(data["keys"], dict):
|
|
data["keys"] = {}
|
|
data["keys"][key_name] = value
|
|
save_xiaohe_config(data)
|
|
|
|
|
|
def remove_key(key_name: str) -> bool:
|
|
"""Remove a key from ``~/.xiaohe/agent/config.json``.
|
|
|
|
Returns ``True`` if the key existed and was removed.
|
|
"""
|
|
data = load_xiaohe_config()
|
|
keys = data.get("keys", {})
|
|
if not isinstance(keys, dict):
|
|
return False
|
|
if key_name not in keys:
|
|
return False
|
|
del keys[key_name]
|
|
if not keys:
|
|
data.pop("keys", None)
|
|
save_xiaohe_config(data)
|
|
return True
|
|
|
|
|
|
def has_key(key_name: str) -> bool:
|
|
"""Return whether a non-empty key exists in either store."""
|
|
return get_key(key_name) is not None
|