api.py

# api.py
# 此檔案定義所有 API 端點
# 包含統計數據、日誌查詢、圖表數據等

from fastapi import FastAPI, Query, Body
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import Optional, List
import os
from groq import Groq

from auth import (
    simulate_click, USER_DB, check_session_timeouts, get_user_status,
    internal_process_login, internal_logout
)
from security import test_send_security_alert, verify_identity
from config import SYSTEM_WEIGHTS, ACTION_MAP, FIXED_TEST_EMAIL
from plotting import get_plot_data_for_api
import pandas as pd
import models  # 使用 models.XXX 存取即時狀態，避免布林值被複製後失效
from utils import get_geo_level, update_system_weights

api_app = FastAPI(title="數位保鏢 API")

# --- Pydantic Models ---
class AnalyzeLogRequest(BaseModel):
    log_data: str  # JSON 格式的 log 字串
    event_desc: Optional[str] = ""
    event_ip: Optional[str] = ""
    event_status: Optional[str] = ""

class WeightsRequest(BaseModel):
    ow: float
    tw: float
    iw: float
    dw: float
    gw: float
    fw: float
    
class LoginRequest(BaseModel):
    account: str
    password: str
    captcha: str
    ip: str
    cookie: str
    time_mode: str
    custom_time: Optional[str] = ""

class StudentActionRequest(BaseModel):
    ip: str
    cookie: str
    time_mode: str
    custom_time: Optional[str] = ""

class StudentActionExtendedRequest(StudentActionRequest):
    account: str

class GenericActionRequest(StudentActionRequest):
    account: str
    action_code: str

# --- 純淨 API 端點專用函式 (供內部呼叫) ---
def internal_get_stats():
    """取得系統狀態與權重"""
    total = models.global_stats["total_count"]
    abnormal = models.global_stats["abnormal_count"]
    rate = round((abnormal / total * 100), 2) if total > 0 else 0
    return {
        "status": "success",
        "ai_ready": models.AI_READY,
        "data": {
            "total_logs": total,
            "abnormal_logs": abnormal,
            "anomaly_rate_percent": rate,
            "current_weights": SYSTEM_WEIGHTS
        }
    }

def internal_get_logs(limit=20):
    """取得最新的 N 筆 Log"""
    return {
        "status": "success",
        "count": len(models.log_history[:limit]),
        "data": models.log_history[:limit]
    }

def _extract_gr_update_value(upd):
    """從 Gradio update 物件中提取 value (如果有)"""
    if hasattr(upd, 'value'):
        return upd.value
    if isinstance(upd, dict) and 'value' in upd:
        return upd['value']
    return ""

# --- FastAPI 路由定義 ---

@api_app.get("/api/stats")
def get_stats():
    return internal_get_stats()

@api_app.get("/api/logs")
def get_logs(limit: int = 20):
    return internal_get_logs(limit)

@api_app.post("/api/login")
def login_endpoint(req: LoginRequest):
    return internal_process_login(
        req.ip, req.cookie, req.account, req.password, 
        req.captcha, req.time_mode, req.custom_time
    )

@api_app.post("/api/logout")
def logout_endpoint(req: StudentActionExtendedRequest):
    return internal_logout(
        req.ip, req.cookie, req.account, req.time_mode, req.custom_time
    )

@api_app.get("/api/chart_data")
def get_chart_data(action_name: str = Query(...)):
    """取得特定系統的圖表原始數據"""
    if not models.DATA_READY:
        return {"status": "error", "message": "歷史基準資料未載入，請確認 baseline_logs.csv 是否存在或雲端同步是否正常"}

    # 嘗試從顯示名稱映射到代碼，如果找不到，假設輸入已經是代碼
    action_code = ACTION_MAP.get(action_name) or action_name
    
    # 驗證 action_code 是否有效 (在 baseline_df 中存在或在 ACTION_MAP 的值中)
    valid_codes = set(ACTION_MAP.values())
    if action_code not in valid_codes and action_code not in models.baseline_df['action'].unique():
        return {"status": "error", "message": f"找不到系統代碼或名稱: {action_name}"}

    df_hist = models.baseline_df[models.baseline_df['action'] == action_code].copy()
    df_live = pd.DataFrame(models.log_history)

    if not df_live.empty:
        df_live['timestamp'] = pd.to_datetime(df_live['timestamp'], errors='coerce')
        df_live = df_live.dropna(subset=['timestamp'])
        df_live = df_live[df_live['action'] == action_code].copy()

    if not df_live.empty:
        df_combined = pd.concat([df_hist, df_live], ignore_index=True)
        df_combined = df_combined.drop_duplicates(subset=['timestamp', 'account', 'action'])
    else:
        df_combined = df_hist.copy()

    if df_combined.empty:
        return {"status": "success", "action": action_code, "data": {"normal": [], "abnormal": []}}

    df_combined['month'] = df_combined['timestamp'].dt.month
    df_combined['time_of_day'] = df_combined['timestamp'].dt.hour + df_combined['timestamp'].dt.minute / 60.0

    is_abnormal = df_combined['status'].astype(str).str.contains("異常", na=False)

    normal_data = df_combined[~is_abnormal][['month', 'time_of_day', 'timestamp']].astype(str).to_dict('records')
    abnormal_data = df_combined[is_abnormal][['month', 'time_of_day', 'timestamp', 'status']].astype(str).to_dict('records')

    return {
        "status": "success",
        "action_code": action_code,
        "data": {"normal": normal_data, "abnormal": abnormal_data}
    }

