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.891.3.0 adds best-effort metadata (dims, aspect ratio, MIME type, size, EXIF10without GPS) read from the image header only. Metadata never fails the node:11on any error the dims are null, exif is {}, `warning` explains why, and12`image` is still returned exactly as before.13"""1415from __future__ import annotations1617import json18import re19import sys20import traceback21from typing import Any2223from gais import Gais2425OUTPUT_DIR = "/data/output"2627_UUID_RE = re.compile(28    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$",29    re.IGNORECASE,30)313233def _normalize_drive_item_id(value: Any) -> str:34    """Accept a bare UUID or a mention://drive_{file,folder}/<uuid> URI."""35    if not isinstance(value, str):36        raise ValueError(f"Drive item id must be a string, got {type(value).__name__}")37    s = value.strip()38    for prefix in ("mention://drive_file/", "mention://drive_folder/"):39        if s.startswith(prefix):40            s = s[len(prefix):]41            break42    if not _UUID_RE.match(s):43        raise ValueError(f"Drive item id is not a valid UUID: {value!r}")44    return s454647def _safe_output_name(name: Any, default: str) -> str:48    """Reduce a Drive-supplied name to a bare filename so it cannot escape49    OUTPUT_DIR (path traversal). The name comes from user-controlled DriveItem50    metadata; os/pathlib are lint-blocked, so strip directory components by51    hand and reject empty / dot names."""52    base = str(name or "").replace("\\", "/").rsplit("/", 1)[-1].strip()53    if not base or base in (".", ".."):54        return default55    return base565758def _copy_to_output(src_path: str, name: str) -> int:59    """Stream-copy the downloaded file into /data/output; return bytes copied.6061    os/shutil are lint-blocked in new-format nodes, so copy with plain open()62    in 1 MiB chunks (never read the whole file into the 512 MB container).63    `name` must already be a sanitized basename (see _safe_output_name).64    """65    total = 066    with open(src_path, "rb") as src, open(f"{OUTPUT_DIR}/{name}", "wb") as dst:67        while True:68            chunk = src.read(1024 * 1024)69            if not chunk:70                break71            dst.write(chunk)72            total += len(chunk)73    return total747576# EXIF tags (numeric ids so this works without importing PIL at module load).77_TAG_MAKE = 0x010F78_TAG_MODEL = 0x011079_TAG_ORIENTATION = 0x011280_TAG_EXIF_IFD = 0x876981_TAG_DATETIME_ORIGINAL = 0x900382# Orientations 5-8 rotate the image by 90/270 degrees: displayed w/h swap.83_SWAPPED_ORIENTATIONS = {5, 6, 7, 8}848586# EXIF text reaches the agent's context through getNodeOutput, and the image87# may come from another user. Make/Model (ASCII by spec, short in practice)88# must be a short run of letters, digits and name punctuation: markup, control89# characters and long values are dropped. That blocks structure, not meaning:90# a short English phrase still passes, so it is untrusted text all the same.91# DateTimeOriginal must be "YYYY:MM:DD HH:MM:SS" in ASCII digits; the all-zero92# "not set" value is dropped.93_EXIF_LABEL_RE = re.compile(r"[A-Za-z0-9 .,_+/()&'#:\[\]-]{1,64}")94_EXIF_DATETIME_RE = re.compile(r"[0-9]{4}:[0-9]{2}:[0-9]{2} [0-9]{2}:[0-9]{2}:[0-9]{2}")95_EXIF_DATETIME_UNSET = "0000:00:00 00:00:00"96# The Drive mimeType is whatever the uploader's client sent.97_MIME_TYPE_RE = re.compile(r"[A-Za-z0-9.+_-]{1,64}/[A-Za-z0-9.+_-]{1,64}")9899100def _exif_text(value: Any, pattern: re.Pattern[str] = _EXIF_LABEL_RE) -> str | None:101    """EXIF ASCII values arrive as str (sometimes bytes) padded with NULs.102103    Control characters are removed, then the value must match ``pattern``104    as a whole; anything else is dropped."""105    if isinstance(value, bytes):106        value = value.decode("utf-8", errors="replace")107    if not isinstance(value, str):108        return None109    text = "".join(ch for ch in value if ch.isprintable()).strip()110    return text if text and pattern.fullmatch(text) else None111112113def _read_image_metadata(path: str, mime_type: str) -> dict[str, Any]:114    """Best-effort header metadata. Never raises.115116    Opens the file lazily (Image.open reads the header; no pixel decode).117    The EXIF output is a whitelist (make, model, datetime_original,118    orientation): GPS and every other tag are never copied out.119    """120    try:121        from PIL import Image  # lazy: a missing PIL is a warning, not a failure122123        with Image.open(path) as im:124            width, height = im.size125            detected_mime = Image.MIME.get(im.format or "", "")126            # PNG's getexif() decodes the whole image when the eXIf chunk is127            # not in the header; skip it so this stays header-only. Accepted128            # limit: orientation from an eXIf chunk after the image data, a129            # text-chunk EXIF or XMP is not applied here, while Image Edit's130            # "Match input image" (which decodes) does apply it.131            if im.format == "PNG" and "exif" not in im.info:132                raw_exif: Any = {}133                exif_ifd: Any = {}134            else:135                raw_exif = im.getexif()136                exif_ifd = raw_exif.get_ifd(_TAG_EXIF_IFD) if raw_exif else {}137            make = _exif_text(raw_exif.get(_TAG_MAKE))138            model = _exif_text(raw_exif.get(_TAG_MODEL))139            taken = _exif_text(exif_ifd.get(_TAG_DATETIME_ORIGINAL), _EXIF_DATETIME_RE)140            if taken == _EXIF_DATETIME_UNSET:141                taken = None142            orientation = raw_exif.get(_TAG_ORIENTATION)143144        exif: dict[str, Any] = {}145        if make:146            exif["make"] = make147        if model:148            exif["model"] = model149        if taken:150            exif["datetime_original"] = taken151        if isinstance(orientation, int) and 1 <= orientation <= 8:152            exif["orientation"] = orientation153            if orientation in _SWAPPED_ORIENTATIONS:154                width, height = height, width155156        if width <= 0 or height <= 0:157            raise ValueError(f"invalid dimensions {width}x{height}")158159        return {160            "width": width,161            "height": height,162            "aspect_ratio": round(width / height, 4),163            "mime_type": mime_type or detected_mime,164            "exif": exif,165            "warning": "",166        }167    except Exception as e:  # noqa: BLE001 — metadata is best-effort by contract168        if type(e).__name__ == "UnidentifiedImageError":169            reason = "unsupported or corrupt image"170        else:171            reason = str(e) or type(e).__name__172        warning = f"Could not read image metadata ({mime_type or 'unknown'}): {reason}"173        print(f"[drive-image-reader] {warning}", file=sys.stderr)174        return {175            "width": None,176            "height": None,177            "aspect_ratio": None,178            "mime_type": mime_type,179            "exif": {},180            "warning": warning,181        }182183184def process(inputs: dict[str, Any]) -> dict[str, Any]:185    """Validate the input, download via gais.drive, stage the Image output."""186    raw = inputs.get("driveItemId")187    if not raw:188        raise ValueError("Required input 'driveItemId' not provided")189190    # Single-image reader: reject multi-item arrays.191    if isinstance(raw, list):192        if len(raw) == 1:193            raw = raw[0]194        else:195            raise ValueError(196                f"Drive Image Reader accepts a single image, got {len(raw)} items. "197                "Use one Drive Image Reader per image."198            )199    elif not isinstance(raw, str):200        raise ValueError(f"driveItemId must be a string, got {type(raw).__name__}")201202    item_id = _normalize_drive_item_id(raw)203204    result = Gais.drive.download(item_id=item_id)205    file_name = _safe_output_name(result.metadata.get("name", "image.png"), "image.png")206    raw_mime = result.metadata.get("mimeType", "") or ""207    raw_mime = raw_mime if isinstance(raw_mime, str) else ""208    # The guard reads the raw value: "text/plain; charset=utf-8" is still not209    # an image. Only the type part, bounded, is quoted back.210    if raw_mime and not raw_mime.startswith("image/"):211        raise ValueError(212            f"Selected Drive item is not an image (mimeType: {raw_mime.split(';')[0].strip()[:64]}). "213            "Please select an image file."214        )215    # Reported as is only when it is a plain type/subtype; otherwise the216    # header-detected type is used.217    mime_type = raw_mime if _MIME_TYPE_RE.fullmatch(raw_mime) else ""218219    size_bytes = _copy_to_output(str(result.file("file")), file_name)220    meta = _read_image_metadata(f"{OUTPUT_DIR}/{file_name}", mime_type)221    return {222        "image": file_name,223        "width": meta["width"],224        "height": meta["height"],225        "aspect_ratio": meta["aspect_ratio"],226        "mime_type": meta["mime_type"],227        "size_bytes": size_bytes,228        "exif": meta["exif"],229        "warning": meta["warning"],230    }231232233def main() -> None:234    try:235        envelope = json.loads(sys.stdin.read() or "{}")236        if not isinstance(envelope, dict):237            envelope = {}238        inputs = envelope.get("inputs", {}) or {}239        json.dump(process(inputs), sys.stdout)240    except Exception as e:241        print(242            json.dumps({243                "error": str(e),244                "errorType": type(e).__name__,245                "traceback": traceback.format_exc(),246            }),247            file=sys.stderr,248        )249        sys.exit(1)250251252if __name__ == "__main__":253    main()254