contracts-flask/deploy/services/upload.py

"""Upload service — accept file, store to DB, parse."""
import uuid, base64, json, io, cgi, os
from db import documents, contracts, supplements
from .parse import parse_file

MAX_FILE_SIZE = 100 * 1024 * 1024  # 100 MB


# ── Чистые функции (тестируемы без HTTP) ──────────────────────

def parse_multipart(raw_bytes: bytes, content_type: str) -> dict:
    """Разобрать multipart/form-data → {filename, file_data, contract_id, batch_id, zip_source}.
    Тестируется без HTTP — передать bytes.
    """
    environ = {
        "REQUEST_METHOD": "POST",
        "CONTENT_TYPE": content_type,
        "CONTENT_LENGTH": str(len(raw_bytes)),
    }
    fs = cgi.FieldStorage(fp=io.BytesIO(raw_bytes), environ=environ, keep_blank_values=True)

    filename = None
    file_data = None
    contract_id = ""
    batch_id = None
    zip_source = None

    if "files" in fs:
        item = fs["files"]
        if isinstance(item, list):
            item = item[0]
        filename = os.path.basename(item.filename) if item.filename else None
        if filename and (".." in filename or "/" in filename or "\\" in filename):
            raise ValueError("invalid filename")
        file_data = item.file.read() if hasattr(item, "file") else item.value
        if isinstance(file_data, str):
            file_data = file_data.encode("utf-8")

    if "contract_id" in fs:
        cid = fs.getfirst("contract_id", "")
        if cid:
            try:
                uuid.UUID(cid)
                contract_id = cid
            except (ValueError, AttributeError):
                pass

    if "batch_id" in fs:
        bid = fs.getfirst("batch_id", None)
        if bid:
            try:
                uuid.UUID(bid)
                batch_id = bid
            except (ValueError, AttributeError):
                pass

    if "zip_source" in fs:
        raw_zip = fs.getfirst("zip_source", "")
        if raw_zip:
            zip_source = os.path.basename(raw_zip)[:255]

    return {
        "filename": filename,
        "file_data": file_data,
        "contract_id": contract_id or "",
        "batch_id": batch_id,
        "zip_source": zip_source,
    }


def store_document(filename: str, file_data: bytes, contract_id: str = "",
                   batch_id: str = None, zip_source: str = None) -> dict:
    """Сохранить документ в БД + распарсить. Возвращает {ok, doc_id, contract_id, parsed, ...}.
    Тестируется с реальной или замоканной БД.
    """
    if not filename or not file_data:
        return {"ok": False, "error": "no file"}

    mime = _mime_for(filename)
    b64 = base64.b64encode(file_data).decode()

    if contract_id:
        try:
            supplements.delete_by_document(contract_id, filename)
        except Exception:
            pass

    doc = documents.insert(filename, mime, b64, batch_id=batch_id, zip_source=zip_source)

    if not contract_id:
        from datetime import datetime
        now = datetime.now()
        c = contracts.insert(f"б/н {now.strftime('%Y%m%d')}-{now.strftime('%H%M')}")
        contract_id = c["id"]
        supp_type = "initial"
    else:
        supp_type = "additional"

    supplements.insert(contract_id, doc["id"], supp_type)

    parsed = parse_file(filename, file_data)
    if parsed and parsed.get("status") == "parsed":
        documents.set_parsed(doc["id"], parsed["elements"])
    elif parsed and parsed.get("status") == "error":
        documents.set_error(doc["id"], parsed.get("error", "parse failed"))

    return {
        "ok": True,
        "doc_id": doc["id"],
        "contract_id": contract_id,
        "filename": filename,
        "size": len(file_data),
        "parsed": parsed,
    }


# ── HTTP-handler (вызывает parse_multipart + store_document) ───

def handle_upload(rfile, content_type, content_length):
    """HTTP-handler: читает rfile → parse_multipart → store_document."""
    try:
        cl = int(content_length)
    except (TypeError, ValueError):
        return {"ok": False, "error": "invalid content-length"}
    if cl <= 0:
        return {"ok": False, "error": "empty request"}
    if cl > MAX_FILE_SIZE:
        return {"ok": False, "error": f"file too large (max {MAX_FILE_SIZE // 1024 // 1024}MB)"}

    raw_bytes = rfile.read(cl)

    try:
        parts = parse_multipart(raw_bytes, content_type)
    except ValueError as e:
        return {"ok": False, "error": str(e)}

    if not parts["filename"] or not parts["file_data"]:
        return {"ok": False, "error": "no file in request"}

    return store_document(
        filename=parts["filename"],
        file_data=parts["file_data"],
        contract_id=parts["contract_id"],
        batch_id=parts["batch_id"],
        zip_source=parts["zip_source"],
    )


def _mime_for(filename):
    ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    mime_map = {
        "pdf": "application/pdf",
        "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "doc": "application/msword",
        "zip": "application/zip",
    }
    return mime_map.get(ext, "application/octet-stream")