@api_app.get("/api/abnormal_logs")
def get_abnormal_logs():
    """取得所有異常 Log"""
    abnormal_logs = [log for log in models.log_history if "異常" in str(log.get("status", "")) or "錯誤" in str(log.get("status", "")) or "受限" in str(log.get("status", ""))]
    return {
        "status": "success",
        "count": len(abnormal_logs),
        "data": abnormal_logs
    }

@api_app.get("/api/get_user_logs")
def get_user_logs(session_id: str = Query(...)):
    """
    獲取指定 session_id (cookie_id) 的所有日誌
    過濾自 Hugging Face 同步的歷史數據與當前記憶體數據
    """
    # 1. 準備過濾後的數據
    filtered_logs = []
    
    # 從歷史數據中過濾 (如果有的話)
    if models.baseline_df is not None and not models.baseline_df.empty:
        # 確保資料格式一致
        df_filtered = models.baseline_df[models.baseline_df['cookie_id'] == session_id].copy()
        if not df_filtered.empty:
            df_filtered['timestamp'] = pd.to_datetime(df_filtered['timestamp']).dt.strftime('%Y-%m-%d %H:%M:%S')
            filtered_logs.extend(df_filtered.to_dict('records'))
            
    # 從當前記憶體日誌中過濾
    mem_logs = [log for log in models.log_history if log.get("cookie_id") == session_id]
    filtered_logs.extend(mem_logs)
    
    # 2. 去重並排序 (依時間倒序)
    # 使用 timestamp + action 作為簡單的去重基準
    seen = set()
    unique_logs = []
    for log in filtered_logs:
        key = (log.get('timestamp'), log.get('action'))
        if key not in seen:
            seen.add(key)
            unique_logs.append(log)
            
    unique_logs.sort(key=lambda x: x.get('timestamp', ''), reverse=True)
    
    # 3. 豐富化資料 (加上地理位置、設備資訊、判定狀態)
    final_logs = []
    for log in unique_logs:
        ip = log.get('ip_address', 'Unknown')
        geo_code = get_geo_level(ip)
        
        # 地理位置映射
        geo_map = {
            0: "台北市 (校內專網)",
            1: "台北市 (校園網路)",
            2: "台灣 (宿舍/寬頻)",
            3: "台灣 (公共網路)",
            4: "美國 (海外連線)",
            5: "未知區域"
        }
        location = geo_map.get(geo_code, "未知區域")
        
        # AI 判定狀態轉為中文
        status_raw = str(log.get('status', 'success'))
        ai_status = "異常" if ("異常" in status_raw or "受限" in status_raw or "錯誤" in status_raw) else "正常"
        
        final_logs.append({
            "time": log.get('timestamp'),
            "action": log.get('action'),
            "ip": ip,
            "location": location,
            "ai_status": ai_status,
            "status_detail": status_raw # 保留原始詳細資訊
        })
        
    return {
        "status": "success",
        "session_id": session_id,
        "count": len(final_logs),
        "data": final_logs
    }

@api_app.get("/api/chart_html")
def get_chart_html(action_name: str = Query(...)):
    """取得圖表的 HTML 內容 (Gradio 專用)"""
    from plotting import get_chart_js_html
    html_data = get_chart_js_html(action_name)
    return {"status": "success", "html": html_data}

