| |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import shutil |
| import subprocess |
| import sys |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from tempfile import NamedTemporaryFile |
| from urllib.request import Request, urlopen |
| from urllib.error import HTTPError, URLError |
|
|
| from codex_ws_auths import enable_codex_auth_websockets |
| from xai_auth_proxy import enforce_config_proxy_url, parse_proxy_urls, stable_proxy_for_auth |
|
|
|
|
| def env_required(name: str) -> str: |
| value = os.environ.get(name, "").strip() |
| if not value: |
| raise SystemExit(f"missing required env: {name}") |
| return value |
|
|
|
|
| MC_BIN = os.environ.get("MC_BIN", "mc") |
| ALIAS = os.environ.get("OBJECTSTORE_ALIAS", "daili-objectstore").strip() or "daili-objectstore" |
| ENDPOINT = env_required("OBJECTSTORE_ENDPOINT") |
| ACCESS_KEY = env_required("OBJECTSTORE_ACCESS_KEY") |
| SECRET_KEY = env_required("OBJECTSTORE_SECRET_KEY") |
| BUCKET = env_required("OBJECTSTORE_BUCKET") |
| ROOT = Path(env_required("OBJECTSTORE_ROOT")).resolve() |
| CONFIG_FALLBACK = os.environ.get("OBJECTSTORE_CONFIG_FALLBACK", "").strip() |
| REMOTE_CONFIG_KEY = "config/config.yaml" |
| REMOTE_AUTHS_PREFIX = "auths" |
| STATE_PATH = Path(os.environ.get("OBJECTSTORE_SYNC_STATE_PATH", "") or (ROOT / ".objectstore-sync-state.json")) |
| XAI_PROXY_URLS = parse_proxy_urls( |
| os.environ.get("XAI_PROXY_URLS", "") or "socks5h://127.0.0.1:9" |
| ) |
|
|
|
|
| def mc_env() -> dict: |
| env = os.environ.copy() |
| for key in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"): |
| env.pop(key, None) |
| return env |
|
|
|
|
| MC_TIMEOUT = int(os.environ.get("MC_TIMEOUT", "30")) |
|
|
|
|
| def run_mc(args: list[str], check: bool = True) -> subprocess.CompletedProcess[str]: |
| return subprocess.run( |
| [MC_BIN, *args], |
| check=check, |
| env=mc_env(), |
| capture_output=True, |
| text=True, |
| timeout=MC_TIMEOUT, |
| ) |
|
|
|
|
| def ensure_alias() -> None: |
| run_mc(["alias", "set", ALIAS, ENDPOINT, ACCESS_KEY, SECRET_KEY, "--api", "S3v4", "--path", "on"]) |
|
|
|
|
| def remote_path(key: str) -> str: |
| return f"{ALIAS}/{BUCKET}/{key}" |
|
|
|
|
| def parse_ts(value: str) -> float: |
| return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() |
|
|
|
|
| def local_path(rel: str) -> Path: |
| return ROOT / rel |
|
|
|
|
| def _is_auth_rel(rel: str) -> bool: |
| return rel.startswith(f"{REMOTE_AUTHS_PREFIX}/") and len(rel) > len(REMOTE_AUTHS_PREFIX) + 1 |
|
|
|
|
| def load_sync_state() -> tuple[set[str], dict[str, dict[str, str]]]: |
| try: |
| data = json.loads(STATE_PATH.read_text()) |
| except (OSError, ValueError): |
| return set(), {} |
| if not isinstance(data, dict): |
| return set(), {} |
| raw_rels = data.get("known_rels") |
| known_rels = {str(rel) for rel in raw_rels if isinstance(rel, str)} if isinstance(raw_rels, list) else set() |
| derived: dict[str, dict[str, str]] = {} |
| raw_derived = data.get("derived_from_remote") |
| if isinstance(raw_derived, dict): |
| for rel, value in raw_derived.items(): |
| if not isinstance(rel, str) or not isinstance(value, dict): |
| continue |
| remote_etag = value.get("remote_etag") |
| local_md5 = value.get("local_md5") |
| if isinstance(remote_etag, str) and isinstance(local_md5, str): |
| derived[rel] = { |
| "remote_etag": remote_etag, |
| "local_md5": local_md5, |
| } |
| return known_rels, derived |
|
|
|
|
| def load_known_rels() -> set[str]: |
| return load_sync_state()[0] |
|
|
|
|
| def save_known_rels( |
| inventory: dict[str, dict], |
| derived_from_remote: dict[str, dict[str, str]] | None = None, |
| ) -> None: |
| derived = derived_from_remote or {} |
| payload = { |
| "version": 2, |
| "updated_at": datetime.now(timezone.utc).isoformat(), |
| "known_rels": sorted(inventory), |
| "derived_from_remote": { |
| rel: derived[rel] |
| for rel in sorted(derived) |
| if rel in inventory |
| }, |
| } |
| STATE_PATH.parent.mkdir(parents=True, exist_ok=True) |
| with NamedTemporaryFile( |
| "w", |
| encoding="utf-8", |
| delete=False, |
| dir=str(STATE_PATH.parent), |
| ) as tmp: |
| json.dump(payload, tmp, ensure_ascii=False, sort_keys=True) |
| tmp.write("\n") |
| tmp_path = Path(tmp.name) |
| tmp_path.replace(STATE_PATH) |
|
|
|
|
| def touch_local(path: Path, ts: float) -> None: |
| os.utime(path, (ts, ts)) |
|
|
|
|
| def prune_empty_dirs(root: Path) -> None: |
| if not root.exists(): |
| return |
| for path in sorted((p for p in root.rglob("*") if p.is_dir()), reverse=True): |
| try: |
| path.rmdir() |
| except OSError: |
| pass |
|
|
|
|
| def file_md5(path: Path) -> str: |
| digest = hashlib.md5() |
| with path.open("rb") as handle: |
| while True: |
| chunk = handle.read(1024 * 1024) |
| if not chunk: |
| break |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def remote_inventory() -> dict[str, dict]: |
| inventory: dict[str, dict] = {} |
|
|
| config_proc = run_mc(["stat", "--json", remote_path(REMOTE_CONFIG_KEY)], check=False) |
| if config_proc.returncode == 0 and config_proc.stdout.strip(): |
| raw = json.loads(config_proc.stdout) |
| inventory[REMOTE_CONFIG_KEY] = { |
| "etag": raw["etag"], |
| "last_modified": parse_ts(raw["lastModified"]), |
| "remote": remote_path(REMOTE_CONFIG_KEY), |
| } |
|
|
| auths_proc = run_mc(["ls", "--json", "--recursive", remote_path(REMOTE_AUTHS_PREFIX)], check=False) |
| if auths_proc.returncode == 0: |
| for line in auths_proc.stdout.splitlines(): |
| if not line.strip(): |
| continue |
| raw = json.loads(line) |
| if raw.get("status") != "success" or raw.get("type") != "file": |
| continue |
| rel = f"{REMOTE_AUTHS_PREFIX}/{raw['key']}" |
| inventory[rel] = { |
| "etag": raw["etag"], |
| "last_modified": parse_ts(raw["lastModified"]), |
| "remote": remote_path(rel), |
| } |
|
|
| return inventory |
|
|
|
|
| def local_inventory() -> dict[str, dict]: |
| inventory: dict[str, dict] = {} |
|
|
| config_path = local_path(REMOTE_CONFIG_KEY) |
| if config_path.is_file(): |
| inventory[REMOTE_CONFIG_KEY] = { |
| "md5": file_md5(config_path), |
| "mtime": config_path.stat().st_mtime, |
| } |
|
|
| auth_root = local_path(REMOTE_AUTHS_PREFIX) |
| if auth_root.is_dir(): |
| for path in sorted(auth_root.rglob("*")): |
| if not path.is_file(): |
| continue |
| rel = path.relative_to(ROOT).as_posix() |
| inventory[rel] = { |
| "md5": file_md5(path), |
| "mtime": path.stat().st_mtime, |
| } |
|
|
| return inventory |
|
|
|
|
| def _enforce_downloaded_xai_proxy(path: Path, auth_name: str) -> bool: |
| try: |
| payload = json.loads(path.read_text(encoding="utf-8")) |
| except (OSError, UnicodeDecodeError, json.JSONDecodeError): |
| return False |
| if not isinstance(payload, dict) or str(payload.get("type") or "").strip().lower() != "xai": |
| return False |
|
|
| desired_proxy = stable_proxy_for_auth(auth_name, XAI_PROXY_URLS) |
| current_proxy = payload.get("proxy_url") |
| current_proxy = current_proxy.strip() if isinstance(current_proxy, str) else "" |
| changed = current_proxy != desired_proxy |
| if changed: |
| payload["proxy_url"] = desired_proxy |
| path.write_text( |
| json.dumps(payload, ensure_ascii=False, indent=2) + "\n", |
| encoding="utf-8", |
| ) |
| path.chmod(0o600) |
| return changed |
|
|
|
|
| def download_file(rel: str, meta: dict) -> bool: |
| dest = local_path(rel) |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| with NamedTemporaryFile(delete=False, dir=str(dest.parent)) as temporary: |
| temporary_path = Path(temporary.name) |
| temporary_path.unlink() |
| replaced = False |
| try: |
| run_mc(["cp", meta["remote"], str(temporary_path)]) |
| changed = False |
| if rel == REMOTE_CONFIG_KEY: |
| changed = enforce_config_proxy_url(temporary_path) |
| elif _is_auth_rel(rel): |
| changed = _enforce_downloaded_xai_proxy(temporary_path, dest.name) |
| touch_local(temporary_path, meta["last_modified"]) |
| temporary_path.replace(dest) |
| replaced = True |
| finally: |
| if not replaced and temporary_path.exists(): |
| temporary_path.unlink() |
| return changed |
|
|
|
|
| def remember_derived_file( |
| derived_from_remote: dict[str, dict[str, str]], |
| rel: str, |
| remote_etag: str, |
| ) -> None: |
| path = local_path(rel) |
| derived_from_remote[rel] = { |
| "remote_etag": remote_etag, |
| "local_md5": file_md5(path), |
| } |
|
|
|
|
| def enforce_config_without_mtime_change(path: Path) -> bool: |
| original_mtime = path.stat().st_mtime |
| changed = enforce_config_proxy_url(path) |
| if changed: |
| touch_local(path, original_mtime) |
| return changed |
|
|
|
|
| def upload_file(rel: str) -> None: |
| src = local_path(rel) |
| if not src.is_file(): |
| return |
| run_mc(["cp", str(src), remote_path(rel)]) |
|
|
|
|
| def delete_remote(rel: str) -> None: |
| run_mc(["rm", "--force", remote_path(rel)], check=False) |
|
|
|
|
| INVALID_STATUS_KEYWORDS = [ |
| "token_invalidated", |
| "token_revoked", |
| "invalidated oauth token", |
| "authentication token has been invalidated", |
| "account has been deactivated", |
| "no_organization", |
| "unauthorized", |
| "401", |
| ] |
|
|
|
|
| def _is_free_auth_file_name(file_name: str) -> bool: |
| return file_name.strip().endswith("-free.json") |
|
|
|
|
| def _free_auth_cleanup_enabled() -> bool: |
| return os.environ.get("CLEANUP_ENABLED", "0").strip() == "1" |
|
|
|
|
| def _fetch_invalid_auth_names() -> set[str]: |
| """Query local CLIProxyAPI management API for invalid free auth names. |
| |
| Team auth lifecycle is managed externally; runtime sync must not evict team |
| auths for transient states such as cooldown or temporary unavailability. |
| """ |
| if not _free_auth_cleanup_enabled(): |
| return set() |
|
|
| mgmt_key = os.environ.get("MANAGEMENT_PASSWORD") or os.environ.get("API_KEY") or "" |
| port = os.environ.get("PORT", "8317") |
| if not mgmt_key: |
| return set() |
| url = f"http://127.0.0.1:{port}/v0/management/auth-files" |
| req = Request(url) |
| req.add_header("Authorization", f"Bearer {mgmt_key}") |
| try: |
| with urlopen(req, timeout=10) as resp: |
| data = json.loads(resp.read()) |
| except (HTTPError, URLError, OSError, ValueError): |
| return set() |
| if not isinstance(data, dict): |
| return set() |
| files = data.get("files", []) |
| if not isinstance(files, list): |
| return set() |
|
|
| invalid_names: set[str] = set() |
| now = datetime.now(timezone.utc).isoformat() |
| for entry in files: |
| if not isinstance(entry, dict): |
| continue |
| name = str(entry.get("name") or "").strip() |
| if not name: |
| continue |
| if not _is_free_auth_file_name(name): |
| continue |
| status = str(entry.get("status") or "").strip().lower() |
| status_message = str(entry.get("status_message") or "").strip().lower() |
| unavailable = entry.get("unavailable", False) |
| |
| next_retry = str(entry.get("next_retry_after") or "").strip() |
| in_cooldown = bool(next_retry and next_retry > now) |
| if status == "error" or unavailable or in_cooldown: |
| invalid_names.add(name) |
| continue |
| if status_message: |
| for kw in INVALID_STATUS_KEYWORDS: |
| if kw in status_message: |
| invalid_names.add(name) |
| break |
| return invalid_names |
|
|
|
|
| def restore() -> None: |
| ensure_alias() |
| ROOT.mkdir(parents=True, exist_ok=True) |
| local_path("config").mkdir(parents=True, exist_ok=True) |
| shutil.rmtree(local_path(REMOTE_AUTHS_PREFIX), ignore_errors=True) |
| local_path(REMOTE_AUTHS_PREFIX).mkdir(parents=True, exist_ok=True) |
|
|
| remote = remote_inventory() |
| derived_from_remote: dict[str, dict[str, str]] = {} |
| if REMOTE_CONFIG_KEY in remote: |
| if download_file(REMOTE_CONFIG_KEY, remote[REMOTE_CONFIG_KEY]): |
| remember_derived_file( |
| derived_from_remote, |
| REMOTE_CONFIG_KEY, |
| remote[REMOTE_CONFIG_KEY]["etag"], |
| ) |
| elif CONFIG_FALLBACK: |
| fallback = Path(CONFIG_FALLBACK) |
| if fallback.is_file(): |
| target = local_path(REMOTE_CONFIG_KEY) |
| target.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copy2(fallback, target) |
|
|
| run_mc(["cp", "--recursive", remote_path(f"{REMOTE_AUTHS_PREFIX}/"), str(local_path(REMOTE_AUTHS_PREFIX))], check=False) |
| for rel, meta in remote.items(): |
| if not _is_auth_rel(rel): |
| continue |
| path = local_path(rel) |
| if not path.is_file(): |
| continue |
| changed = _enforce_downloaded_xai_proxy(path, path.name) |
| touch_local(path, meta["last_modified"]) |
| if changed: |
| remember_derived_file(derived_from_remote, rel, meta["etag"]) |
| enable_codex_auth_websockets(local_path(REMOTE_AUTHS_PREFIX)) |
| config_path = local_path(REMOTE_CONFIG_KEY) |
| if config_path.is_file(): |
| enforce_config_without_mtime_change(config_path) |
| save_known_rels(local_inventory(), derived_from_remote) |
|
|
|
|
| def sync() -> None: |
| ensure_alias() |
| ROOT.mkdir(parents=True, exist_ok=True) |
| enable_codex_auth_websockets(local_path(REMOTE_AUTHS_PREFIX)) |
| config_path = local_path(REMOTE_CONFIG_KEY) |
| if config_path.is_file(): |
| enforce_config_without_mtime_change(config_path) |
|
|
| invalid_names = _fetch_invalid_auth_names() |
| known_rels, derived_from_remote = load_sync_state() |
|
|
| remote = remote_inventory() |
| local = local_inventory() |
|
|
| |
| |
| |
| |
| |
| |
| for rel, meta in remote.items(): |
| |
| file_name = rel.rsplit("/", 1)[-1] if "/" in rel else rel |
| if file_name in invalid_names: |
| derived_from_remote.pop(rel, None) |
| delete_remote(rel) |
| dest = local_path(rel) |
| if dest.is_file(): |
| dest.unlink() |
| continue |
|
|
| local_meta = local.get(rel) |
| if local_meta is None: |
| if _is_auth_rel(rel) and rel in known_rels: |
| derived_from_remote.pop(rel, None) |
| delete_remote(rel) |
| continue |
| if download_file(rel, meta): |
| remember_derived_file(derived_from_remote, rel, meta["etag"]) |
| else: |
| derived_from_remote.pop(rel, None) |
| continue |
| if local_meta["md5"] == meta["etag"]: |
| derived_from_remote.pop(rel, None) |
| continue |
| derived = derived_from_remote.get(rel) |
| if ( |
| derived is not None |
| and derived.get("remote_etag") == meta["etag"] |
| and derived.get("local_md5") == local_meta["md5"] |
| ): |
| continue |
| if local_meta["mtime"] > meta["last_modified"]: |
| derived_from_remote.pop(rel, None) |
| upload_file(rel) |
| else: |
| if download_file(rel, meta): |
| remember_derived_file(derived_from_remote, rel, meta["etag"]) |
| else: |
| derived_from_remote.pop(rel, None) |
|
|
| |
| for rel in sorted(set(local) - set(remote)): |
| file_name = rel.rsplit("/", 1)[-1] if "/" in rel else rel |
| if file_name in invalid_names: |
| derived_from_remote.pop(rel, None) |
| dest = local_path(rel) |
| if dest.is_file(): |
| dest.unlink() |
| continue |
| derived_from_remote.pop(rel, None) |
| upload_file(rel) |
|
|
| if config_path.is_file(): |
| enforce_config_without_mtime_change(config_path) |
| prune_empty_dirs(local_path(REMOTE_AUTHS_PREFIX)) |
| save_known_rels(local_inventory(), derived_from_remote) |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("mode", choices=("restore", "sync")) |
| args = parser.parse_args() |
|
|
| if args.mode == "restore": |
| restore() |
| else: |
| sync() |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
