Screenshot

scripts/take_screenshot.py

#!/usr/bin/env python3
"""Cross-platform screenshot helper for Codex skills."""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import platform
import shutil
import subprocess
import tempfile
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
MAC_PERM_SCRIPT = SCRIPT_DIR / "macos_permissions.swift"
MAC_PERM_HELPER = SCRIPT_DIR / "ensure_macos_permissions.sh"
MAC_WINDOW_SCRIPT = SCRIPT_DIR / "macos_window_info.swift"
MAC_DISPLAY_SCRIPT = SCRIPT_DIR / "macos_display_info.swift"
TEST_MODE_ENV = "CODEX_SCREENSHOT_TEST_MODE"
TEST_PLATFORM_ENV = "CODEX_SCREENSHOT_TEST_PLATFORM"
TEST_WINDOWS_ENV = "CODEX_SCREENSHOT_TEST_WINDOWS"
TEST_DISPLAYS_ENV = "CODEX_SCREENSHOT_TEST_DISPLAYS"
TEST_PNG = (
    b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
    b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\x0cIDAT\x08\xd7c"
    b"\xf8\xff\xff?\x00\x05\xfe\x02\xfeA\xad\x1c\x1c\x00\x00\x00\x00IEND"
    b"\xaeB`\x82"
)


def parse_region(value: str) -> tuple[int, int, int, int]:
    parts = [p.strip() for p in value.split(",")]
    if len(parts) != 4:
        raise argparse.ArgumentTypeError("region must be x,y,w,h")
    try:
        x, y, w, h = (int(p) for p in parts)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("region values must be integers") from exc
    if w <= 0 or h <= 0:
        raise argparse.ArgumentTypeError("region width and height must be positive")
    return x, y, w, h


def test_mode_enabled() -> bool:
    value = os.environ.get(TEST_MODE_ENV, "")
    return value.lower() in {"1", "true", "yes", "on"}


def normalize_platform(value: str) -> str:
    lowered = value.strip().lower()
    if lowered in {"darwin", "mac", "macos", "osx"}:
        return "Darwin"
    if lowered in {"linux", "ubuntu"}:
        return "Linux"
    if lowered in {"windows", "win"}:
        return "Windows"
    return value


def test_platform_override() -> str | None:
    value = os.environ.get(TEST_PLATFORM_ENV)
    if value:
        return normalize_platform(value)
    return None


def parse_int_list(value: str) -> list[int]:
    results: list[int] = []
    for part in value.split(","):
        part = part.strip()
        if not part:
            continue
        try:
            results.append(int(part))
        except ValueError:
            continue
    return results


def test_window_ids() -> list[int]:
    value = os.environ.get(TEST_WINDOWS_ENV, "101,102")
    ids = parse_int_list(value)
    return ids or [101]


def test_display_ids() -> list[int]:
    value = os.environ.get(TEST_DISPLAYS_ENV, "1,2")
    ids = parse_int_list(value)
    return ids or [1]


def write_test_png(path: Path) -> None:
    ensure_parent(path)
    path.write_bytes(TEST_PNG)


def timestamp() -> str:
    return dt.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")


def default_filename(fmt: str, prefix: str = "screenshot") -> str:
    return f"{prefix}-{timestamp()}.{fmt}"


def mac_default_dir() -> Path:
    desktop = Path.home() / "Desktop"
    try:
        proc = subprocess.run(
            ["defaults", "read", "com.apple.screencapture", "location"],
            check=False,
            capture_output=True,
            text=True,
        )
        location = proc.stdout.strip()
        if location:
            return Path(location).expanduser()
    except OSError:
        pass
    return desktop


def default_dir(system: str) -> Path:
    home = Path.home()
    if system == "Darwin":
        return mac_default_dir()
    if system == "Windows":
        pictures = home / "Pictures"
        screenshots = pictures / "Screenshots"
        if screenshots.exists():
            return screenshots
        if pictures.exists():
            return pictures
        return home
    pictures = home / "Pictures"
    screenshots = pictures / "Screenshots"
    if screenshots.exists():
        return screenshots
    if pictures.exists():
        return pictures
    return home


def ensure_parent(path: Path) -> None:
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
    except OSError:
        # Fall back to letting the capture command report a clearer error.
        pass


