@agxnte/reidx 2.0.2

package/src/reidx/provider_manager/database.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import json
+import os
+import uuid
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+
+from reidx.diagnostics.logger import get_logger
+from reidx.provider_manager import keychain
+
+log = get_logger("reidx.provider_manager.database")
+
+SCHEMA_VERSION = 1
+DB_FILENAME = "providers.db"
+LEGACY_FILENAME = "providers.json"
+
+
+@dataclass
+class StoredKey:
+    id: str
+    label: str
+    encrypted_key: str
+
+    def decrypt(self) -> str:
+        return keychain.decrypt(self.encrypted_key)
+
+
+@dataclass
+class StoredProvider:
+    name: str
+    kind: str
+    base_url: str = ""
+    default_model: str = ""
+    auth_method: str = "bearer"
+    extra_headers: dict[str, str] = field(default_factory=dict)
+    catalog_id: str | None = None
+    keys: list[StoredKey] = field(default_factory=list)
+    active_key_id: str | None = None
+
+    def active_key(self) -> StoredKey | None:
+        if not self.active_key_id:
+            return self.keys[0] if self.keys else None
+        for k in self.keys:
+            if k.id == self.active_key_id:
+                return k
+        return self.keys[0] if self.keys else None
+
+    def decrypted_api_key(self) -> str:
+        k = self.active_key()
+        return k.decrypt() if k else ""
+
+
+class ProviderDatabase:
+    def __init__(self, storage_root: Path) -> None:
+        self.path = Path(storage_root) / DB_FILENAME
+        self.legacy_path = Path(storage_root) / LEGACY_FILENAME
+        self._migrate_legacy_if_needed()
+
+    def _migrate_legacy_if_needed(self) -> None:
+        if self.path.exists() or not self.legacy_path.exists():
+            return
+        try:
+            data = json.loads(self.legacy_path.read_text(encoding="utf-8"))
+        except (OSError, ValueError):
+            log.warning("failed to read legacy providers.json; skipping migration")
+            return
+        records = data.get("providers", [])
+        if not records:
+            return
+        log.info("migrating %d provider(s) from legacy providers.json", len(records))
+        providers: list[dict] = []
+        for entry in records:
+            name = entry.get("name", "")
+            if not name:
+                continue
+            api_key = entry.get("api_key", "")
+            keys: list[dict] = []
+            active_key_id = None
+            if api_key:
+                kid = uuid.uuid4().hex[:12]
+                keys.append({
+                    "id": kid,
+                    "label": "Imported",
+                    "encrypted_key": keychain.encrypt(api_key),
+                })
+                active_key_id = kid
+            providers.append({
+                "name": name,
+                "kind": entry.get("kind", "openai-compatible"),
+                "base_url": entry.get("base_url", ""),
+                "default_model": entry.get("default_model", ""),
+                "auth_method": "bearer",
+                "extra_headers": {},
+                "catalog_id": None,
+                "keys": keys,
+                "active_key_id": active_key_id,
+            })
+        self._write_raw({"version": SCHEMA_VERSION, "providers": providers})
+        try:
+            os.rename(self.legacy_path, self.legacy_path.with_suffix(".json.bak"))
+        except OSError:
+            pass
+        log.info("migration complete; legacy file backed up")
+
+    def _read_raw(self) -> dict:
+        if not self.path.exists():
+            return {"version": SCHEMA_VERSION, "providers": []}
+        try:
+            data = json.loads(self.path.read_text(encoding="utf-8"))
+        except (OSError, ValueError):
+            log.exception("failed to read providers.db; treating as empty")
+            return {"version": SCHEMA_VERSION, "providers": []}
+        if data.get("version", 0) > SCHEMA_VERSION:
+            log.warning("providers.db version %s > supported %s; some features may not work",
+                        data.get("version"), SCHEMA_VERSION)
+        return data
+
+    def _write_raw(self, data: dict) -> None:
+        keychain.secure_write(self.path, json.dumps(data, indent=2))
+
+    def _to_stored(self, entry: dict) -> StoredProvider:
+        keys = [StoredKey(
+            id=k.get("id", ""),
+            label=k.get("label", ""),
+            encrypted_key=k.get("encrypted_key", ""),
+        ) for k in entry.get("keys", [])]
+        return StoredProvider(
+            name=entry.get("name", ""),
+            kind=entry.get("kind", "openai-compatible"),
+            base_url=entry.get("base_url", ""),
+            default_model=entry.get("default_model", ""),
+            auth_method=entry.get("auth_method", "bearer"),
+            extra_headers=entry.get("extra_headers", {}),
+            catalog_id=entry.get("catalog_id"),
+            keys=keys,
+            active_key_id=entry.get("active_key_id"),
+        )
+
+    def _to_dict(self, p: StoredProvider) -> dict:
+        return {
+            "name": p.name,
+            "kind": p.kind,
+            "base_url": p.base_url,
+            "default_model": p.default_model,
+            "auth_method": p.auth_method,
+            "extra_headers": p.extra_headers,
+            "catalog_id": p.catalog_id,
+            "keys": [asdict(k) for k in p.keys],
+            "active_key_id": p.active_key_id,
+        }
+
+    def list_providers(self) -> list[StoredProvider]:
+        data = self._read_raw()
+        return [self._to_stored(e) for e in data.get("providers", [])]
+
+    def get_provider(self, name: str) -> StoredProvider | None:
+        for p in self.list_providers():
+            if p.name == name:
+                return p
+        return None
+
+    def save_provider(self, provider: StoredProvider) -> None:
+        data = self._read_raw()
+        providers = data.get("providers", [])
+        providers = [e for e in providers if e.get("name") != provider.name]
+        providers.append(self._to_dict(provider))
+        data["providers"] = providers
+        self._write_raw(data)
+
+    def remove_provider(self, name: str) -> bool:
+        data = self._read_raw()
+        providers = data.get("providers", [])
+        remaining = [e for e in providers if e.get("name") != name]
+        if len(remaining) == len(providers):
+            return False
+        data["providers"] = remaining
+        self._write_raw(data)
+        return True
+
+    def add_key(self, provider_name: str, label: str, api_key: str) -> StoredKey | None:
+        p = self.get_provider(provider_name)
+        if p is None:
+            return None
+        key = StoredKey(
+            id=uuid.uuid4().hex[:12],
+            label=label,
+            encrypted_key=keychain.encrypt(api_key),
+        )
+        p.keys.append(key)
+        if p.active_key_id is None:
+            p.active_key_id = key.id
+        self.save_provider(p)
+        return key
+
+    def remove_key(self, provider_name: str, key_id: str) -> bool:
+        p = self.get_provider(provider_name)
+        if p is None:
+            return False
+        before = len(p.keys)
+        p.keys = [k for k in p.keys if k.id != key_id]
+        if len(p.keys) == before:
+            return False
+        if p.active_key_id == key_id:
+            p.active_key_id = p.keys[0].id if p.keys else None
+        self.save_provider(p)
+        return True
+
+    def rename_key(self, provider_name: str, key_id: str, new_label: str) -> bool:
+        p = self.get_provider(provider_name)
+        if p is None:
+            return False
+        for k in p.keys:
+            if k.id == key_id:
+                k.label = new_label
+                self.save_provider(p)
+                return True
+        return False
+
+    def set_active_key(self, provider_name: str, key_id: str) -> bool:
+        p = self.get_provider(provider_name)
+        if p is None:
+            return False
+        if not any(k.id == key_id for k in p.keys):
+            return False
+        p.active_key_id = key_id
+        self.save_provider(p)
+        return True
+
+    def get_decrypted_key(self, provider_name: str) -> str:
+        p = self.get_provider(provider_name)
+        if p is None:
+            return ""
+        return p.decrypted_api_key()

