spb/uqo-chat

1"""File storage: uploads and tool artifacts on local disk (S3 later), TTL, signed URLs."""23from __future__ import annotations45import base646import hashlib7import hmac8import re9import time10from datetime import timedelta11from pathlib import Path1213from sqlalchemy import delete, select1415from app.core.config import get_settings16from app.db import SessionLocal17from app.models import StoredFile, new_id, utcnow1819ALLOWED_UPLOAD = {20    "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",21    "xls": "application/vnd.ms-excel",22    "csv": "text/csv",23    "pdf": "application/pdf",24    "png": "image/png",25    "jpg": "image/jpeg",26    "jpeg": "image/jpeg",27    "webp": "image/webp",28    "txt": "text/plain",29    "md": "text/markdown",30    "json": "application/json",31    "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",32}33MIME_BY_TYPE = {34    "xlsx": ALLOWED_UPLOAD["xlsx"], "image": "image/png", "csv": "text/csv",35    "pdf": "application/pdf", "file": "application/octet-stream", "docx": ALLOWED_UPLOAD["docx"],36}373839def safe_name(name: str) -> str:40    name = re.sub(r"[^\w.\-() àâäçéèêëîïôöùûüÿœæÀÂÄÇÉÈÊËÎÏÔÖÙÛÜŸŒÆ]", "_", name).strip()41    return name[:120] or "fichier"424344def _root() -> Path:45    return get_settings().DATA_DIR / "files"464748def _path(storage_key: str) -> Path:49    return _root() / storage_key505152def ext_of(filename: str) -> str:53    return filename.rsplit(".", 1)[-1].lower() if "." in filename else ""545556async def store_upload(user_id: str, conversation_id: str | None, filename: str,57                       data: bytes) -> StoredFile:58    s = get_settings()59    filename = safe_name(filename)60    ext = ext_of(filename)61    if ext not in ALLOWED_UPLOAD:62        raise ValueError("Type de fichier non autorisé.")63    if len(data) > s.UPLOAD_MAX_MB * 1024 * 1024:64        raise ValueError(f"Fichier trop volumineux (max {s.UPLOAD_MAX_MB} Mo).")65    fid = new_id()66    key = f"{user_id[:8]}/{fid}.{ext}"67    p = _path(key)68    p.parent.mkdir(parents=True, exist_ok=True)69    p.write_bytes(data)70    ftype = "image" if ext in {"png", "jpg", "jpeg", "webp"} else ext71    rec = StoredFile(id=fid, user_id=user_id, conversation_id=conversation_id, kind="upload",72                     type=ftype, filename=filename, mime=ALLOWED_UPLOAD[ext], storage_key=key,73                     size_bytes=len(data), expires_at=utcnow() + timedelta(days=s.FILE_TTL_DAYS))74    async with SessionLocal() as session:75        session.add(rec)76        await session.commit()77    return rec787980async def store_artifact(user_id: str, conversation_id: str | None, filename: str, data: bytes,81                         ftype: str = "file", tool_call_id: str | None = None) -> StoredFile:82    s = get_settings()83    filename = safe_name(filename)84    ext = ext_of(filename) or "bin"85    fid = new_id()86    key = f"{user_id[:8]}/{fid}.{ext}"87    p = _path(key)88    p.parent.mkdir(parents=True, exist_ok=True)89    p.write_bytes(data)90    mime = ALLOWED_UPLOAD.get(ext, MIME_BY_TYPE.get(ftype, "application/octet-stream"))91    rec = StoredFile(id=fid, user_id=user_id, conversation_id=conversation_id, kind="artifact",92                     tool_call_id=tool_call_id, type=ftype, filename=filename, mime=mime,93                     storage_key=key, size_bytes=len(data),94                     expires_at=utcnow() + timedelta(days=s.FILE_TTL_DAYS))95    async with SessionLocal() as session:96        session.add(rec)97        await session.commit()98    return rec99100101async def get_file(file_id: str) -> StoredFile | None:102    async with SessionLocal() as session:103        return await session.get(StoredFile, file_id)104105106def read_bytes(rec: StoredFile) -> bytes:107    return _path(rec.storage_key).read_bytes()108109110async def list_for_conversation(conversation_id: str) -> list[StoredFile]:111    async with SessionLocal() as session:112        rows = await session.execute(113            select(StoredFile).where(StoredFile.conversation_id == conversation_id)114            .order_by(StoredFile.created_at))115        return list(rows.scalars())116117118async def sandbox_inputs(user_id: str, file_ids: list[str]) -> list[dict[str, str]]:119    out: list[dict[str, str]] = []120    for fid in file_ids[:10]:121        rec = await get_file(fid)122        if rec and rec.user_id == user_id and _path(rec.storage_key).exists():123            out.append({"name": rec.filename,124                        "content_b64": base64.b64encode(read_bytes(rec)).decode()})125    return out126127128async def set_pinned(file_id: str, user_id: str, pinned: bool) -> bool:129    s = get_settings()130    async with SessionLocal() as session:131        rec = await session.get(StoredFile, file_id)132        if not rec or rec.user_id != user_id:133            return False134        rec.pinned = pinned135        days = s.FILE_PINNED_TTL_DAYS if pinned else s.FILE_TTL_DAYS136        rec.expires_at = rec.created_at + timedelta(days=days)137        await session.commit()138        return True139140141async def purge_expired() -> int:142    n = 0143    async with SessionLocal() as session:144        rows = await session.execute(select(StoredFile).where(StoredFile.expires_at < utcnow()))145        for rec in rows.scalars():146            _path(rec.storage_key).unlink(missing_ok=True)147            n += 1148        await session.execute(delete(StoredFile).where(StoredFile.expires_at < utcnow()))149        await session.commit()150    return n151152153async def purge_user(user_id: str) -> int:154    n = 0155    async with SessionLocal() as session:156        rows = await session.execute(select(StoredFile).where(StoredFile.user_id == user_id))157        for rec in rows.scalars():158            _path(rec.storage_key).unlink(missing_ok=True)159            n += 1160        await session.execute(delete(StoredFile).where(StoredFile.user_id == user_id))161        await session.commit()162    return n163164165# --- signed download links (10 min) --------------------------------------------------------166def sign(file_id: str, ttl_s: int = 600) -> str:167    s = get_settings()168    exp = int(time.time()) + ttl_s169    mac = hmac.new(s.JWT_SECRET.get_secret_value().encode(), f"{file_id}:{exp}".encode(),170                   hashlib.sha256).hexdigest()[:32]171    return f"{exp}.{mac}"172173174def verify_signature(file_id: str, sig: str) -> bool:175    try:176        exp_s, mac = sig.split(".", 1)177        exp = int(exp_s)178    except ValueError:179        return False180    if exp < time.time():181        return False182    s = get_settings()183    expected = hmac.new(s.JWT_SECRET.get_secret_value().encode(), f"{file_id}:{exp}".encode(),184                        hashlib.sha256).hexdigest()[:32]185    return hmac.compare_digest(expected, mac)186