from functools import wraps from datetime import datetime from flask import request, jsonify, session from backend.core.security import verify_access_token from backend.database.db_manager import get_user_db_conn def jwt_required(f): """Decorator: require valid JWT OR session auth.""" @wraps(f) def decorated(*args, **kwargs): # Try JWT first auth_header = request.headers.get("Authorization", "") if auth_header.startswith("Bearer "): token = auth_header[7:] payload = verify_access_token(token) if payload: request._jwt_user = { "id": int(payload["sub"]), "username": payload["username"], "vip_status": payload["vip"] } return f(*args, **kwargs) # Fall back to session auth if "user_id" in session: request._jwt_user = { "id": session["user_id"], "username": session["username"], "vip_status": session.get("vip_status", 0) } return f(*args, **kwargs) return jsonify({"error": "Authentication required"}), 401 return decorated def require_api_key_auth(f): """Decorator: require valid Developer API Key or VIP Validation Code auth.""" @wraps(f) def decorated(*args, **kwargs): # 1. Try to find the key in Authorization header auth_header = request.headers.get("Authorization", "") api_key = "" if auth_header.startswith("Bearer "): api_key = auth_header[7:].strip() # If the api_key from Authorization header looks like a JWT token (not sk-tc- and not a valid VIP code), # clear it so we fall back to other headers or the body. from backend.config import Config if api_key and not api_key.startswith("sk-tc-") and api_key not in Config.VALID_VIP_CODES: api_key = "" # 2. Try X-VIP-Key header if not api_key: api_key = request.headers.get("X-VIP-Key", "").strip() # 3. Try X-VIP-Code header if not api_key: api_key = request.headers.get("X-VIP-Code", "").strip() # 4. Try request JSON body if not api_key: try: data = request.json or {} api_key = data.get("vip_key", "") or data.get("vip_code", "") or data.get("api_key", "") if isinstance(api_key, str): api_key = api_key.strip() else: api_key = "" except: pass # 5. Check if it's empty if not api_key: return jsonify({"error": "Missing API Key or VIP Code. Please authenticate."}), 401 # 6. Check if it is a valid static VIP code from backend.config import Config if api_key in Config.VALID_VIP_CODES: # Bypassed - treat as static VIP access with unlimited balance request.api_key = api_key request.api_user_id = None request.api_balance = 999999.0 return f(*args, **kwargs) # 7. Otherwise, it must be a developer key starting with sk-tc- if not api_key.startswith("sk-tc-"): return jsonify({"error": "Invalid API key format. Key must start with 'sk-tc-' or be a valid VIP Code."}), 401 conn = get_user_db_conn() key_record = conn.execute( "SELECT k.*, u.api_balance FROM api_keys k JOIN users u ON k.user_id = u.id WHERE k.api_key = ? AND k.status = 'active'", (api_key,) ).fetchone() if not key_record: conn.close() return jsonify({"error": "API Key not found, inactive, or revoked."}), 401 balance = key_record["api_balance"] if key_record["api_balance"] is not None else 0.0 if balance <= 0.0: conn.close() return jsonify({ "error": f"Tài khoản hết số dư API (Số dư hiện tại: {balance:.2f}đ). Vui lòng nạp thêm tiền tại website để tiếp tục sử dụng." }), 402 # Update last_used_at now = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") conn.execute("UPDATE api_keys SET last_used_at = ? WHERE api_key = ?", (now, api_key)) conn.commit() conn.close() request.api_key = api_key request.api_user_id = key_record["user_id"] request.api_balance = balance return f(*args, **kwargs) return decorated def get_current_user(): """Get current user from JWT or session.""" if hasattr(request, '_jwt_user'): return request._jwt_user auth_header = request.headers.get("Authorization", "") if auth_header.startswith("Bearer "): token = auth_header[7:] payload = verify_access_token(token) if payload: return { "id": int(payload["sub"]), "username": payload["username"], "vip_status": payload["vip"] } if "user_id" in session: return { "id": session["user_id"], "username": session["username"], "vip_status": session.get("vip_status", 0) } return None