Rhinox13/chatapi
0
1from __future__ import annotations
2
3import base64
4import binascii
5import hashlib
6import json
7import re
8from pathlib import Path
9from typing import Any
10
11from ..repositories import SystemConfigStore
12from ..repositories.users import UserStore
13
14
15_DATA_IMAGE_RE = re.compile(r"^data:(image/[a-zA-Z0-9.+-]+);base64,(.+)$", re.IGNORECASE | re.DOTALL)
16_UPLOAD_URL_RE = re.compile(r"/api/uploads/imgs/([A-Za-z0-9._-]+)(?:\?.*)?$", re.IGNORECASE)
17ALLOWED_IMAGE_MIME_TYPES = {
18 "image/png",
19 "image/jpeg",
20 "image/jpg",
21 "image/gif",
22 "image/webp",
23 "image/avif",
24 "image/bmp",
25 "image/tiff",
26}
27
28
29def _normalize_key(value: str) -> str:
30 return re.sub(r"[^a-z0-9]+", "", value.lower())
31
32
33def _mime_to_extension(mime_type: str) -> str:
34 normalized = mime_type.lower().split(";", 1)[0].strip()
35 mapping = {
36 "image/jpeg": "jpg",
37 "image/jpg": "jpg",
38 "image/png": "png",
39 "image/gif": "gif",
40 "image/webp": "webp",
41 "image/avif": "avif",
42 "image/bmp": "bmp",
43 "image/tiff": "tiff",
44 }
45 return mapping.get(normalized, normalized.rsplit("/", 1)[-1] or "img")
46
47
48def _mime_from_magic(image_bytes: bytes) -> str | None:
49 if image_bytes.startswith(b"\x89PNG\r\n\x1a\n"):
50 return "image/png"
51 if image_bytes.startswith(b"\xff\xd8\xff"):
52 return "image/jpeg"
53 if image_bytes.startswith(b"GIF87a") or image_bytes.startswith(b"GIF89a"):
54 return "image/gif"
55 if image_bytes.startswith(b"RIFF") and image_bytes[8:12] == b"WEBP":
56 return "image/webp"
57 if image_bytes.startswith(b"BM"):
58 return "image/bmp"
59 if image_bytes.startswith(b"\x00\x00\x00") and b"ftypavif" in image_bytes[:32]:
60 return "image/avif"
61 if image_bytes.lstrip().startswith(b"<?xml") or b"<svg" in image_bytes[:256].lower():
62 return None
63 return None
64
65
66def _is_image_key(key: str) -> bool:
67 normalized = _normalize_key(key)
68 return any(token in normalized for token in ("image", "img", "photo", "picture", "thumbnail", "avatar", "filedata", "basedata", "binary"))
69
70
71def _looks_like_base64(value: str) -> bool:
72 candidate = re.sub(r"\s+", "", value)
73 if len(candidate) < 32 or len(candidate) % 4 != 0:
74 return False
75 return bool(re.fullmatch(r"[A-Za-z0-9+/=]+", candidate))
76
77
78def _try_parse_structured_content(raw: str) -> Any | None:
79 try:
80 return json.loads(raw)
81 except (TypeError, json.JSONDecodeError):
82 pass
83
84 trimmed = raw.strip()
85 if not trimmed or trimmed[0] not in "[{":
86 return None
87
88 normalized = ""
89 in_single_quote = False
90 in_double_quote = False
91 escape_next = False
92
93 for char in trimmed:
94 if escape_next:
95 normalized += char
96 escape_next = False
97 continue
98 if char == "\\":
99 normalized += char
100 escape_next = True
101 continue
102 if char == "'" and not in_double_quote:
103 normalized += '"'
104 in_single_quote = not in_single_quote
105 continue
106 if char == '"' and not in_single_quote:
107 normalized += char
108 in_double_quote = not in_double_quote
109 continue
110 normalized += '\\"' if in_single_quote and char == '"' else char
111
112 normalized = normalized.replace("None", "null").replace("True", "true").replace("False", "false")
113 try:
114 return json.loads(normalized)
115 except (TypeError, json.JSONDecodeError):
116 return None
117
118
119def _decode_base64_image(value: str, *, mime_type_hint: str | None = None) -> tuple[bytes, str] | None:
120 text = value.strip()
121 data_url_match = _DATA_IMAGE_RE.match(text)
122 if data_url_match:
123 mime_type = data_url_match.group(1)
124 payload = data_url_match.group(2)
125 else:
126 if not _looks_like_base64(text):
127 return None
128 mime_type = mime_type_hint or ""
129 payload = text
130 try:
131 image_bytes = base64.b64decode(payload, validate=False)
132 except (binascii.Error, ValueError):
133 return None
134 detected_mime = _mime_from_magic(image_bytes)
135 if detected_mime is not None:
136 mime_type = detected_mime
137 if not mime_type.startswith("image/"):
138 return None
139 return image_bytes, mime_type
140
141
142class ImageAssetStore:
143 def __init__(
144 self,
145 base_dir: Path,
146 route_prefix: str = "/api/uploads/imgs",
147 system_config_store: SystemConfigStore | None = None,
148 user_store: UserStore | None = None,
149 ):
150 self.base_dir = base_dir
151 self.route_prefix = route_prefix.rstrip("/")
152 self.system_config_store = system_config_store
153 self.user_store = user_store
154 self._request_bytes = 0
155 self.base_dir.mkdir(parents=True, exist_ok=True)
156
157 @property
158 def route_root(self) -> str:
159 return self.route_prefix
160
161 def public_url(self, filename: str) -> str:
162 return f"{self.route_root}/{filename}"
163
164 def _filename_for(self, image_bytes: bytes, mime_type: str) -> str:
165 digest = hashlib.sha256(image_bytes).hexdigest()
166 ext = _mime_to_extension(mime_type)
167 return f"{digest}.{ext}"
168
169 def _config_int(self, key: str, default: int = 0) -> int:
170 if self.system_config_store is None:
171 return default
172 try:
173 return max(0, int(self.system_config_store.get_system_config(key, str(default)) or default))
174 except ValueError:
175 return default
176
177 def _placeholder_for_expired(self, filename: str) -> str:
178 return f"{self.route_root}/{filename}"
179
180 def _can_store(self, image_bytes: bytes) -> bool:
181 size = len(image_bytes)
182 single_limit = self._config_int("value.image_max_single_bytes", 0)
183 if single_limit > 0 and size > single_limit:
184 return False
185
186 request_limit = self._config_int("value.image_max_request_bytes", 0)
187 if request_limit > 0 and self._request_bytes + size > request_limit:
188 return False
189
190 total_limit = self._config_int("value.image_max_total_bytes", 0)
191 if total_limit > 0:
192 self.prune_oldest_to_fit(max(0, total_limit - size))
193 if self.storage_usage()["total_bytes"] + size > total_limit:
194 return False
195 return True
196
197 def _track_owner(self, filename: str, owner_id: str | None, mime_type: str) -> None:
198 if self.user_store is None or not owner_id:
199 return
200 self.user_store.set_uploaded_image_owner(filename, owner_id, mime_type)
201
202 def store_data_url(self, value: str, *, owner_id: str | None = None) -> str | None:
203 decoded = _decode_base64_image(value)
204 if decoded is None:
205 return None
206 image_bytes, mime_type = decoded
207 if mime_type not in ALLOWED_IMAGE_MIME_TYPES:
208 return None
209 filename = self._filename_for(image_bytes, mime_type)
210 if not self._can_store(image_bytes):
211 return self._placeholder_for_expired(filename)
212 path = self.base_dir / filename
213 if not path.exists():
214 path.write_bytes(image_bytes)
215 self._track_owner(filename, owner_id, mime_type)
216 self._request_bytes += len(image_bytes)
217 return self.public_url(filename)
218
219 def normalize_request_data(self, value: Any, *, owner_id: str | None = None) -> Any:
220 self._request_bytes = 0
221 return self._normalize_node(value, owner_id=owner_id)
222
223 def _rewrite_string(self, value: str, *, key: str = "", owner_id: str | None = None) -> str:
224 data_url_url = self.store_data_url(value, owner_id=owner_id)
225 if data_url_url is not None:
226 return data_url_url
227
228 decoded = _decode_base64_image(value)
229 if decoded is not None:
230 image_bytes, mime_type = decoded
231 if mime_type not in ALLOWED_IMAGE_MIME_TYPES:
232 return value
233 filename = self._filename_for(image_bytes, mime_type)
234 if not self._can_store(image_bytes):
235 return self._placeholder_for_expired(filename)
236 path = self.base_dir / filename
237 if not path.exists():
238 path.write_bytes(image_bytes)
239 self._track_owner(filename, owner_id, mime_type)
240 self._request_bytes += len(image_bytes)
241 return self.public_url(filename)
242
243 if _is_image_key(key):
244 decoded = _decode_base64_image(value)
245 if decoded is not None:
246 image_bytes, mime_type = decoded
247 if mime_type not in ALLOWED_IMAGE_MIME_TYPES:
248 return value
249 filename = self._filename_for(image_bytes, mime_type)
250 if not self._can_store(image_bytes):
251 return self._placeholder_for_expired(filename)
252 path = self.base_dir / filename
253 if not path.exists():
254 path.write_bytes(image_bytes)
255 self._track_owner(filename, owner_id, mime_type)
256 self._request_bytes += len(image_bytes)
257 return self.public_url(filename)
258 return value
259
260 def _normalize_node(self, value: Any, *, key: str = "", owner_id: str | None = None) -> Any:
261 if isinstance(value, str):
262 parsed = _try_parse_structured_content(value)
263 if parsed is not None:
264 return self._normalize_node(parsed, key=key, owner_id=owner_id)
265 return self._rewrite_string(value, key=key, owner_id=owner_id)
266 if isinstance(value, list):
267 return [self._normalize_node(item, key=key, owner_id=owner_id) for item in value]
268 if isinstance(value, dict):
269 rewritten: dict[str, Any] = {}
270 for item_key, item_value in value.items():
271 rewritten[item_key] = self._normalize_node(item_value, key=str(item_key), owner_id=owner_id)
272 return rewritten
273 return value
274
275 def rewrite_value(self, value: Any, *, key: str = "", owner_id: str | None = None) -> Any:
276 return self._normalize_node(value, key=key, owner_id=owner_id)
277
278 def normalize_content(self, content: Any, *, owner_id: str | None = None) -> str:
279 if isinstance(content, str):
280 parsed = _try_parse_structured_content(content)
281 if parsed is not None:
282 rewritten = self.rewrite_value(parsed, key="content", owner_id=owner_id)
283 return json.dumps(rewritten, ensure_ascii=False, separators=(",", ":"))
284 rewritten = self.rewrite_value(content, key="content", owner_id=owner_id)
285 if isinstance(rewritten, str):
286 return rewritten.replace("\r\n", "\n").replace("\\r\\n", "\n").replace("\\n", "\n")
287 return json.dumps(rewritten, ensure_ascii=False, separators=(",", ":"))
288
289 def backfill_owners_from_messages(self, messages: list[Any], owner_lookup: dict[str, str]) -> None:
290 if self.user_store is None:
291 return
292 for message in messages:
293 owner_id = owner_lookup.get(str(getattr(message, "conversation_id", "")).strip(), "")
294 if not owner_id:
295 continue
296 for match in _UPLOAD_URL_RE.finditer(str(getattr(message, "content", "") or "")):
297 filename = match.group(1)
298 if filename:
299 mime_type = ""
300 suffix = Path(filename).suffix.lower()
301 if suffix == ".png":
302 mime_type = "image/png"
303 elif suffix in {".jpg", ".jpeg"}:
304 mime_type = "image/jpeg"
305 elif suffix == ".gif":
306 mime_type = "image/gif"
307 elif suffix == ".webp":
308 mime_type = "image/webp"
309 elif suffix == ".avif":
310 mime_type = "image/avif"
311 elif suffix == ".bmp":
312 mime_type = "image/bmp"
313 elif suffix in {".tif", ".tiff"}:
314 mime_type = "image/tiff"
315 self.user_store.set_uploaded_image_owner(filename, owner_id, mime_type)
316
317 def referenced_filenames(self, content: str) -> set[str]:
318 if not content:
319 return set()
320 pattern = re.compile(
321 rf"(?:https?://[^\s\"']+)?{re.escape(self.route_root)}/([A-Za-z0-9._-]+)",
322 re.IGNORECASE,
323 )
324 return {match.group(1) for match in pattern.finditer(content)}
325
326 def cleanup_orphans(self, messages: list[Any]) -> list[str]:
327 referenced: set[str] = set()
328 for message in messages:
329 content = getattr(message, "content", "")
330 if isinstance(content, str):
331 referenced.update(self.referenced_filenames(content))
332
333 deleted: list[str] = []
334 if not self.base_dir.exists():
335 return deleted
336
337 for path in self.base_dir.iterdir():
338 if not path.is_file():
339 continue
340 if path.name in referenced:
341 continue
342 path.unlink(missing_ok=True)
343 deleted.append(path.name)
344 return deleted
345
346 def storage_usage(self, messages: list[Any] | None = None) -> dict[str, int]:
347 total_bytes = 0
348 file_count = 0
349 orphan_bytes = 0
350 referenced: set[str] | None = None
351 if messages is not None:
352 referenced = set()
353 for message in messages:
354 content = getattr(message, "content", "")
355 if isinstance(content, str):
356 referenced.update(self.referenced_filenames(content))
357
358 if not self.base_dir.exists():
359 return {
360 "total_bytes": 0,
361 "file_count": 0,
362 "orphan_bytes": 0,
363 "orphan_count": 0,
364 }
365
366 orphan_count = 0
367 for path in self.base_dir.iterdir():
368 if not path.is_file():
369 continue
370 size = path.stat().st_size
371 total_bytes += size
372 file_count += 1
373 if referenced is not None and path.name not in referenced:
374 orphan_bytes += size
375 orphan_count += 1
376
377 return {
378 "total_bytes": total_bytes,
379 "file_count": file_count,
380 "orphan_bytes": orphan_bytes,
381 "orphan_count": orphan_count,
382 }
383
384 def prune_oldest_to_fit(self, target_bytes: int) -> list[str]:
385 target_bytes = max(0, int(target_bytes))
386 files = [
387 path
388 for path in self.base_dir.iterdir()
389 if path.is_file()
390 ] if self.base_dir.exists() else []
391 total = sum(path.stat().st_size for path in files)
392 if total <= target_bytes:
393 return []
394
395 deleted: list[str] = []
396 for path in sorted(files, key=lambda item: item.stat().st_mtime):
397 if total <= target_bytes:
398 break
399 size = path.stat().st_size
400 path.unlink(missing_ok=True)
401 total -= size
402 deleted.append(path.name)
403 return deleted
404