agent-zero/tests/test_image_get_security.py

import asyncio
import base64
import sys
import threading
from pathlib import Path

from flask import Flask, request

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from api import image_get


def _patch_base_dir(monkeypatch, base_dir: Path, *, development: bool = False) -> None:
    base_dir.mkdir(parents=True, exist_ok=True)

    def fake_get_abs_path(*parts: str) -> str:
        if len(parts) == 1 and Path(str(parts[0])).is_absolute():
            return str(Path(str(parts[0])))
        return str(base_dir.joinpath(*(str(part) for part in parts)))

    monkeypatch.setattr(image_get.files, "get_base_dir", lambda: str(base_dir))
    monkeypatch.setattr(image_get.files, "get_abs_path", fake_get_abs_path)
    monkeypatch.setattr(image_get.runtime, "is_development", lambda: development)


async def _request_image(path: str):
    app = Flask("test_image_get_security")
    handler = image_get.ImageGet(app, threading.Lock())
    with app.test_request_context("/api/image_get"):
        return await handler.process({"path": path}, request)


def test_image_get_serves_images_inside_base_dir(tmp_path, monkeypatch):
    base_dir = tmp_path / "a0"
    _patch_base_dir(monkeypatch, base_dir)
    image_path = base_dir / "usr" / "uploads" / "safe.png"
    image_path.parent.mkdir(parents=True)
    image_path.write_bytes(b"\x89PNG\r\n\x1a\n")

    response = asyncio.run(_request_image(str(image_path)))

    assert response.status_code == 200
    assert response.headers["X-File-Type"] == "image"
    assert response.headers["X-Content-Type-Options"] == "nosniff"


def test_image_get_blocks_image_paths_outside_base_dir(tmp_path, monkeypatch):
    base_dir = tmp_path / "a0"
    _patch_base_dir(monkeypatch, base_dir)
    outside_image = tmp_path / "outside.png"
    outside_image.write_bytes(b"outside")

    response = asyncio.run(_request_image(str(outside_image)))

    assert response.status_code == 403
    assert response.get_data(as_text=True) == "Path is outside of allowed directory"


def test_image_get_blocks_symlink_escape_from_base_dir(tmp_path, monkeypatch):
    base_dir = tmp_path / "a0"
    _patch_base_dir(monkeypatch, base_dir)
    outside_image = tmp_path / "secret.png"
    outside_image.write_bytes(b"secret")
    link_path = base_dir / "usr" / "uploads" / "linked.png"
    link_path.parent.mkdir(parents=True)
    link_path.symlink_to(outside_image)

    response = asyncio.run(_request_image(str(link_path)))

    assert response.status_code == 403


def test_image_get_hardens_svg_responses(tmp_path, monkeypatch):
    base_dir = tmp_path / "a0"
    _patch_base_dir(monkeypatch, base_dir)
    svg_path = base_dir / "usr" / "uploads" / "payload.svg"
    svg_path.parent.mkdir(parents=True)
    svg_path.write_text(
        '<svg xmlns="http://www.w3.org/2000/svg"><script>alert(1)</script></svg>',
        encoding="utf-8",
    )

    response = asyncio.run(_request_image(str(svg_path)))

    assert response.status_code == 200
    assert response.headers["Content-Security-Policy"].startswith("sandbox;")
    assert "script-src 'none'" in response.headers["Content-Security-Policy"]
    assert response.headers["X-Content-Type-Options"] == "nosniff"


def test_image_get_development_fallback_validates_remote_path(tmp_path, monkeypatch):
    base_dir = tmp_path / "a0"
    _patch_base_dir(monkeypatch, base_dir, development=True)
    calls = []

    async def fake_call_development_function(func, *args, **kwargs):
        calls.append(func.__name__)
        if func is image_get._resolve_allowed_image_path:
            return "/a0/usr/uploads/remote.png"
        if func is image_get.files.exists:
            return True
        if func is image_get.files.read_file_base64:
            return base64.b64encode(b"\x89PNG\r\n\x1a\n").decode("ascii")
        raise AssertionError(f"Unexpected remote call: {func.__name__}")

    monkeypatch.setattr(
        image_get.runtime,
        "call_development_function",
        fake_call_development_function,
    )

    response = asyncio.run(_request_image("/a0/usr/uploads/remote.png"))

    assert response.status_code == 200
    assert response.headers["X-File-Type"] == "image"
    assert calls == ["_resolve_allowed_image_path", "exists", "read_file_base64"]


def test_image_get_development_fallback_does_not_read_rejected_remote_path(
    tmp_path,
    monkeypatch,
):
    base_dir = tmp_path / "a0"
    _patch_base_dir(monkeypatch, base_dir, development=True)
    calls = []

    async def fake_call_development_function(func, *args, **kwargs):
        calls.append(func.__name__)
        if func is image_get._resolve_allowed_image_path:
            raise ValueError("Path is outside of allowed directory")
        raise AssertionError(f"Unexpected remote call after validation: {func.__name__}")

    monkeypatch.setattr(
        image_get.runtime,
        "call_development_function",
        fake_call_development_function,
    )
    monkeypatch.setattr(
        image_get,
        "_send_fallback_icon",
        lambda _icon_name: image_get.Response("fallback", status=200),
    )

    response = asyncio.run(_request_image("/a0/usr/uploads/rejected.png"))

    assert response.status_code == 200
    assert response.get_data(as_text=True) == "fallback"
    assert calls == ["_resolve_allowed_image_path"]