package/src/reidx/provider_manager/keychain.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import base64
+import os
+import sys
+from pathlib import Path
+
+from reidx.diagnostics.logger import get_logger
+
+log = get_logger("reidx.provider_manager.keychain")
+
+_IS_WINDOWS = sys.platform == "win32"
+
+if _IS_WINDOWS:
+    import ctypes
+    import ctypes.wintypes
+
+    class _DATA_BLOB(ctypes.Structure):
+        _fields_ = [
+            ("cbData", ctypes.wintypes.DWORD),
+            ("pbData", ctypes.POINTER(ctypes.c_char)),
+        ]
+
+    def _dpapi_encrypt(plaintext: bytes) -> bytes:
+        blob_in = _DATA_BLOB()
+        blob_in.pbData = ctypes.cast(
+            ctypes.create_string_buffer(plaintext, len(plaintext)),
+            ctypes.POINTER(ctypes.c_char),
+        )
+        blob_in.cbData = len(plaintext)
+        blob_out = _DATA_BLOB()
+        ok = ctypes.windll.crypt32.CryptProtectData(
+            ctypes.byref(blob_in),
+            None,
+            None,
+            None,
+            None,
+            0,
+            ctypes.byref(blob_out),
+        )
+        if not ok:
+            raise OSError("DPAPI CryptProtectData failed")
+        try:
+            return ctypes.string_at(blob_out.pbData, blob_out.cbData)
+        finally:
+            ctypes.windll.kernel32.LocalFree(blob_out.pbData)
+
+    def _dpapi_decrypt(ciphertext: bytes) -> bytes:
+        blob_in = _DATA_BLOB()
+        blob_in.pbData = ctypes.cast(
+            ctypes.create_string_buffer(ciphertext, len(ciphertext)),
+            ctypes.POINTER(ctypes.c_char),
+        )
+        blob_in.cbData = len(ciphertext)
+        blob_out = _DATA_BLOB()
+        ok = ctypes.windll.crypt32.CryptUnprotectData(
+            ctypes.byref(blob_in),
+            None,
+            None,
+            None,
+            None,
+            0,
+            ctypes.byref(blob_out),
+        )
+        if not ok:
+            raise OSError("DPAPI CryptUnprotectData failed")
+        try:
+            return ctypes.string_at(blob_out.pbData, blob_out.cbData)
+        finally:
+            ctypes.windll.kernel32.LocalFree(blob_out.pbData)
+
+
+def encrypt(plaintext: str) -> str:
+    if not plaintext:
+        return ""
+    raw = plaintext.encode("utf-8")
+    if _IS_WINDOWS:
+        encrypted = _dpapi_encrypt(raw)
+    else:
+        log.warning("DPAPI not available on this platform; storing key without OS encryption")
+        encrypted = raw
+    return base64.b64encode(encrypted).decode("ascii")
+
+
+def decrypt(ciphertext: str) -> str:
+    if not ciphertext:
+        return ""
+    raw = base64.b64decode(ciphertext)
+    if _IS_WINDOWS:
+        decrypted = _dpapi_decrypt(raw)
+    else:
+        decrypted = raw
+    return decrypted.decode("utf-8")
+
+
+def secure_write(path: Path, data: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(data, encoding="utf-8")
+    try:
+        os.chmod(path, 0o600)
+    except OSError:
+        pass