@api_app.get("/api/plot_image")
def get_plot_image(action_name: str = Query(...)):
    """取得圖表的 Base64 圖片"""
    if not models.DATA_READY:
        return {"status": "error", "message": "歷史基準資料未載入"}
    img_data = get_plot_data_for_api(action_name)
    return {"status": "success", "image": img_data}

@api_app.post("/api/update_weights")
def update_weights(req: WeightsRequest):
    """更新系統權重"""
    msg = update_system_weights(req.ow, req.tw, req.iw, req.dw, req.gw, req.fw)
    return {"status": "success", "message": msg}

@api_app.get("/verify_identity")
def verify_identity_endpoint(token: str = Query(...), choice: str = Query(...)):
    return verify_identity(token, choice)

@api_app.get("/test_send_security_alert")
def test_send_security_alert_endpoint():
    return test_send_security_alert()

# --- Student Actions ---

def _run_student_action_helper(ip, c, acc, action_code, tm, ct):
    # 每次請求時檢查超時
    check_session_timeouts(timeout_seconds=300)
    
    # 直接以帳戶名稱作為 key 查詢封鎖狀態
    current_status = get_user_status(acc)
    
    if current_status == "verifying":
        return {
            "status": "pending_verify",
            "message": "偵測到異常，請於5分鐘內確認 Email 以繼續使用"
        }
    elif current_status == "blocked":
        return {
            "status": "blocked",
            "message": "帳號已被封鎖，請聯絡系統管理員或寄驗證信進行解封"
        }

    df_all, df_abn, inp_upd = simulate_click(ip, c, acc, action_code, tm, ct)
    val = _extract_gr_update_value(inp_upd)
    return {
        "status": "success",
        "anomaly_msg": val or "",
        "action": action_code
    }

@api_app.post("/api/student/eval")
def student_eval(req: StudentActionExtendedRequest):
    return _run_student_action_helper(req.ip, req.cookie, req.account, "click_eval_system", req.time_mode, req.custom_time)

@api_app.post("/api/student/pre_select")
def student_pre_select(req: StudentActionExtendedRequest):
    return _run_student_action_helper(req.ip, req.cookie, req.account, "click_pre_select_system", req.time_mode, req.custom_time)

@api_app.post("/api/student/leave")
def student_leave(req: StudentActionExtendedRequest):
    return _run_student_action_helper(req.ip, req.cookie, req.account, "click_leave_system", req.time_mode, req.custom_time)

@api_app.post("/api/student/dorm")
def student_dorm(req: StudentActionExtendedRequest):
    return _run_student_action_helper(req.ip, req.cookie, req.account, "click_dorm_system", req.time_mode, req.custom_time)

@api_app.post("/api/student/webmail")
def student_webmail(req: StudentActionExtendedRequest):
    return _run_student_action_helper(req.ip, req.cookie, req.account, "click_webmail", req.time_mode, req.custom_time)

@api_app.post("/api/student/vdesk")
def student_vdesk(req: StudentActionExtendedRequest):
    return _run_student_action_helper(req.ip, req.cookie, req.account, "click_vdesk", req.time_mode, req.custom_time)

@api_app.post("/api/student/sql_injection")
def student_sql_injection(req: StudentActionExtendedRequest):
    return _run_student_action_helper(req.ip, req.cookie, req.account, "malicious_sql_injection", req.time_mode, req.custom_time)

@api_app.post("/api/student/run_action")
def student_run_action(req: GenericActionRequest):
    """通用學生操作端點"""
    # 每次請求時檢查超時
    check_session_timeouts(timeout_seconds=300)
    
    # 直接以帳戶名稱作為 key 查詢封鎖狀態
    current_status = get_user_status(req.account)
    
    if current_status == "verifying":
        return {
            "status": "pending_verify",
            "message": "偵測到異常，請於5分鐘內確認 Email 以繼續使用"
        }
    elif current_status == "blocked":
        return {
            "status": "blocked",
            "message": "帳號已被封鎖，請聯絡系統管理員或寄驗證信進行解封"
        }

    df_all, df_abn, inp_upd = simulate_click(req.ip, req.cookie, req.account, req.action_code, req.time_mode, req.custom_time)
    val = _extract_gr_update_value(inp_upd)
    return {
        "status": "success",
        "anomaly_msg": val or "",
        # 為了 Vue 前端可能需要的資料更新，回傳簡化後的 log
        "log_entry": {
            "timestamp": pd.Timestamp.now().strftime("%Y-%m-%d %H:%M:%S"), # 這裡簡化處理
            "action": req.action_code,
            "status": val or "success"
        }
    }

