@samitouri / QOSami_Agents_Zero-2
Code Commits Issues Pulls Sessions CI/CD
main
main
py 151 lines 5.26 KB
Raw
1 import asyncio
2 import base64
3 import sys
4 import threading
5 from pathlib import Path
6
7 from flask import Flask, request
8
9 PROJECT_ROOT = Path(__file__).resolve().parents[1]
10 if str(PROJECT_ROOT) not in sys.path:
11 sys.path.insert(0, str(PROJECT_ROOT))
12
13 from api import image_get
14
15
16 def _patch_base_dir(monkeypatch, base_dir: Path, *, development: bool = False) -> None:
17 base_dir.mkdir(parents=True, exist_ok=True)
18
19 def fake_get_abs_path(*parts: str) -> str:
20 if len(parts) == 1 and Path(str(parts[0])).is_absolute():
21 return str(Path(str(parts[0])))
22 return str(base_dir.joinpath(*(str(part) for part in parts)))
23
24 monkeypatch.setattr(image_get.files, "get_base_dir", lambda: str(base_dir))
25 monkeypatch.setattr(image_get.files, "get_abs_path", fake_get_abs_path)
26 monkeypatch.setattr(image_get.runtime, "is_development", lambda: development)
27
28
29 async def _request_image(path: str):
30 app = Flask("test_image_get_security")
31 handler = image_get.ImageGet(app, threading.Lock())
32 with app.test_request_context("/api/image_get"):
33 return await handler.process({"path": path}, request)
34
35
36 def test_image_get_serves_images_inside_base_dir(tmp_path, monkeypatch):
37 base_dir = tmp_path / "a0"
38 _patch_base_dir(monkeypatch, base_dir)
39 image_path = base_dir / "usr" / "uploads" / "safe.png"
40 image_path.parent.mkdir(parents=True)
41 image_path.write_bytes(b"\x89PNG\r\n\x1a\n")
42
43 response = asyncio.run(_request_image(str(image_path)))
44
45 assert response.status_code == 200
46 assert response.headers["X-File-Type"] == "image"
47 assert response.headers["X-Content-Type-Options"] == "nosniff"
48
49
50 def test_image_get_blocks_image_paths_outside_base_dir(tmp_path, monkeypatch):
51 base_dir = tmp_path / "a0"
52 _patch_base_dir(monkeypatch, base_dir)
53 outside_image = tmp_path / "outside.png"
54 outside_image.write_bytes(b"outside")
55
56 response = asyncio.run(_request_image(str(outside_image)))
57
58 assert response.status_code == 403
59 assert response.get_data(as_text=True) == "Path is outside of allowed directory"
60
61
62 def test_image_get_blocks_symlink_escape_from_base_dir(tmp_path, monkeypatch):
63 base_dir = tmp_path / "a0"
64 _patch_base_dir(monkeypatch, base_dir)
65 outside_image = tmp_path / "secret.png"
66 outside_image.write_bytes(b"secret")
67 link_path = base_dir / "usr" / "uploads" / "linked.png"
68 link_path.parent.mkdir(parents=True)
69 link_path.symlink_to(outside_image)
70
71 response = asyncio.run(_request_image(str(link_path)))
72
73 assert response.status_code == 403
74
75
76 def test_image_get_hardens_svg_responses(tmp_path, monkeypatch):
77 base_dir = tmp_path / "a0"
78 _patch_base_dir(monkeypatch, base_dir)
79 svg_path = base_dir / "usr" / "uploads" / "payload.svg"
80 svg_path.parent.mkdir(parents=True)
81 svg_path.write_text(
82 '<svg xmlns="http://www.w3.org/2000/svg"><script>alert(1)</script></svg>',
83 encoding="utf-8",
84 )
85
86 response = asyncio.run(_request_image(str(svg_path)))
87
88 assert response.status_code == 200
89 assert response.headers["Content-Security-Policy"].startswith("sandbox;")
90 assert "script-src 'none'" in response.headers["Content-Security-Policy"]
91 assert response.headers["X-Content-Type-Options"] == "nosniff"
92
93
94 def test_image_get_development_fallback_validates_remote_path(tmp_path, monkeypatch):
95 base_dir = tmp_path / "a0"
96 _patch_base_dir(monkeypatch, base_dir, development=True)
97 calls = []
98
99 async def fake_call_development_function(func, *args, **kwargs):
100 calls.append(func.__name__)
101 if func is image_get._resolve_allowed_image_path:
102 return "/a0/usr/uploads/remote.png"
103 if func is image_get.files.exists:
104 return True
105 if func is image_get.files.read_file_base64:
106 return base64.b64encode(b"\x89PNG\r\n\x1a\n").decode("ascii")
107 raise AssertionError(f"Unexpected remote call: {func.__name__}")
108
109 monkeypatch.setattr(
110 image_get.runtime,
111 "call_development_function",
112 fake_call_development_function,
113 )
114
115 response = asyncio.run(_request_image("/a0/usr/uploads/remote.png"))
116
117 assert response.status_code == 200
118 assert response.headers["X-File-Type"] == "image"
119 assert calls == ["_resolve_allowed_image_path", "exists", "read_file_base64"]
120
121
122 def test_image_get_development_fallback_does_not_read_rejected_remote_path(
123 tmp_path,
124 monkeypatch,
125 ):
126 base_dir = tmp_path / "a0"
127 _patch_base_dir(monkeypatch, base_dir, development=True)
128 calls = []
129
130 async def fake_call_development_function(func, *args, **kwargs):
131 calls.append(func.__name__)
132 if func is image_get._resolve_allowed_image_path:
133 raise ValueError("Path is outside of allowed directory")
134 raise AssertionError(f"Unexpected remote call after validation: {func.__name__}")
135
136 monkeypatch.setattr(
137 image_get.runtime,
138 "call_development_function",
139 fake_call_development_function,
140 )
141 monkeypatch.setattr(
142 image_get,
143 "_send_fallback_icon",
144 lambda _icon_name: image_get.Response("fallback", status=200),
145 )
146
147 response = asyncio.run(_request_image("/a0/usr/uploads/rejected.png"))
148
149 assert response.status_code == 200
150 assert response.get_data(as_text=True) == "fallback"
151 assert calls == ["_resolve_allowed_image_path"]