Drive Image Reader

1"""drive-image-reader — read an image from Drive, output it as an Image file.23Migrated from the legacy monolithic template to the new multi-file format.4Privileged Drive IO is delegated to the gais.drive SDK (USER_TOKEN + Drive API5live there); this node imports no HTTP client and uses no os/shutil (both are6lint-blocked). The downloaded file is stream-copied into /data/output with7plain open() so the executor uploads it as the Image output.8"""910from __future__ import annotations1112import json13import re14import sys15import traceback16from typing import Any1718from gais import Gais1920OUTPUT_DIR = "/data/output"2122_UUID_RE = re.compile(23    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$",24    re.IGNORECASE,25)262728def _normalize_drive_item_id(value: Any) -> str:29    """Accept a bare UUID or a mention://drive_{file,folder}/<uuid> URI."""30    if not isinstance(value, str):31        raise ValueError(f"Drive item id must be a string, got {type(value).__name__}")32    s = value.strip()33    for prefix in ("mention://drive_file/", "mention://drive_folder/"):34        if s.startswith(prefix):35            s = s[len(prefix):]36            break37    if not _UUID_RE.match(s):38        raise ValueError(f"Drive item id is not a valid UUID: {value!r}")39    return s404142def _safe_output_name(name: Any, default: str) -> str:43    """Reduce a Drive-supplied name to a bare filename so it cannot escape44    OUTPUT_DIR (path traversal). The name comes from user-controlled DriveItem45    metadata; os/pathlib are lint-blocked, so strip directory components by46    hand and reject empty / dot names."""47    base = str(name or "").replace("\\", "/").rsplit("/", 1)[-1].strip()48    if not base or base in (".", ".."):49        return default50    return base515253def _copy_to_output(src_path: str, name: str) -> None:54    """Stream-copy the downloaded file into /data/output.5556    os/shutil are lint-blocked in new-format nodes, so copy with plain open()57    in 1 MiB chunks (never read the whole file into the 512 MB container).58    `name` must already be a sanitized basename (see _safe_output_name).59    """60    with open(src_path, "rb") as src, open(f"{OUTPUT_DIR}/{name}", "wb") as dst:61        while True:62            chunk = src.read(1024 * 1024)63            if not chunk:64                break65            dst.write(chunk)666768def process(inputs: dict[str, Any]) -> dict[str, Any]:69    """Validate the input, download via gais.drive, stage the Image output."""70    raw = inputs.get("driveItemId")71    if not raw:72        raise ValueError("Required input 'driveItemId' not provided")7374    # Single-image reader: reject multi-item arrays.75    if isinstance(raw, list):76        if len(raw) == 1:77            raw = raw[0]78        else:79            raise ValueError(80                f"Drive Image Reader accepts a single image, got {len(raw)} items. "81                "Use one Drive Image Reader per image."82            )83    elif not isinstance(raw, str):84        raise ValueError(f"driveItemId must be a string, got {type(raw).__name__}")8586    item_id = _normalize_drive_item_id(raw)8788    result = Gais.drive.download(item_id=item_id)89    file_name = _safe_output_name(result.metadata.get("name", "image.png"), "image.png")90    mime_type = result.metadata.get("mimeType", "") or ""91    if mime_type and not mime_type.startswith("image/"):92        raise ValueError(93            f"Selected Drive item is not an image (mimeType: {mime_type}). "94            "Please select an image file."95        )9697    _copy_to_output(str(result.file("file")), file_name)98    return {"image": file_name}99100101def main() -> None:102    try:103        envelope = json.loads(sys.stdin.read() or "{}")104        if not isinstance(envelope, dict):105            envelope = {}106        inputs = envelope.get("inputs", {}) or {}107        json.dump(process(inputs), sys.stdout)108    except Exception as e:109        print(110            json.dumps({111                "error": str(e),112                "errorType": type(e).__name__,113                "traceback": traceback.format_exc(),114            }),115            file=sys.stderr,116        )117        sys.exit(1)118119120if __name__ == "__main__":121    main()