# --- 使用者管理 API 端點 ---
class AccountRequest(BaseModel):
    account: str

@api_app.get("/api/users")
def get_users_endpoint():
    """取得所有使用者帳戶清單與其狀態"""
    users_list = []
    for acc, info in USER_DB.items():
        role_en = "admin" if info[1] == "管理" else "student"
        status = get_user_status(acc)
        users_list.append({
            "account": acc,
            "role": role_en,
            "role_zh": info[1],
            "name": info[2],
            "status": status
        })
    return {"status": "success", "data": users_list}

@api_app.post("/api/users/block")
def block_user_endpoint(req: AccountRequest):
    """封鎖特定使用者"""
    acc = req.account.strip()
    if acc not in USER_DB:
        return {"status": "error", "message": "帳號不存在"}
    
    from auth import session_cache_lock, session_cache
    with session_cache_lock:
        if acc not in session_cache:
            session_cache[acc] = {}
        session_cache[acc]["status"] = "blocked"
        session_cache[acc]["pending_token"] = None
        session_cache[acc]["anomaly_timestamp"] = None
    
    return {"status": "success", "message": f"帳號 {acc} 已成功封鎖"}

@api_app.post("/api/users/unblock")
def unblock_user_endpoint(req: AccountRequest):
    """解鎖特定使用者"""
    acc = req.account.strip()
    if acc not in USER_DB:
        return {"status": "error", "message": "帳號不存在"}
    
    from auth import unblock_user
    unblock_user(acc)
    
    return {"status": "success", "message": f"帳號 {acc} 已成功解封"}


# --- Groq AI 分析端點 ---

@api_app.post("/api/analyze_log")
def analyze_log_with_groq(req: AnalyzeLogRequest):
    """
    使用 Groq AI 分析異常 Log 事件。
    接收原始 Log JSON 字串，呼叫 Groq API 進行資安分析，回傳中文解析結果。
    """
    groq_api_key = os.getenv("GROQ_API_KEY", "")
    if not groq_api_key:
        return {
            "status": "error",
            "message": "Groq API Key 未設定，請確認後端 .env 檔案中的 GROQ_API_KEY"
        }

    try:
        client = Groq(api_key=groq_api_key)

        # 建構 Prompt
        prompt_parts = []
        prompt_parts.append("你是一位資訊安全專家，正在分析一筆校園系統的異常存取 Log。")
        prompt_parts.append("")
        
        if req.event_ip:
            prompt_parts.append(f"來源 IP：{req.event_ip}")
        if req.event_status:
            prompt_parts.append(f"AI 判定狀態：{req.event_status}")
        if req.event_desc:
            prompt_parts.append(f"事件描述：{req.event_desc}")
        
        prompt_parts.append("")
        prompt_parts.append("原始 Log 數據（JSON 格式）：")
        prompt_parts.append(req.log_data)
        prompt_parts.append("")
        prompt_parts.append("請根據以上資訊，以繁體中文提供：")
        prompt_parts.append("1. 可能的攻擊類型或異常行為說明（1-2句）")
        prompt_parts.append("2. 風險等級評估（低/中/高）並說明原因")
        prompt_parts.append("3. 具體的建議處理動作（2-4點條列）")
        prompt_parts.append("")
        prompt_parts.append("請直接輸出分析內容，格式簡潔專業，不要加入多餘的標題或引號。")
        
        full_prompt = "\n".join(prompt_parts)

        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "你是一位熟悉台灣校園資訊安全的資安分析師，專精於識別異常存取行為、帳號盜用與惡意攻擊。請以繁體中文回答。"
                },
                {
                    "role": "user",
                    "content": full_prompt
                }
            ],
            model="llama-3.3-70b-versatile",
            temperature=0.4,
            max_tokens=512,
        )

        analysis_text = chat_completion.choices[0].message.content or "無法生成分析結果。"

        return {
            "status": "success",
            "analysis": analysis_text,
            "model": chat_completion.model,
            "usage": {
                "prompt_tokens": chat_completion.usage.prompt_tokens,
                "completion_tokens": chat_completion.usage.completion_tokens,
            }
        }

    except Exception as e:
        error_msg = str(e)
        print(f"[Groq API Error] {error_msg}")
        return {
            "status": "error",
            "message": f"Groq AI 分析失敗：{error_msg}"
        }