import io import re import zipfile from pathlib import Path from typing import Any from urllib.parse import urlparse import requests from config import Config from task_service import TaskError SESSION = requests.Session() SESSION.trust_env = False def _get(path: str, params: dict[str, Any] | None = None) -> Any: try: response = SESSION.get( f"{Config.CRAWLER_BASE_URL}/api/v1{path}", params={key: value for key, value in (params or {}).items() if value not in (None, "")}, timeout=Config.CRAWLER_TIMEOUT_SECONDS, ) response.raise_for_status() payload = response.json() except (requests.RequestException, ValueError) as exc: raise TaskError("OCR_DASHBOARD_UNAVAILABLE", "OCR 看板数据服务暂时不可用", 502) from exc if not payload.get("success") or int(payload.get("code") or 0) != 200: code = int(payload.get("code") or 500) raise TaskError( "OCR_DASHBOARD_UPSTREAM_ERROR", str(payload.get("msg") or "OCR 看板数据查询失败"), 404 if code == 404 else 422 if code == 400 else 502, ) return payload.get("data") def list_batches() -> list[dict[str, Any]]: return _get("/xhs/ocr/batches") or [] def list_collections(batch_key: str) -> list[dict[str, Any]]: if not batch_key.strip(): raise TaskError("OCR_BATCH_REQUIRED", "batch_key 不能为空", 422) return _get(f"/xhs/ocr/batches/{batch_key.strip()}/collections") or [] def get_summary(batch_key: str = "", collection_key: str = "", start_date: str = "", end_date: str = "") -> dict[str, Any]: return _get("/xhs/ocr/dashboard/summary", { "batch_key": batch_key, "collection_key": collection_key, "start_date": start_date, "end_date": end_date, }) or {} def get_series(batch_key: str = "", start_date: str = "", end_date: str = "") -> dict[str, Any]: return _get("/xhs/ocr/dashboard/series", { "batch_key": batch_key, "start_date": start_date, "end_date": end_date, }) or {"items": [], "subjects": []} def get_comparison(batch_key: str = "", collection_key: str = "", start_date: str = "", end_date: str = "") -> dict[str, Any]: return _get("/xhs/ocr/dashboard/comparison", { "batch_key": batch_key, "collection_key": collection_key, "start_date": start_date, "end_date": end_date, }) or {} def list_results(params: dict[str, Any]) -> dict[str, Any]: allowed = { "batch_key", "collection_key", "start_date", "end_date", "question_code", "search_keyword", "city", "department_name", "minsheng_matched", "cobranded_matched", "is_valid", "page", "page_size", } query = {key: value for key, value in params.items() if key in allowed} page = int(query.get("page") or 1) page_size = int(query.get("page_size") or 50) if page < 1 or not 1 <= page_size <= 200: raise TaskError("OCR_INVALID_PAGINATION", "分页参数无效", 422) query.update({"page": page, "page_size": page_size}) return _get("/xhs/ocr/dashboard/results", query) or { "items": [], "total": 0, "page": page, "page_size": page_size, } def get_result_matches(result_id: int) -> dict[str, Any]: if result_id < 1: raise TaskError("OCR_RESULT_ID_INVALID", "result_id 必须大于 0", 422) return _get(f"/xhs/ocr/results/{result_id}/matches") or {} def list_evidences(params: dict[str, Any]) -> dict[str, Any]: allowed = {"batch_key", "collection_key", "start_date", "end_date", "question_code", "evidence_type", "limit"} query = {key: value for key, value in params.items() if key in allowed} if not (str(query.get("batch_key") or "").strip() or str(query.get("collection_key") or "").strip() or str(query.get("start_date") or "").strip() or str(query.get("end_date") or "").strip()): raise TaskError("OCR_EVIDENCE_SCOPE_REQUIRED", "batch_key/collection_key 或 start_date/end_date 至少传入一种", 422) limit = int(query.get("limit") or 10) if not 1 <= limit <= 100: raise TaskError("OCR_EVIDENCE_LIMIT_INVALID", "limit 必须在 1 到 100 之间", 422) query["limit"] = limit return _get("/xhs/ocr/dashboard/evidences", query) or { "items": [], "total": 0, "limit": limit, "is_complete": False, } def export_evidences(params: dict[str, Any]) -> tuple[io.BytesIO, str, int]: query = dict(params) query["limit"] = 100 data = list_evidences(query) items = data.get("items") or [] if not items: raise TaskError("OCR_EVIDENCE_EMPTY", "当前筛选范围没有可导出的截图", 404) output = io.BytesIO() used_names: set[str] = set() with zipfile.ZipFile(output, "w", zipfile.ZIP_DEFLATED) as archive: for index, item in enumerate(items, 1): url = str(item.get("cos_url") or "").strip() host = (urlparse(url).hostname or "").lower() if not url.startswith("https://") or not host.endswith(".myqcloud.com"): continue try: response = SESSION.get(url, timeout=Config.CRAWLER_TIMEOUT_SECONDS) response.raise_for_status() except requests.RequestException: continue suffix = Path(str(item.get("source_file_name") or urlparse(url).path)).suffix.lower() if suffix not in {".jpg", ".jpeg", ".png", ".webp"}: suffix = ".jpg" parts = [item.get("question_code"), item.get("search_keyword"), item.get("device_code"), item.get("city"), str(item.get("sort_order") or index)] stem = "_".join(re.sub(r"[^0-9A-Za-z\u4e00-\u9fff_-]+", "-", str(part)) for part in parts if part not in (None, "")) name = f"{stem or f'evidence-{index}'}{suffix}" if name in used_names: name = f"{Path(name).stem}-{index}{suffix}" used_names.add(name) archive.writestr(name, response.content) if not used_names: raise TaskError("OCR_EVIDENCE_EXPORT_FAILED", "截图下载失败,未生成导出文件", 502) output.seek(0) scope = str(query.get("question_code") or "all") return output, f"holy-crab-evidences-{scope}.zip", len(used_names)