97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Store or remove the API key using a protected local store."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
SERVICE = "codex-image-gen-custom"
|
|
|
|
|
|
def config_dir() -> Path:
|
|
documents = os.environ.get("CODEX_DOCUMENTS") or str(Path.home() / "Documents")
|
|
return Path(documents) / "Codex" / ".shared" / "image-gen-custom"
|
|
|
|
|
|
def base_url_path() -> Path:
|
|
return config_dir() / "base_url.txt"
|
|
|
|
|
|
def api_key_path() -> Path:
|
|
return config_dir() / "api_key.xml"
|
|
|
|
|
|
def store_api_key(key: str, user: str) -> None:
|
|
if os.name == "nt":
|
|
raise RuntimeError("Use configure.ps1 on Windows to store the API key securely.")
|
|
keyring = secure_keyring()
|
|
keyring.set_password(SERVICE, user, key)
|
|
|
|
|
|
def delete_api_key(user: str) -> None:
|
|
if os.name == "nt":
|
|
api_key_path().unlink(missing_ok=True)
|
|
return
|
|
keyring = secure_keyring()
|
|
keyring.delete_password(SERVICE, user)
|
|
|
|
|
|
def load_saved_api_key(user: str) -> str:
|
|
if os.name == "nt":
|
|
raise RuntimeError("Use generate-image.ps1 on Windows to read the saved API key securely.")
|
|
keyring = secure_keyring()
|
|
return (keyring.get_password(SERVICE, user) or "").strip()
|
|
|
|
|
|
def secure_keyring():
|
|
import keyring
|
|
|
|
backend = keyring.get_keyring()
|
|
backend_id = f"{backend.__class__.__module__}.{backend.__class__.__name__}".lower()
|
|
allowed = ("windows", "macos", "secretservice", "libsecret", "kwallet")
|
|
if not any(marker in backend_id for marker in allowed):
|
|
raise RuntimeError("The configured keyring is not a recognized protected OS credential store")
|
|
return keyring
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--user", default=os.environ.get("IMAGE_GEN_KEYRING_USER", "default"))
|
|
parser.add_argument("--delete", action="store_true")
|
|
parser.add_argument("--base-url", default="")
|
|
parser.add_argument("--delete-base-url", action="store_true")
|
|
args = parser.parse_args()
|
|
if args.delete_base_url:
|
|
base_url_path().unlink(missing_ok=True)
|
|
print("API base URL removed from the local configuration.")
|
|
return 0
|
|
if args.base_url:
|
|
path = base_url_path()
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(args.base_url.rstrip("/") + "\n", encoding="utf-8")
|
|
print("API base URL stored in the local configuration.")
|
|
return 0
|
|
try:
|
|
if args.delete:
|
|
delete_api_key(args.user)
|
|
print("API key removed from the local secret store.")
|
|
else:
|
|
key = input("OpenAI-compatible API key: ").strip()
|
|
if not key:
|
|
raise ValueError("API key cannot be empty")
|
|
store_api_key(key, args.user)
|
|
if load_saved_api_key(args.user) != key:
|
|
raise RuntimeError("Stored key could not be read back")
|
|
print("API key stored in the local secret store.")
|
|
except Exception as exc:
|
|
print(f"Error: {exc}. Set OPENAI_API_KEY for a one-off run or run configure again.", file=sys.stderr)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|