first commit

2026-09-05 21:08:13 +08:00
commit 7f8d810ccf
14 changed files with 781 additions and 0 deletions
@@ -0,0 +1,258 @@
+#!/usr/bin/env python3
+"""Generate an image through an OpenAI-compatible Images API."""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import binascii
+import json
+import os
+import shutil
+import subprocess
+from pathlib import Path
+import sys
+import tempfile
+import time
+from typing import Any, Callable
+from urllib.error import HTTPError, URLError
+from urllib.parse import urljoin
+from urllib.parse import urlsplit
+from urllib.request import Request, urlopen
+
+MAX_IMAGE_BYTES = 50 * 1024 * 1024
+DOWNLOAD_TIMEOUT_SECONDS = 120
+API_TIMEOUT_SECONDS = 180
+KEYRING_SERVICE = "codex-image-gen-custom"
+KEYRING_USER = "default"
+
+
+def config_dir() -> Path:
+    documents = os.environ.get("CODEX_DOCUMENTS") or str(Path.home() / "Documents")
+    return Path(documents) / "Codex" / ".shared" / "image-gen-custom"
+
+
+def _windows_saved_key_path() -> Path:
+    return config_dir() / "api_key.dpapi"
+
+
+def _decrypt_windows_saved_key(path: Path) -> str:
+    shell = shutil.which("powershell.exe") or shutil.which("pwsh")
+    if not shell:
+        raise RuntimeError("PowerShell is required to read the saved Windows API key")
+    script = r"""
+$ErrorActionPreference = 'Stop'
+Add-Type -AssemblyName System.Security
+$protected = [Convert]::FromBase64String((Get-Content -Raw -LiteralPath $env:CODEX_IMAGEGEN_KEY_PATH).Trim())
+$plain = [System.Security.Cryptography.ProtectedData]::Unprotect(
+  $protected,
+  $null,
+  [System.Security.Cryptography.DataProtectionScope]::LocalMachine
+)
+[Text.Encoding]::UTF8.GetString($plain)
+"""
+    env = os.environ.copy()
+    env["CODEX_IMAGEGEN_KEY_PATH"] = str(path)
+    command = [shell, "-NoProfile", "-NonInteractive", "-ExecutionPolicy", "Bypass", "-Command", script]
+    result = subprocess.run(command, check=False, capture_output=True, text=True, env=env)
+    if result.returncode != 0:
+        raise RuntimeError((result.stderr or "PowerShell could not decrypt the saved API key").strip())
+    return result.stdout.strip()
+
+
+def _secure_keyring():
+    try:
+        import keyring
+    except ImportError as exc:
+        raise RuntimeError("OPENAI_API_KEY is not set and the keyring package is unavailable") from exc
+    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 _field(item: Any, name: str) -> Any:
+    if isinstance(item, dict):
+        return item.get(name)
+    return getattr(item, name, None)
+
+
+def _validate_image(raw: bytes) -> None:
+    signatures = (
+        raw.startswith(b"\x89PNG\r\n\x1a\n"),
+        raw.startswith(b"\xff\xd8\xff"),
+        raw.startswith((b"GIF87a", b"GIF89a")),
+        len(raw) >= 12 and raw[:4] in (b"RIFF", b"RIFX") and raw[8:12] == b"WEBP",
+    )
+    if not raw or not any(signatures):
+        raise ValueError("Image API returned data without a supported image signature")
+
+
+def _download_image(url: str, *, opener: Callable[..., Any] = urlopen) -> bytes:
+    parsed = urlsplit(url)
+    if parsed.scheme not in ("http", "https") or not parsed.netloc:
+        raise ValueError("Image API returned an invalid image URL")
+    request = Request(url, headers={"User-Agent": "codex-image-gen-custom/1"})
+    with opener(request, timeout=DOWNLOAD_TIMEOUT_SECONDS) as response:
+        content_length = response.headers.get("Content-Length")
+        if content_length and int(content_length) > MAX_IMAGE_BYTES:
+            raise ValueError(f"Image URL exceeds the {MAX_IMAGE_BYTES}-byte limit")
+        raw = response.read(MAX_IMAGE_BYTES + 1)
+    if len(raw) > MAX_IMAGE_BYTES:
+        raise ValueError(f"Downloaded image exceeds the {MAX_IMAGE_BYTES}-byte limit")
+    return raw
+
+
+def extract_image_bytes(item: Any, *, opener: Callable[..., Any] = urlopen) -> bytes:
+    image_b64 = _field(item, "b64_json")
+    if image_b64:
+        try:
+            raw = base64.b64decode(image_b64, validate=True)
+        except (binascii.Error, TypeError, ValueError) as exc:
+            raise ValueError("Image API returned invalid b64_json") from exc
+    else:
+        image_url = _field(item, "url")
+        if not image_url:
+            raise ValueError("Image API response contained neither b64_json nor url")
+        raw = _download_image(str(image_url), opener=opener)
+    _validate_image(raw)
+    return raw
+
+
+def write_image(raw: bytes, output: Path, *, force: bool) -> None:
+    if output.exists() and not force:
+        raise FileExistsError(f"Output already exists: {output} (use --force to overwrite)")
+    output.parent.mkdir(parents=True, exist_ok=True)
+    fd, temporary_name = tempfile.mkstemp(prefix=f".{output.name}.", suffix=".tmp", dir=output.parent)
+    temporary = Path(temporary_name)
+    try:
+        with os.fdopen(fd, "wb") as stream:
+            stream.write(raw)
+            stream.flush()
+            os.fsync(stream.fileno())
+        temporary.replace(output)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def load_api_key() -> str:
+    environment_key = os.environ.get("OPENAI_API_KEY", "").strip()
+    if environment_key:
+        return environment_key
+    if os.name == "nt":
+        key_path = _windows_saved_key_path()
+        if not key_path.exists():
+            raise RuntimeError("No API key found. Run configure or set OPENAI_API_KEY")
+        key = _decrypt_windows_saved_key(key_path)
+        if not key:
+            raise RuntimeError("No API key found. Run configure or set OPENAI_API_KEY")
+        return key
+    try:
+        keyring = _secure_keyring()
+        key = keyring.get_password(KEYRING_SERVICE, KEYRING_USER)
+    except Exception as exc:
+        raise RuntimeError("The OS keyring is unavailable; set OPENAI_API_KEY for this run") from exc
+    if not key:
+        raise RuntimeError("No API key found. Run configure or set OPENAI_API_KEY")
+    return key
+
+
+def _configured_base_url() -> str:
+    environment_url = os.environ.get("OPENAI_BASE_URL", "").strip()
+    if environment_url:
+        return environment_url
+    config_path = config_dir() / "base_url.txt"
+    try:
+        configured_url = config_path.read_text(encoding="utf-8-sig").strip()
+    except OSError:
+        configured_url = ""
+    return configured_url
+
+
+def _base_url(override: str = "") -> str:
+    base_url = override.strip() or _configured_base_url() or "https://api.openai.com/v1"
+    base_url = base_url.rstrip("/")
+    if not base_url.lower().endswith("/v1"):
+        base_url += "/v1"
+    return base_url
+
+
+def _request_generation(*, api_key: str, base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
+    endpoint = urljoin(f"{base_url.rstrip('/')}/", "images/generations")
+    request = Request(
+        endpoint,
+        data=json.dumps(payload).encode("utf-8"),
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+            "User-Agent": "codex-image-gen-custom/1",
+        },
+        method="POST",
+    )
+    try:
+        with urlopen(request, timeout=API_TIMEOUT_SECONDS) as response:
+            raw = response.read()
+    except HTTPError as exc:
+        detail = exc.read(4096).decode("utf-8", errors="replace").strip()
+        suffix = f": {detail}" if detail else ""
+        raise RuntimeError(f"Image API returned HTTP {exc.code}{suffix}") from exc
+    except URLError as exc:
+        raise RuntimeError(f"Could not reach the Image API: {exc.reason}") from exc
+    try:
+        result = json.loads(raw.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise ValueError("Image API returned invalid JSON") from exc
+    if not isinstance(result, dict):
+        raise ValueError("Image API returned a non-object JSON response")
+    return result
+
+
+def generate(args: argparse.Namespace) -> None:
+    api_key = load_api_key()
+    base_url = _base_url(args.base_url)
+    started = time.time()
+    print("Calling Image API. This can take up to a couple of minutes.", file=sys.stderr)
+    result = _request_generation(
+        api_key=api_key,
+        base_url=base_url,
+        payload={
+            "model": args.model,
+            "prompt": args.prompt,
+            "size": args.size,
+            "quality": args.quality,
+        },
+    )
+    print(f"Generation completed in {time.time() - started:.1f}s.", file=sys.stderr)
+    data = result.get("data")
+    if not isinstance(data, list) or not data:
+        raise ValueError("Image API returned an empty data list")
+    write_image(extract_image_bytes(data[0]), Path(args.out), force=args.force)
+    print(f"Wrote {args.out}")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--base-url", default="", help="OpenAI-compatible API root")
+    parser.add_argument("--model", default="gpt-image-2")
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--size", default="1024x1024")
+    parser.add_argument("--quality", choices=("low", "medium", "high", "auto"), default="low")
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--force", action="store_true")
+    return parser
+
+
+def main() -> int:
+    try:
+        generate(build_parser().parse_args())
+    except Exception as exc:
+        print(f"Error: {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())