image-gen-custom-skill/skill/image-gen-custom/scripts/generate_image.py

#!/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())