def resolve_output_path(
    requested_path: str | None, mode: str, fmt: str, system: str
) -> Path:
    if requested_path:
        path = Path(requested_path).expanduser()
        if path.exists() and path.is_dir():
            path = path / default_filename(fmt)
        elif requested_path.endswith(("/", "\\")) and not path.exists():
            path.mkdir(parents=True, exist_ok=True)
            path = path / default_filename(fmt)
        elif path.suffix == "":
            path = path.with_suffix(f".{fmt}")
        ensure_parent(path)
        return path

    if mode == "temp":
        tmp_dir = Path(tempfile.gettempdir())
        tmp_path = tmp_dir / default_filename(fmt, prefix="codex-shot")
        ensure_parent(tmp_path)
        return tmp_path

    dest_dir = default_dir(system)
    dest_path = dest_dir / default_filename(fmt)
    ensure_parent(dest_path)
    return dest_path


def multi_output_paths(base: Path, suffixes: list[str]) -> list[Path]:
    if len(suffixes) <= 1:
        return [base]
    paths: list[Path] = []
    for suffix in suffixes:
        candidate = base.with_name(f"{base.stem}-{suffix}{base.suffix}")
        ensure_parent(candidate)
        paths.append(candidate)
    return paths


def run(cmd: list[str]) -> None:
    try:
        subprocess.run(cmd, check=True)
    except FileNotFoundError as exc:
        raise SystemExit(f"required command not found: {cmd[0]}") from exc
    except subprocess.CalledProcessError as exc:
        raise SystemExit(f"command failed ({exc.returncode}): {' '.join(cmd)}") from exc


def swift_json(script: Path, extra_args: list[str] | None = None) -> dict:
    module_cache = Path(tempfile.gettempdir()) / "codex-swift-module-cache"
    module_cache.mkdir(parents=True, exist_ok=True)
    cmd = ["swift", "-module-cache-path", str(module_cache), str(script)]
    if extra_args:
        cmd.extend(extra_args)
    try:
        proc = subprocess.run(cmd, check=True, capture_output=True, text=True)
    except FileNotFoundError as exc:
        raise SystemExit("swift not found; install Xcode command line tools") from exc
    except subprocess.CalledProcessError as exc:
        stderr = (exc.stderr or "").strip()
        if "ModuleCache" in stderr and "Operation not permitted" in stderr:
            raise SystemExit(
                "swift needs module-cache access; rerun with escalated permissions"
            ) from exc
        msg = stderr or (exc.stdout or "").strip() or "swift helper failed"
        raise SystemExit(msg) from exc
    try:
        return json.loads(proc.stdout)
    except json.JSONDecodeError as exc:
        raise SystemExit(f"swift helper returned invalid JSON: {proc.stdout.strip()}") from exc


def macos_screen_capture_granted(request: bool = False) -> bool:
    args = ["--request"] if request else []
    payload = swift_json(MAC_PERM_SCRIPT, args)
    return bool(payload.get("screenCapture"))


def ensure_macos_permissions() -> None:
    if os.environ.get("CODEX_SANDBOX"):
        raise SystemExit(
            "screen capture checks are blocked in the sandbox; rerun with escalated permissions"
        )
    if macos_screen_capture_granted():
        return
    subprocess.run(["bash", str(MAC_PERM_HELPER)], check=False)
    if not macos_screen_capture_granted():
        raise SystemExit(
            "Screen Recording permission is required; enable it in System Settings and retry"
        )


def activate_app(app: str) -> None:
    safe_app = app.replace('"', '\\"')
    script = f'tell application "{safe_app}" to activate'
    subprocess.run(["osascript", "-e", script], check=False, capture_output=True, text=True)


def macos_window_payload(args: argparse.Namespace, frontmost: bool, include_list: bool) -> dict:
    flags: list[str] = []
    if frontmost:
        flags.append("--frontmost")
    if args.app:
        flags.extend(["--app", args.app])
    if args.window_name:
        flags.extend(["--window-name", args.window_name])
    if include_list:
        flags.append("--list")
    return swift_json(MAC_WINDOW_SCRIPT, flags)


def macos_display_indexes() -> list[int]:
    payload = swift_json(MAC_DISPLAY_SCRIPT)
    displays = payload.get("displays") or []
    indexes: list[int] = []
    for item in displays:
        try:
            value = int(item)
        except (TypeError, ValueError):
            continue
        if value > 0:
            indexes.append(value)
    return indexes or [1]


