""" Shared dependencies, guards, and helper utilities for all routers. """ from __future__ import annotations import asyncio import json import os import tempfile import time import uuid from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, Dict, List, Optional, Tuple from fastapi import Depends, File, HTTPException, Request, UploadFile from geoalchemy2.functions import ST_Intersects from shapely.geometry import mapping, shape from shapely.ops import unary_union from sqlalchemy import delete, inspect as sa_inspect from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from ..auth_service import ( ROLE_ADMIN, SESSION_COOKIE_NAME, add_audit_log, get_user_by_session_token, ) from ..config import read_int_env, settings # Trusted proxy IPs whose X-Forwarded-For header is accepted. # Defaults to localhost (nginx on the same host). Override via TRUSTED_PROXY_IPS env var # (comma-separated, e.g. "127.0.0.1,10.0.0.1"). _TRUSTED_PROXY_IPS: frozenset[str] = frozenset( ip.strip() for ip in (settings.TRUSTED_PROXY_IPS or "127.0.0.1").split(",") if ip.strip() ) from ..database import get_db from ..license_service import check_license from ..models import AuthRateLimitORM, AuthUserORM, DinsarTaskBatchORM, DinsarTaskItemORM, PsTaskBatchORM, PsTaskItemORM # --------------------------------------------------------------------------- # Path classification constants # --------------------------------------------------------------------------- _LICENSE_EXEMPT_PATHS = { "/api/license/status", "/api/license/upload", "/api/license/refresh", "/api/health", "/api/auth/login", "/api/auth/logout", "/api/auth/me", } READ_ONLY_METHODS = {"GET", "HEAD", "OPTIONS"} READ_SAFE_POST_PATHS = { "/api/radar-data/search", } PUBLIC_AUTH_PATHS = { "/api/license/status", "/api/health", "/api/auth/login", "/api/auth/logout", } HIGH_RISK_WRITE_PATH_PREFIXES = ( "/api/auth/users", "/api/license/upload", "/api/license/refresh", "/api/workflow/runs", "/api/task-batches/", "/api/tools/", "/api/unpack/run", "/api/monitor/", "/api/scan-data", "/api/scan-dinsar-results", "/api/dinsar-results/", "/api/dinsar-production", "/api/sbas-insar-production", "/api/ai/", "/api/idl/launch-workbench", "/api/idl/worker/", "/api/hazard-points/scan", "/api/radar-data/", ) # --------------------------------------------------------------------------- # Login throttle constants & locks # --------------------------------------------------------------------------- AUTH_LOGIN_MAX_FAILURES = read_int_env("AUTH_LOGIN_MAX_FAILURES", 5, minimum=1, maximum=100) AUTH_LOGIN_WINDOW_SECONDS = read_int_env("AUTH_LOGIN_WINDOW_SECONDS", 900, minimum=5, maximum=86400) AUTH_LOGIN_LOCK_SECONDS = read_int_env("AUTH_LOGIN_LOCK_SECONDS", 900, minimum=5, maximum=86400) LOGIN_THROTTLE_CLEANUP_INTERVAL_SECONDS = read_int_env( "AUTH_LOGIN_CLEANUP_INTERVAL_SECONDS", 300, minimum=30, maximum=3600, ) _LOGIN_THROTTLE_LOCK = asyncio.Lock() _LOGIN_THROTTLE_LAST_CLEANUP_MONO = 0.0 _LICENSE_UPLOAD_LOCK = asyncio.Lock() MAX_LICENSE_UPLOAD_BYTES = read_int_env( "MAX_LICENSE_UPLOAD_BYTES", 1024 * 1024, minimum=1024, maximum=20 * 1024 * 1024, ) # --------------------------------------------------------------------------- # Statistics cache # --------------------------------------------------------------------------- STATS_CACHE_TTL_SECONDS = read_int_env( "STATS_CACHE_TTL_SECONDS", 120, minimum=0, maximum=3600, ) _STATS_CACHE_LOCK = asyncio.Lock() _STATS_CACHE_DATA: Optional[Dict[str, Any]] = None _STATS_CACHE_EXPIRES_AT = 0.0 _STATS_CACHE_GENERATED_AT_UTC: Optional[str] = None # --------------------------------------------------------------------------- # AOI token store # --------------------------------------------------------------------------- _AOI_TOKEN_STORE: Dict[str, Dict[str, Any]] = {} _AOI_TOKEN_LOCK = asyncio.Lock() AOI_TOKEN_TTL_SECONDS = read_int_env( "AOI_TOKEN_TTL_SECONDS", 1800, minimum=60, maximum=24 * 3600, ) AOI_UPLOAD_MAX_FILES = read_int_env( "AOI_UPLOAD_MAX_FILES", 10, minimum=1, maximum=200, ) AOI_UPLOAD_MAX_SINGLE_FILE_BYTES = read_int_env( "AOI_UPLOAD_MAX_SINGLE_FILE_BYTES", 20 * 1024 * 1024, minimum=1024, maximum=500 * 1024 * 1024, ) AOI_UPLOAD_MAX_TOTAL_BYTES = max( read_int_env( "AOI_UPLOAD_MAX_TOTAL_BYTES", 100 * 1024 * 1024, minimum=1024, maximum=2 * 1024 * 1024 * 1024, ), AOI_UPLOAD_MAX_SINGLE_FILE_BYTES, ) AOI_UPLOAD_STREAM_CHUNK_BYTES = 1024 * 1024 _SHAPEFILE_READ_LOCK = asyncio.Lock() # --------------------------------------------------------------------------- # Region index caches # --------------------------------------------------------------------------- _REGION_CHILDREN_CACHE: Optional[Dict[str, List[Dict[str, Any]]]] = None _REGION_BY_ID_CACHE: Optional[Dict[str, Dict[str, Any]]] = None _REGION_GEOMETRY_BY_ID_CACHE: Optional[Dict[str, List[Dict[str, Any]]]] = None # --------------------------------------------------------------------------- # Export/root-dir path validation # --------------------------------------------------------------------------- import logging as _logging _dep_logger = _logging.getLogger(__name__) # Optional whitelist of allowed export directories. # Set ALLOWED_EXPORT_DIRS in .env as comma-separated paths to restrict exports. _ALLOWED_EXPORT_DIRS_RAW = (settings.ALLOWED_EXPORT_DIRS or "").strip() ALLOWED_EXPORT_DIRS: Optional[List[str]] = None if _ALLOWED_EXPORT_DIRS_RAW: ALLOWED_EXPORT_DIRS = [ os.path.normpath(p.strip()) for p in _ALLOWED_EXPORT_DIRS_RAW.split(",") if p.strip() ] _SYSTEM_DIRS_WIN = { os.path.normpath(p) for p in ["C:\\Windows", "C:\\Program Files", "C:\\Program Files (x86)", "C:\\ProgramData"] } _SYSTEM_DIRS_UNIX = {"/bin", "/sbin", "/usr", "/etc", "/lib", "/lib64", "/boot", "/proc", "/sys"} def _is_system_directory(path: str) -> bool: """Check if path points to a system directory.""" normed = os.path.normpath(path) for sd in _SYSTEM_DIRS_WIN | _SYSTEM_DIRS_UNIX: if normed == sd or normed.startswith(sd + os.sep): return True return False def _validate_export_path(path: str, param_name: str = "target_dir") -> str: """ Validate an export/destination path: - Rejects path traversal (..) - Rejects system directories - Optionally checks against ALLOWED_EXPORT_DIRS whitelist Returns the normalized path. """ if not path or not path.strip(): raise HTTPException(status_code=400, detail=f"{param_name} 不能为空") normed = os.path.normpath(path.strip()) if ".." in normed.split(os.sep): raise HTTPException(status_code=400, detail=f"{param_name} 包含非法路径遍历") if _is_system_directory(normed): raise HTTPException(status_code=400, detail=f"{param_name} 指向系统目录，操作被拒绝") if ALLOWED_EXPORT_DIRS is not None: if not any(normed.startswith(allowed) for allowed in ALLOWED_EXPORT_DIRS): raise HTTPException(status_code=400, detail=f"{param_name} 不在允许的导出目录内") return normed def _validate_root_dir(path: str, param_name: str = "root_dir") -> str: """ Validate a root directory parameter: - Rejects path traversal (..) - Rejects system directories - Verifies directory exists Returns the normalized path. """ if not path or not path.strip(): raise HTTPException(status_code=400, detail=f"{param_name} 不能为空") normed = os.path.normpath(path.strip()) if ".." in normed.split(os.sep): raise HTTPException(status_code=400, detail=f"{param_name} 包含非法路径遍历") if _is_system_directory(normed): raise HTTPException(status_code=400, detail=f"{param_name} 指向系统目录，操作被拒绝") if not os.path.isdir(normed): raise HTTPException(status_code=400, detail=f"{param_name} 目录不存在") return normed # --------------------------------------------------------------------------- # Path helpers # --------------------------------------------------------------------------- def _normalize_request_path(path: str) -> str: normalized = (path or "").rstrip("/") return normalized or "/" def _format_size_limit(size_bytes: int) -> str: if size_bytes >= 1024 * 1024: return f"{size_bytes / (1024 * 1024):.1f} MB" if size_bytes >= 1024: return f"{size_bytes / 1024:.1f} KB" return f"{size_bytes} B" def _is_public_auth_path(path: str) -> bool: return _normalize_request_path(path) in PUBLIC_AUTH_PATHS def _is_read_safe_post_path(path: str, method: str) -> bool: return (method or "").upper() == "POST" and _normalize_request_path(path) in READ_SAFE_POST_PATHS def _is_read_only_operation(path: str, method: str) -> bool: upper_method = (method or "").upper() if upper_method in READ_ONLY_METHODS: return True return _is_read_safe_post_path(path, upper_method) def _is_high_risk_write_path(path: str, method: str) -> bool: normalized = _normalize_request_path(path) upper_method = (method or "").upper() if _is_read_only_operation(normalized, upper_method): return False return any(normalized.startswith(prefix) for prefix in HIGH_RISK_WRITE_PATH_PREFIXES) def _get_client_ip(request: Request) -> Optional[str]: direct_ip = request.client.host if request.client else None if direct_ip in _TRUSTED_PROXY_IPS: forwarded = request.headers.get("x-forwarded-for") if forwarded: return forwarded.split(",")[0].strip() return direct_ip # --------------------------------------------------------------------------- # Audit log helper # --------------------------------------------------------------------------- async def _add_operation_audit_log( db: AsyncSession, request: Request, action: str, resource: str, detail: Optional[Dict[str, Any]] = None, user: Optional[AuthUserORM] = None, user_id: Optional[int] = None, username: Optional[str] = None, ) -> None: actor = user or getattr(request.state, "current_user", None) resolved_user_id = user_id resolved_username = username if actor is not None: if resolved_user_id is None: try: identity = sa_inspect(actor).identity if identity and len(identity) > 0: resolved_user_id = int(identity[0]) except Exception: try: resolved_user_id = actor.id except Exception: resolved_user_id = None if resolved_username is None: try: actor_state_dict = sa_inspect(actor).dict resolved_username = actor_state_dict.get("username") except Exception: try: resolved_username = actor.username except Exception: resolved_username = None await add_audit_log( db, action=action, user=None, user_id=resolved_user_id, username=resolved_username, resource=resource, detail=detail or None, ip_address=_get_client_ip(request), ) # --------------------------------------------------------------------------- # License guard # --------------------------------------------------------------------------- def _require_license(request: Request): """ 授权校验：未授权时拒绝所有 API。 使用精确路径集合匹配，防止 endswith 绕过攻击。 """ path = (request.url.path or "").rstrip("/") or "/" if path in _LICENSE_EXEMPT_PATHS: return result = check_license() if not result.get("ok"): raise HTTPException(status_code=403, detail=f"License required: {result.get('reason')}") # --------------------------------------------------------------------------- # Auth guards # --------------------------------------------------------------------------- async def _require_auth(request: Request, db: AsyncSession = Depends(get_db)): path = _normalize_request_path(request.url.path) if _is_public_auth_path(path): return method = request.method.upper() token = request.cookies.get(SESSION_COOKIE_NAME) user = await get_user_by_session_token(db, token) if not user: if _is_high_risk_write_path(path, method): await add_audit_log( db, action="write_auth_required", resource=path, detail={"method": method}, ip_address=_get_client_ip(request), ) await db.commit() raise HTTPException(status_code=401, detail="Authentication required.") if (not _is_read_only_operation(path, method)) and user.role != ROLE_ADMIN: if _is_high_risk_write_path(path, method): await add_audit_log( db, action="write_blocked_readonly", user=user, resource=path, detail={"method": method, "role": user.role}, ip_address=_get_client_ip(request), ) await db.commit() raise HTTPException(status_code=403, detail="Read-only account cannot perform this operation.") if _is_high_risk_write_path(path, method): await add_audit_log( db, action="write_access_granted", user=user, resource=path, detail={"method": method}, ip_address=_get_client_ip(request), ) await db.commit() refreshed_user = await get_user_by_session_token(db, token) if not refreshed_user: raise HTTPException(status_code=401, detail="Authentication required.") user = refreshed_user request.state.current_user = user async def _get_current_user(request: Request, db: AsyncSession = Depends(get_db)) -> AuthUserORM: user = getattr(request.state, "current_user", None) if user: return user token = request.cookies.get(SESSION_COOKIE_NAME) user = await get_user_by_session_token(db, token) if not user: raise HTTPException(status_code=401, detail="Authentication required.") request.state.current_user = user return user async def _get_optional_session_user( request: Request, db: AsyncSession, ) -> Optional[AuthUserORM]: cached_user = getattr(request.state, "current_user", None) if cached_user is not None: return cached_user token = request.cookies.get(SESSION_COOKIE_NAME) user = await get_user_by_session_token(db, token) if user is not None: request.state.current_user = user return user async def _require_admin(current_user: AuthUserORM = Depends(_get_current_user)) -> AuthUserORM: if current_user.role != ROLE_ADMIN: raise HTTPException(status_code=403, detail="Admin role required.") return current_user # --------------------------------------------------------------------------- # Login throttle helpers # --------------------------------------------------------------------------- def _build_login_throttle_key(username: str, client_ip: Optional[str]) -> str: # Throttle by username only to prevent IP-switching bypass normalized_username = (username or "").strip().lower() or "" return f"user:{normalized_username}" def _remaining_seconds(deadline: datetime, now: datetime) -> int: remaining = max(0.0, (deadline - now).total_seconds()) rounded = int(remaining) if rounded < remaining: rounded += 1 return max(1, rounded) def _login_throttle_retention_seconds() -> int: return max(AUTH_LOGIN_WINDOW_SECONDS, AUTH_LOGIN_LOCK_SECONDS) * 2 async def _maybe_cleanup_login_throttle_records(db: AsyncSession) -> None: global _LOGIN_THROTTLE_LAST_CLEANUP_MONO if LOGIN_THROTTLE_CLEANUP_INTERVAL_SECONDS <= 0: return now_mono = time.monotonic() if ( _LOGIN_THROTTLE_LAST_CLEANUP_MONO > 0 and now_mono - _LOGIN_THROTTLE_LAST_CLEANUP_MONO < LOGIN_THROTTLE_CLEANUP_INTERVAL_SECONDS ): return now = datetime.now(timezone.utc).replace(tzinfo=None) cutoff = now - timedelta(seconds=_login_throttle_retention_seconds()) await db.execute( delete(AuthRateLimitORM).where( AuthRateLimitORM.updated_at < cutoff, (AuthRateLimitORM.locked_until.is_(None)) | (AuthRateLimitORM.locked_until <= now), ) ) _LOGIN_THROTTLE_LAST_CLEANUP_MONO = now_mono async def _get_login_retry_after_seconds(throttle_key: str, db: AsyncSession) -> int: async with _LOGIN_THROTTLE_LOCK: await _maybe_cleanup_login_throttle_records(db) now = datetime.now(timezone.utc).replace(tzinfo=None) result = await db.execute( select(AuthRateLimitORM).where(AuthRateLimitORM.throttle_key == throttle_key) ) record = result.scalar_one_or_none() if not record or not record.locked_until: return 0 if record.locked_until <= now: return 0 return _remaining_seconds(record.locked_until, now) async def _record_login_failure(throttle_key: str, db: AsyncSession) -> int: async with _LOGIN_THROTTLE_LOCK: await _maybe_cleanup_login_throttle_records(db) now = datetime.now(timezone.utc).replace(tzinfo=None) failure_cutoff = now - timedelta(seconds=AUTH_LOGIN_WINDOW_SECONDS) result = await db.execute( select(AuthRateLimitORM).where(AuthRateLimitORM.throttle_key == throttle_key) ) record = result.scalar_one_or_none() if record is None: record = AuthRateLimitORM(throttle_key=throttle_key, failure_timestamps=[]) db.add(record) timestamps = [ts for ts in (record.failure_timestamps or []) if ts > failure_cutoff.timestamp()] timestamps.append(now.timestamp()) if len(timestamps) >= AUTH_LOGIN_MAX_FAILURES: locked_until = now + timedelta(seconds=AUTH_LOGIN_LOCK_SECONDS) record.locked_until = locked_until record.failure_timestamps = [] await db.flush() return _remaining_seconds(locked_until, now) record.failure_timestamps = timestamps record.locked_until = None await db.flush() return 0 async def _clear_login_failure_state(throttle_key: str, db: AsyncSession) -> None: async with _LOGIN_THROTTLE_LOCK: result = await db.execute( select(AuthRateLimitORM).where(AuthRateLimitORM.throttle_key == throttle_key) ) record = result.scalar_one_or_none() if record is not None: record.failure_timestamps = [] record.locked_until = None await db.flush() # --------------------------------------------------------------------------- # Batch summary helpers # --------------------------------------------------------------------------- async def _refresh_dinsar_batch_summary(db: AsyncSession, batch_id: str) -> None: from sqlalchemy import func total_res = await db.execute( select(func.count(DinsarTaskItemORM.id)).where(DinsarTaskItemORM.batch_id == batch_id) ) completed_res = await db.execute( select(func.count(DinsarTaskItemORM.id)).where( DinsarTaskItemORM.batch_id == batch_id, func.upper(func.coalesce(DinsarTaskItemORM.status, "")) == "COMPLETED", ) ) total = total_res.scalar_one() or 0 completed = completed_res.scalar_one() or 0 status = "PENDING" if total > 0 and completed == total: status = "COMPLETED" elif completed > 0: status = "IN_PROGRESS" await db.execute( DinsarTaskBatchORM.__table__.update() .where(DinsarTaskBatchORM.batch_id == batch_id) .values(total_items=total, completed_items=completed, status=status) ) async def _refresh_ps_batch_summary(db: AsyncSession, batch_id: str) -> None: from sqlalchemy import func total_res = await db.execute( select(func.count(PsTaskItemORM.id)).where(PsTaskItemORM.batch_id == batch_id) ) completed_res = await db.execute( select(func.count(PsTaskItemORM.id)).where( PsTaskItemORM.batch_id == batch_id, func.upper(func.coalesce(PsTaskItemORM.status, "")) == "COMPLETED", ) ) total = total_res.scalar_one() or 0 completed = completed_res.scalar_one() or 0 status = "PENDING" if total > 0 and completed == total: status = "COMPLETED" elif completed > 0: status = "IN_PROGRESS" await db.execute( PsTaskBatchORM.__table__.update() .where(PsTaskBatchORM.batch_id == batch_id) .values(total_items=total, completed_items=completed, status=status) ) # --------------------------------------------------------------------------- # AOI helpers # --------------------------------------------------------------------------- def _infer_region_level(tree_id: str) -> str: depth = len((tree_id or "").split("-")) if depth == 1: return "country" if depth == 2: return "province" if depth == 3: return "city" if depth == 4: return "district" return "unknown" def _normalize_region_node(raw: Dict[str, Any]) -> Optional[Dict[str, Any]]: tree_id = str(raw.get("treeID") or "").strip() if not tree_id: return None parent_raw = raw.get("parent") parent_tree_id = str(parent_raw).strip() if parent_raw is not None else None if parent_tree_id == "": parent_tree_id = None name = str(raw.get("name") or tree_id).strip() return { "tree_id": tree_id, "parent_tree_id": parent_tree_id, "name": name, "level": _infer_region_level(tree_id), } def _load_region_index() -> None: global _REGION_CHILDREN_CACHE, _REGION_BY_ID_CACHE if _REGION_CHILDREN_CACHE is not None and _REGION_BY_ID_CACHE is not None: return index_path = Path(settings.AOI_REGION_INDEX_FILE or "") if not index_path.is_file(): raise HTTPException( status_code=503, detail=f"行政区索引文件不存在，请检查 AOI_REGION_INDEX_FILE: {index_path}", ) try: raw_data = json.loads(index_path.read_text(encoding="utf-8")) except Exception as exc: raise HTTPException(status_code=500, detail=f"加载行政区索引失败: {exc}") from exc if not isinstance(raw_data, list): raise HTTPException(status_code=500, detail="行政区索引文件格式错误：应为 JSON 数组。") children_cache: Dict[str, List[Dict[str, Any]]] = {} node_cache: Dict[str, Dict[str, Any]] = {} for item in raw_data: if not isinstance(item, dict): continue normalized = _normalize_region_node(item) if not normalized: continue tree_id = normalized["tree_id"] parent_tree_id = normalized["parent_tree_id"] node_cache[tree_id] = normalized if parent_tree_id: children_cache.setdefault(parent_tree_id, []).append(normalized) for parent_tree_id in children_cache: children_cache[parent_tree_id].sort(key=lambda row: row["tree_id"]) _REGION_CHILDREN_CACHE = children_cache _REGION_BY_ID_CACHE = node_cache def _load_region_geometry_index() -> None: global _REGION_GEOMETRY_BY_ID_CACHE if _REGION_GEOMETRY_BY_ID_CACHE is not None: return geometry_path = Path(settings.AOI_REGION_GEOJSON_FILE or "") if not geometry_path.is_file(): raise HTTPException( status_code=503, detail=( "行政区边界文件不存在，请准备标准 GeoJSON（FeatureCollection），" f"并配置 AOI_REGION_GEOJSON_FILE: {geometry_path}" ), ) try: raw_data = json.loads(geometry_path.read_text(encoding="utf-8")) except Exception as exc: raise HTTPException(status_code=500, detail=f"加载行政区边界数据失败: {exc}") from exc if isinstance(raw_data, dict) and raw_data.get("type") == "FeatureCollection": features = raw_data.get("features", []) elif isinstance(raw_data, list): features = raw_data else: raise HTTPException(status_code=500, detail="行政区边界文件格式错误：应为 FeatureCollection 或 Feature 数组。") feature_index: Dict[str, List[Dict[str, Any]]] = {} for feature in features: if not isinstance(feature, dict): continue if feature.get("type") != "Feature": continue properties = feature.get("properties") or {} tree_id = str( properties.get("treeID") or properties.get("tree_id") or properties.get("treeId") or "" ).strip() if not tree_id: continue feature_index.setdefault(tree_id, []).append(feature) if not feature_index: raise HTTPException( status_code=500, detail="行政区边界文件中未找到可用 treeID 字段（支持 treeID/tree_id/treeId）。", ) _REGION_GEOMETRY_BY_ID_CACHE = feature_index def _normalize_geojson_to_feature_collection(payload: Any) -> Dict[str, Any]: if not payload: raise HTTPException(status_code=400, detail="AOI GeoJSON 不能为空。") if not isinstance(payload, dict): raise HTTPException(status_code=400, detail="AOI GeoJSON 必须是对象。") payload_type = payload.get("type") if payload_type == "FeatureCollection": features = payload.get("features") or [] if not isinstance(features, list): raise HTTPException(status_code=400, detail="FeatureCollection.features 必须是数组。") return {"type": "FeatureCollection", "features": features} if payload_type == "Feature": return {"type": "FeatureCollection", "features": [payload]} if payload_type in {"Polygon", "MultiPolygon", "LineString", "MultiLineString", "Point", "MultiPoint"}: return {"type": "FeatureCollection", "features": [{"type": "Feature", "properties": {}, "geometry": payload}]} raise HTTPException(status_code=400, detail=f"不支持的 AOI GeoJSON 类型: {payload_type}") def _feature_collection_to_union_geometry(feature_collection: Dict[str, Any]): features = feature_collection.get("features") or [] geometries = [] for feature in features: if not isinstance(feature, dict): continue geometry = feature.get("geometry") if not geometry: continue try: geom = shape(geometry) except Exception: continue if geom.is_empty: continue geometries.append(geom) if not geometries: raise HTTPException(status_code=400, detail="AOI GeoJSON 中未解析到有效几何。") merged = unary_union(geometries) if merged.is_empty: raise HTTPException(status_code=400, detail="AOI 几何为空，无法用于筛选。") return merged def _parse_aoi_geojson_form_value(aoi_geojson: Optional[str]) -> Optional[Tuple[str, Dict[str, Any]]]: if not aoi_geojson: return None try: payload = json.loads(aoi_geojson) except Exception as exc: raise HTTPException(status_code=400, detail=f"AOI GeoJSON 解析失败: {exc}") from exc feature_collection = _normalize_geojson_to_feature_collection(payload) merged_geometry = _feature_collection_to_union_geometry(feature_collection) return merged_geometry.wkt, feature_collection def _read_aoi_shapefile_with_restore_shx(shp_path: str): import geopandas as gpd previous_restore_shx = os.environ.get("SHAPE_RESTORE_SHX") os.environ["SHAPE_RESTORE_SHX"] = "YES" try: return gpd.read_file(shp_path, engine="pyogrio") finally: if previous_restore_shx is None: os.environ.pop("SHAPE_RESTORE_SHX", None) else: os.environ["SHAPE_RESTORE_SHX"] = previous_restore_shx async def _parse_aoi_from_files(files: Optional[List[UploadFile]]) -> Optional[Tuple[str, Dict[str, Any]]]: if not files: return None valid_files = [file for file in files if file and file.filename and str(file.filename).strip()] if not valid_files: return None if len(valid_files) > AOI_UPLOAD_MAX_FILES: raise HTTPException( status_code=400, detail=f"AOI 上传文件数量超限，最多允许 {AOI_UPLOAD_MAX_FILES} 个文件。", ) with tempfile.TemporaryDirectory() as temp_dir: shp_path: Optional[str] = None geojson_payload: Optional[Dict[str, Any]] = None total_bytes = 0 for index, file in enumerate(valid_files): raw_name = str(file.filename).strip() base_name = os.path.basename(raw_name) or f"upload_{index}" dest_name = f"{index:02d}_{base_name}" dest_path = os.path.join(temp_dir, dest_name) file_bytes = 0 try: with open(dest_path, "wb") as buffer: while True: chunk = await file.read(AOI_UPLOAD_STREAM_CHUNK_BYTES) if not chunk: break chunk_size = len(chunk) file_bytes += chunk_size total_bytes += chunk_size if file_bytes > AOI_UPLOAD_MAX_SINGLE_FILE_BYTES: raise HTTPException( status_code=400, detail=( f"AOI 文件 {base_name} 超过单文件大小限制（" f"{_format_size_limit(AOI_UPLOAD_MAX_SINGLE_FILE_BYTES)}）。" ), ) if total_bytes > AOI_UPLOAD_MAX_TOTAL_BYTES: raise HTTPException( status_code=400, detail=( "AOI 上传文件总大小超限（" f"{_format_size_limit(AOI_UPLOAD_MAX_TOTAL_BYTES)}）。" ), ) buffer.write(chunk) finally: await file.close() lower_dest = dest_path.lower() if lower_dest.endswith(".shp"): shp_path = dest_path elif lower_dest.endswith(".geojson") or lower_dest.endswith(".json"): try: geojson_payload = json.loads(Path(dest_path).read_text(encoding="utf-8")) except UnicodeDecodeError: geojson_payload = json.loads(Path(dest_path).read_text(encoding="gbk")) if shp_path: try: async with _SHAPEFILE_READ_LOCK: gdf = await asyncio.to_thread(_read_aoi_shapefile_with_restore_shx, shp_path) except Exception as exc: raise HTTPException( status_code=400, detail=( "AOI Shapefile 读取失败。系统已尝试自动恢复缺失的 .shx 索引；" f"请确认已上传 .shp/.dbf/.prj/.shx 或可恢复的标准 Shapefile。原始错误: {exc}" ), ) from exc if gdf.crs and gdf.crs.to_epsg() != 4326: gdf = gdf.to_crs(epsg=4326) feature_collection = json.loads(gdf.to_json()) merged_geometry = _feature_collection_to_union_geometry(feature_collection) return merged_geometry.wkt, feature_collection if geojson_payload: feature_collection = _normalize_geojson_to_feature_collection(geojson_payload) merged_geometry = _feature_collection_to_union_geometry(feature_collection) return merged_geometry.wkt, feature_collection return None def _cleanup_expired_aoi_tokens(now: float) -> None: expired_tokens = [ token for token, payload in _AOI_TOKEN_STORE.items() if float(payload.get("expires_at", 0.0)) <= now ] for token in expired_tokens: _AOI_TOKEN_STORE.pop(token, None) async def _store_aoi_token(aoi_wkt: str, feature_collection: Dict[str, Any]) -> str: token = uuid.uuid4().hex now = time.monotonic() async with _AOI_TOKEN_LOCK: _cleanup_expired_aoi_tokens(now) _AOI_TOKEN_STORE[token] = { "aoi_wkt": aoi_wkt, "aoi_geojson": feature_collection, "expires_at": now + AOI_TOKEN_TTL_SECONDS, } return token async def _get_aoi_from_token(aoi_token: Optional[str]) -> Optional[Tuple[str, Dict[str, Any]]]: if not aoi_token: return None normalized_token = str(aoi_token).strip() if not normalized_token: return None now = time.monotonic() async with _AOI_TOKEN_LOCK: _cleanup_expired_aoi_tokens(now) payload = _AOI_TOKEN_STORE.get(normalized_token) if not payload: return None payload["expires_at"] = now + AOI_TOKEN_TTL_SECONDS return payload.get("aoi_wkt"), payload.get("aoi_geojson") def _resolve_region_aoi_payload(tree_id: str) -> Dict[str, Any]: normalized_tree_id = (tree_id or "").strip() if not normalized_tree_id: raise HTTPException(status_code=400, detail="tree_id 不能为空。") _load_region_index() _load_region_geometry_index() node = (_REGION_BY_ID_CACHE or {}).get(normalized_tree_id) if not node: raise HTTPException(status_code=404, detail=f"未找到行政区: {normalized_tree_id}") geometry_index = _REGION_GEOMETRY_BY_ID_CACHE or {} def _collect_features_by_tree_id(target_tree_id: str) -> Tuple[List[Dict[str, Any]], str]: exact_features = list(geometry_index.get(target_tree_id) or []) if exact_features: return exact_features, "exact" descendant_features: List[Dict[str, Any]] = [] prefix = f"{target_tree_id}-" for feature_tree_id, feature_list in geometry_index.items(): if feature_tree_id.startswith(prefix): descendant_features.extend(feature_list) if descendant_features: return descendant_features, "descendants" return [], "none" matched_tree_id = normalized_tree_id features, source = _collect_features_by_tree_id(normalized_tree_id) if not features: current_tree_id = node.get("parent_tree_id") while current_tree_id: ancestor_features, ancestor_source = _collect_features_by_tree_id(current_tree_id) if ancestor_features: features = ancestor_features matched_tree_id = current_tree_id source = f"ancestor_{ancestor_source}" break current_node = (_REGION_BY_ID_CACHE or {}).get(current_tree_id) current_tree_id = current_node.get("parent_tree_id") if current_node else None if not features: raise HTTPException( status_code=404, detail=( f"行政区 {normalized_tree_id} 未匹配到边界数据。" "请检查 AOI_REGION_GEOJSON_FILE 中 features[*].properties.treeID。" ), ) raw_feature_collection = { "type": "FeatureCollection", "features": features, } merged_geometry = _feature_collection_to_union_geometry(raw_feature_collection) merged_feature_collection = { "type": "FeatureCollection", "features": [ { "type": "Feature", "properties": { "treeID": normalized_tree_id, "matched_treeID": matched_tree_id, "name": node.get("name"), "level": node.get("level"), "source": source, }, "geometry": mapping(merged_geometry), } ], } return { "tree_id": normalized_tree_id, "name": node.get("name"), "level": node.get("level"), "source": source, "matched_tree_id": matched_tree_id, "feature_count": len(features), "aoi_geojson": merged_feature_collection, }