def macos_window_ids(args: argparse.Namespace, capture_all: bool) -> list[int]:
    payload = macos_window_payload(
        args,
        frontmost=args.active_window,
        include_list=capture_all,
    )
    if capture_all:
        windows = payload.get("windows") or []
        ids: list[int] = []
        for item in windows:
            win_id = item.get("id")
            if win_id is None:
                continue
            try:
                ids.append(int(win_id))
            except (TypeError, ValueError):
                continue
        if ids:
            return ids
    selected = payload.get("selected") or {}
    win_id = selected.get("id")
    if win_id is not None:
        try:
            return [int(win_id)]
        except (TypeError, ValueError):
            pass
    raise SystemExit("no matching macOS window found; try --list-windows to inspect ids")


def list_macos_windows(args: argparse.Namespace) -> None:
    payload = macos_window_payload(args, frontmost=args.active_window, include_list=True)
    windows = payload.get("windows") or []
    if not windows:
        print("no matching windows found")
        return
    for item in windows:
        bounds = item.get("bounds") or {}
        name = item.get("name") or ""
        width = bounds.get("width", 0)
        height = bounds.get("height", 0)
        x = bounds.get("x", 0)
        y = bounds.get("y", 0)
        print(f"{item.get('id')}\t{item.get('owner')}\t{name}\t{width}x{height}+{x}+{y}")


def list_test_macos_windows(args: argparse.Namespace) -> None:
    owner = args.app or "TestApp"
    name = args.window_name or ""
    ids = test_window_ids()
    if args.active_window and ids:
        ids = [ids[0]]
    for idx, win_id in enumerate(ids, start=1):
        window_name = name or f"Window {idx}"
        print(f"{win_id}\t{owner}\t{window_name}\t800x600+0+0")


def resolve_macos_windows(args: argparse.Namespace) -> list[int]:
    if args.app:
        activate_app(args.app)
    capture_all = not args.active_window
    return macos_window_ids(args, capture_all=capture_all)


def resolve_test_macos_windows(args: argparse.Namespace) -> list[int]:
    ids = test_window_ids()
    if args.active_window and ids:
        return [ids[0]]
    return ids


def capture_macos(
    args: argparse.Namespace,
    output: Path,
    *,
    window_id: int | None = None,
    display: int | None = None,
) -> None:
    cmd = ["screencapture", "-x", f"-t{args.format}"]
    if args.interactive:
        cmd.append("-i")
    if display is not None:
        cmd.append(f"-D{display}")
    effective_window_id = window_id if window_id is not None else args.window_id
    if effective_window_id is not None:
        cmd.append(f"-l{effective_window_id}")
    elif args.region is not None:
        x, y, w, h = args.region
        cmd.append(f"-R{x},{y},{w},{h}")
    cmd.append(str(output))
    run(cmd)


def capture_linux(args: argparse.Namespace, output: Path) -> None:
    scrot = shutil.which("scrot")
    gnome = shutil.which("gnome-screenshot")
    imagemagick = shutil.which("import")
    xdotool = shutil.which("xdotool")

    if args.region is not None:
        x, y, w, h = args.region
        if scrot:
            run(["scrot", "-a", f"{x},{y},{w},{h}", str(output)])
            return
        if imagemagick:
            geometry = f"{w}x{h}+{x}+{y}"
            run(["import", "-window", "root", "-crop", geometry, str(output)])
            return
        raise SystemExit("region capture requires scrot or ImageMagick (import)")

    if args.window_id is not None:
        if imagemagick:
            run(["import", "-window", str(args.window_id), str(output)])
            return
        raise SystemExit("window-id capture requires ImageMagick (import)")

    if args.active_window:
        if scrot:
            run(["scrot", "-u", str(output)])
            return
        if gnome:
            run(["gnome-screenshot", "-w", "-f", str(output)])
            return
        if imagemagick and xdotool:
            win_id = (
                subprocess.check_output(["xdotool", "getactivewindow"], text=True)
                .strip()
            )
            run(["import", "-window", win_id, str(output)])
            return
        raise SystemExit("active-window capture requires scrot, gnome-screenshot, or import+xdotool")

    if scrot:
        run(["scrot", str(output)])
        return
    if gnome:
        run(["gnome-screenshot", "-f", str(output)])
        return
    if imagemagick:
        run(["import", "-window", "root", str(output)])
        return
    raise SystemExit("no supported screenshot tool found (scrot, gnome-screenshot, or import)")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--path",
        help="output file path or directory; overrides --mode",
    )
    parser.add_argument(
        "--mode",
        choices=("default", "temp"),
        default="default",
        help="default saves to the OS screenshot location; temp saves to the temp dir",
    )
    parser.add_argument(
        "--format",
        default="png",
        help="image format/extension (default: png)",
    )
    parser.add_argument(
        "--app",
        help="macOS only: capture all matching on-screen windows for this app name",
    )
    parser.add_argument(
        "--window-name",
        help="macOS only: substring match for a window title (optionally scoped by --app)",
    )
    parser.add_argument(
        "--list-windows",
        action="store_true",
        help="macOS only: list matching window ids instead of capturing",
    )
    parser.add_argument(
        "--region",
        type=parse_region,
        help="capture region as x,y,w,h (pixel coordinates)",
    )
    parser.add_argument(
        "--window-id",
        type=int,
        help="capture a specific window id when supported",
    )
    parser.add_argument(
        "--active-window",
        action="store_true",
        help="capture the focused/active window only when supported",
    )
    parser.add_argument(
        "--interactive",
        action="store_true",
        help="use interactive selection where the OS tool supports it",
    )
    args = parser.parse_args()

    if args.region and args.window_id is not None:
        raise SystemExit("choose either --region or --window-id, not both")
    if args.region and args.active_window:
        raise SystemExit("choose either --region or --active-window, not both")
    if args.window_id is not None and args.active_window:
        raise SystemExit("choose either --window-id or --active-window, not both")
    if args.app and args.window_id is not None:
        raise SystemExit("choose either --app or --window-id, not both")
    if args.region and args.app:
        raise SystemExit("choose either --region or --app, not both")
    if args.region and args.window_name:
        raise SystemExit("choose either --region or --window-name, not both")
    if args.interactive and args.app:
        raise SystemExit("choose either --interactive or --app, not both")
    if args.interactive and args.window_name:
        raise SystemExit("choose either --interactive or --window-name, not both")
    if args.interactive and args.window_id is not None:
        raise SystemExit("choose either --interactive or --window-id, not both")
    if args.interactive and args.active_window:
        raise SystemExit("choose either --interactive or --active-window, not both")
    if args.list_windows and (args.region or args.window_id is not None or args.interactive):
        raise SystemExit("--list-windows only supports --app, --window-name, and --active-window")

    test_mode = test_mode_enabled()
    system = platform.system()
    if test_mode:
        override = test_platform_override()
        if override:
            system = override
    window_ids: list[int] = []
    display_ids: list[int] = []

    if system != "Darwin" and (args.app or args.window_name or args.list_windows):
        raise SystemExit("--app/--window-name/--list-windows are supported on macOS only")

    if system == "Darwin":
        if test_mode:
            if args.list_windows:
                list_test_macos_windows(args)
                return
            if args.window_id is not None:
                window_ids = [args.window_id]
            elif args.app or args.window_name or args.active_window:
                window_ids = resolve_test_macos_windows(args)
            elif args.region is None and not args.interactive:
                display_ids = test_display_ids()
        else:
            ensure_macos_permissions()
            if args.list_windows:
                list_macos_windows(args)
                return
            if args.window_id is not None:
                window_ids = [args.window_id]
            elif args.app or args.window_name or args.active_window:
                window_ids = resolve_macos_windows(args)
            elif args.region is None and not args.interactive:
                display_ids = macos_display_indexes()

    output = resolve_output_path(args.path, args.mode, args.format, system)

    if test_mode:
        if system == "Darwin":
            if window_ids:
                suffixes = [f"w{wid}" for wid in window_ids]
                paths = multi_output_paths(output, suffixes)
                for path in paths:
                    write_test_png(path)
                for path in paths:
                    print(path)
                return
            if len(display_ids) > 1:
                suffixes = [f"d{did}" for did in display_ids]
                paths = multi_output_paths(output, suffixes)
                for path in paths:
                    write_test_png(path)
                for path in paths:
                    print(path)
                return
        write_test_png(output)
        print(output)
        return

    if system == "Darwin":
        if window_ids:
            suffixes = [f"w{wid}" for wid in window_ids]
            paths = multi_output_paths(output, suffixes)
            for wid, path in zip(window_ids, paths):
                capture_macos(args, path, window_id=wid)
            for path in paths:
                print(path)
            return
        if len(display_ids) > 1:
            suffixes = [f"d{did}" for did in display_ids]
            paths = multi_output_paths(output, suffixes)
            for did, path in zip(display_ids, paths):
                capture_macos(args, path, display=did)
            for path in paths:
                print(path)
            return
        capture_macos(args, output)
    elif system == "Linux":
        capture_linux(args, output)
    elif system == "Windows":
        raise SystemExit(
            "Windows support lives in scripts/take_screenshot.ps1; run it with PowerShell"
        )
    else:
        raise SystemExit(f"unsupported platform: {system}")

    print(output)


if __name__ == "__main__":
    main()