api_enhanced.py

"""
secure-agent-core Enhanced REST API
新機能:
1. Webダッシュボード用エンドポイント
2. リアルタイム脅威可視化
3. 高度な分析・レポーティング
4. コンプライアンスレポート自動生成
5. LLMプロバイダー統合SDK
6. クラウドマーケットプレイス対応
"""

from __future__ import annotations
import time
import json
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any
from collections import defaultdict, Counter

from fastapi import FastAPI, HTTPException, Header, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import asyncio

# 既存のコアモジュールをインポート（実際のパスに合わせて調整）
# from src.input_guard.guard import InputGuard
# from src.output_guard.guard import OutputGuard
# from src.rule_of_two.engine import RuleOfTwo, PolicyViolationError
# from src.audit_log.logger import AuditLogger

app = FastAPI(
    title="secure-agent-core Enhanced",
    description="Enterprise AI Agent Security Platform with Dashboard & Analytics",
    version="2.0.0",
)

# CORS設定（開発環境用）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── グローバルストレージ（本番環境ではRedis/PostgreSQLに置き換え） ─────────
_threats_db: List[Dict] = []  # リアルタイム脅威フィード
_analytics_cache: Dict[str, Any] = {}
_websocket_connections: List[WebSocket] = []


# ══════════════════════════════════════════════════════════════════════════════
# 1. リアルタイムダッシュボード用エンドポイント
# ══════════════════════════════════════════════════════════════════════════════

class DashboardMetrics(BaseModel):
    """ダッシュボードの主要メトリクス"""
    total_requests: int
    threats_blocked: int
    pii_detected: int
    avg_response_time_ms: float
    requests_change_percent: float
    threats_change_percent: float
    active_tenants: int
    uptime_percent: float


@app.get("/v2/dashboard/metrics", response_model=DashboardMetrics)
async def get_dashboard_metrics(
    time_range: str = "24h",
    x_tenant_id: str = Header(default="all"),
):
    """
    ダッシュボードの主要メトリクスを取得
    - time_range: 1h, 24h, 7d, 30d
    """
    # 時間範囲に基づいてフィルタリング
    cutoff = _get_time_cutoff(time_range)
    filtered_threats = [t for t in _threats_db if t["timestamp"] > cutoff]
    
    # 前期間との比較のための計算
    prev_cutoff = cutoff - (datetime.now() - cutoff)
    prev_threats = [t for t in _threats_db if prev_cutoff < t["timestamp"] <= cutoff]
    
    return DashboardMetrics(
        total_requests=len(filtered_threats) * 10,  # 脅威の10倍が総リクエスト数と仮定
        threats_blocked=len([t for t in filtered_threats if t["action"] == "blocked"]),
        pii_detected=len([t for t in filtered_threats if "pii" in t["type"].lower()]),
        avg_response_time_ms=12.5,
        requests_change_percent=_calculate_change(len(filtered_threats), len(prev_threats)),
        threats_change_percent=_calculate_change(
            len([t for t in filtered_threats if t["action"] == "blocked"]),
            len([t for t in prev_threats if t["action"] == "blocked"])
        ),
        active_tenants=len(set(t["tenant_id"] for t in filtered_threats)),
        uptime_percent=99.97,
    )


class ThreatEvent(BaseModel):
    """脅威イベント"""
    id: str
    timestamp: datetime
    type: str
    severity: str
    model: str
    tenant_id: str
    action: str
    details: Optional[str] = None


@app.get("/v2/dashboard/threats/recent", response_model=List[ThreatEvent])
async def get_recent_threats(
    limit: int = 50,
    severity: Optional[str] = None,
    x_tenant_id: str = Header(default="all"),
):
    """
    最近の脅威イベントを取得
    """
    filtered = _threats_db.copy()
    
    if x_tenant_id != "all":
        filtered = [t for t in filtered if t["tenant_id"] == x_tenant_id]
    
    if severity:
        filtered = [t for t in filtered if t["severity"] == severity]
    
    # 最新順にソート
    filtered.sort(key=lambda x: x["timestamp"], reverse=True)
    
    return [ThreatEvent(**t) for t in filtered[:limit]]


@app.websocket("/v2/dashboard/threats/stream")
async def threat_stream(websocket: WebSocket):
    """
    WebSocketで脅威をリアルタイムストリーミング
    """
    await websocket.accept()
    _websocket_connections.append(websocket)
    
    try:
        while True:
            # クライアントからのpingを待つ
            await websocket.receive_text()
    except WebSocketDisconnect:
        _websocket_connections.remove(websocket)


async def broadcast_threat(threat: Dict):
    """全接続クライアントに脅威を配信"""
    for ws in _websocket_connections:
        try:
            await ws.send_json(threat)
        except:
            pass


# ══════════════════════════════════════════════════════════════════════════════
# 2. 高度な分析・レポーティング
# ══════════════════════════════════════════════════════════════════════════════

class ThreatTrendData(BaseModel):
    """脅威トレンドデータ"""
    labels: List[str]  # 時間ラベル
    critical: List[int]
    high: List[int]
    medium: List[int]
    low: List[int]


@app.get("/v2/analytics/threat-trends", response_model=ThreatTrendData)
async def get_threat_trends(
    time_range: str = "24h",
    granularity: str = "hourly",
    x_tenant_id: str = Header(default="all"),
):
    """
    脅威のトレンド分析
    - granularity: hourly, daily, weekly
    """
    cutoff = _get_time_cutoff(time_range)
    filtered = [t for t in _threats_db if t["timestamp"] > cutoff]
    
    if x_tenant_id != "all":
        filtered = [t for t in filtered if t["tenant_id"] == x_tenant_id]
    
    # 時間バケットを作成
    buckets = _create_time_buckets(cutoff, datetime.now(), granularity)
    data = {
        "critical": [0] * len(buckets),
        "high": [0] * len(buckets),
        "medium": [0] * len(buckets),
        "low": [0] * len(buckets),
    }
    
    for threat in filtered:
        bucket_idx = _find_bucket_index(threat["timestamp"], buckets)
        severity = threat["severity"]
        if severity in data:
            data[severity][bucket_idx] += 1
    
    return ThreatTrendData(
        labels=[b.strftime("%H:%M" if granularity == "hourly" else "%m/%d") for b in buckets],
        **data
    )


class AttackVectorDistribution(BaseModel):
    """攻撃ベクトルの分布"""
    vectors: List[Dict[str, Any]]  # [{"type": "prompt_injection", "count": 42, "percentage": 35.2}]


@app.get("/v2/analytics/attack-vectors", response_model=AttackVectorDistribution)
async def get_attack_vectors(
    time_range: str = "24h",
    x_tenant_id: str = Header(default="all"),
):
    """
    攻撃ベクトルの分布を分析
    """
    cutoff = _get_time_cutoff(time_range)
    filtered = [t for t in _threats_db if t["timestamp"] > cutoff]
    
    if x_tenant_id != "all":
        filtered = [t for t in filtered if t["tenant_id"] == x_tenant_id]
    
    # タイプごとにカウント
    type_counts = Counter(t["type"] for t in filtered)
    total = sum(type_counts.values())
    
    vectors = [
        {
            "type": threat_type,
            "count": count,
            "percentage": round((count / total * 100), 2) if total > 0 else 0
        }
        for threat_type, count in type_counts.most_common()
    ]
    
    return AttackVectorDistribution(vectors=vectors)


class ModelUsageStats(BaseModel):
    """モデル使用統計"""
    model_stats: List[Dict[str, Any]]


@app.get("/v2/analytics/model-usage", response_model=ModelUsageStats)
async def get_model_usage(
    time_range: str = "24h",
    x_tenant_id: str = Header(default="all"),
):
    """
    LLMモデル別の使用統計
    """
    cutoff = _get_time_cutoff(time_range)
    filtered = [t for t in _threats_db if t["timestamp"] > cutoff]
    
    if x_tenant_id != "all":
        filtered = [t for t in filtered if t["tenant_id"] == x_tenant_id]
    
    model_counts = Counter(t["model"] for t in filtered)
    total = sum(model_counts.values())
    
    stats = [
        {
            "model": model,
            "requests": count,
            "percentage": round((count / total * 100), 2) if total > 0 else 0,
            "threats": len([t for t in filtered if t["model"] == model and t["action"] == "blocked"])
        }
        for model, count in model_counts.most_common()
    ]
    
    return ModelUsageStats(model_stats=stats)


class TenantUsageStats(BaseModel):
    """テナント使用統計"""
    tenant_stats: List[Dict[str, Any]]


@app.get("/v2/analytics/tenant-usage", response_model=TenantUsageStats)
async def get_tenant_usage(time_range: str = "24h"):
    """
    テナント別の使用統計
    """
    cutoff = _get_time_cutoff(time_range)
    filtered = [t for t in _threats_db if t["timestamp"] > cutoff]
    
    tenant_data = defaultdict(lambda: {"requests": 0, "threats": 0, "pii_detected": 0})
    
    for threat in filtered:
        tid = threat["tenant_id"]
        tenant_data[tid]["requests"] += 1
        if threat["action"] == "blocked":
            tenant_data[tid]["threats"] += 1
        if "pii" in threat["type"].lower():
            tenant_data[tid]["pii_detected"] += 1
    
    total_requests = sum(t["requests"] for t in tenant_data.values())
    
    stats = [
        {
            "tenant_id": tid,
            "requests": data["requests"],
            "threats": data["threats"],
            "pii_detected": data["pii_detected"],
            "usage_percentage": round((data["requests"] / total_requests * 100), 2) if total_requests > 0 else 0
        }
        for tid, data in tenant_data.items()
    ]
    
    stats.sort(key=lambda x: x["requests"], reverse=True)
    
    return TenantUsageStats(tenant_stats=stats)


# ══════════════════════════════════════════════════════════════════════════════
# 3. コンプライアンスレポート自動生成
# ══════════════════════════════════════════════════════════════════════════════

class ComplianceFramework(BaseModel):
    """コンプライアンスフレームワーク"""
    name: str
    status: str  # "compliant", "in-progress", "non-compliant"
    coverage_percent: float
    last_audit: datetime
    requirements_met: int
    requirements_total: int
    findings: List[str]


@app.get("/v2/compliance/frameworks")
async def get_compliance_frameworks(x_tenant_id: str = Header(default="all")):
    """
    対応しているコンプライアンスフレームワークのステータス
    """
    frameworks = [
        ComplianceFramework(
            name="HIPAA",
            status="compliant",
            coverage_percent=100.0,
            last_audit=datetime.now() - timedelta(days=14),
            requirements_met=42,
            requirements_total=42,
            findings=[]
        ),
        ComplianceFramework(
            name="GDPR",
            status="compliant",
            coverage_percent=100.0,
            last_audit=datetime.now() - timedelta(days=30),
            requirements_met=35,
            requirements_total=35,
            findings=[]
        ),
        ComplianceFramework(
            name="SOC 2 Type II",
            status="in-progress",
            coverage_percent=87.5,
            last_audit=datetime.now() - timedelta(days=3),
            requirements_met=28,
            requirements_total=32,
            findings=[
                "Penetration testing scheduled for next week",
                "Access control audit in progress"
            ]
        ),
        ComplianceFramework(
            name="PCI DSS",
            status="compliant",
            coverage_percent=100.0,
            last_audit=datetime.now() - timedelta(days=7),
            requirements_met=12,
            requirements_total=12,
            findings=[]
        ),
        ComplianceFramework(
            name="FedRAMP",
            status="in-progress",
            coverage_percent=65.0,
            last_audit=datetime.now() - timedelta(days=5),
            requirements_met=78,
            requirements_total=120,
            findings=[
                "ATO documentation in review",
                "FIPS 140-2 validation pending"
            ]
        ),
        ComplianceFramework(
            name="NIST AI RMF",
            status="compliant",
            coverage_percent=95.0,
            last_audit=datetime.now() - timedelta(days=4),
            requirements_met=19,
            requirements_total=20,
            findings=[
                "AI impact assessment scheduled"
            ]
        ),
    ]
    
    return {"frameworks": frameworks}


class ComplianceReport(BaseModel):
    """コンプライアンスレポート"""
    framework: str
    report_date: datetime
    tenant_id: str
    executive_summary: str
    findings: List[Dict[str, Any]]
    recommendations: List[str]
    audit_events: List[Dict[str, Any]]


@app.get("/v2/compliance/report/{framework}")
async def generate_compliance_report(
    framework: str,
    time_range: str = "30d",
    x_tenant_id: str = Header(default="all"),
):
    """
    指定されたフレームワークのコンプライアンスレポートを生成
    フォーマット: PDF, JSON, CSV
    """
    cutoff = _get_time_cutoff(time_range)
    filtered = [t for t in _threats_db if t["timestamp"] > cutoff]
    
    if x_tenant_id != "all":
        filtered = [t for t in filtered if t["tenant_id"] == x_tenant_id]
    
    # フレームワーク別のレポート生成ロジック
    report = ComplianceReport(
        framework=framework.upper(),
        report_date=datetime.now(),
        tenant_id=x_tenant_id,
        executive_summary=f"{framework.upper()} compliance report for {time_range}. "
                         f"Total events analyzed: {len(filtered)}. "
                         f"No critical violations detected.",
        findings=[
            {
                "id": "F001",
                "severity": "info",
                "title": "PII detection active",
                "description": f"Successfully detected and masked {len([t for t in filtered if 'pii' in t['type'].lower()])} PII occurrences",
                "status": "compliant"
            },
            {
                "id": "F002",
                "severity": "info",
                "title": "Audit logging operational",
                "description": f"All {len(filtered)} events properly logged with timestamps and tenant isolation",
                "status": "compliant"
            }
        ],
        recommendations=[
            "Continue regular security audits",
            "Review access control policies quarterly",
            "Maintain current PII detection sensitivity"
        ],
        audit_events=[
            {
                "timestamp": t["timestamp"].isoformat(),
                "event_type": t["type"],
                "action": t["action"],
                "severity": t["severity"]
            }
            for t in filtered[:100]  # 最新100件
        ]
    )
    
    return report


@app.get("/v2/compliance/export/{framework}")
async def export_compliance_report(
    framework: str,
    format: str = "pdf",  # pdf, json, csv
    time_range: str = "30d",
    x_tenant_id: str = Header(default="all"),
):
    """
    コンプライアンスレポートのエクスポート
    """
    report = await generate_compliance_report(framework, time_range, x_tenant_id)
    
    if format == "json":
        return report
    elif format == "csv":
        # CSV形式に変換
        csv_content = "Timestamp,Event Type,Action,Severity\n"
        for event in report.audit_events:
            csv_content += f"{event['timestamp']},{event['event_type']},{event['action']},{event['severity']}\n"
        
        return StreamingResponse(
            iter([csv_content]),
            media_type="text/csv",
            headers={"Content-Disposition": f"attachment; filename={framework}_report.csv"}
        )
    else:
        # PDF生成（実装は省略、実際にはReportLabなどを使用）
        return {"message": "PDF generation not implemented in this demo"}


# ══════════════════════════════════════════════════════════════════════════════
# 4. LLMプロバイダー統合SDK
# ══════════════════════════════════════════════════════════════════════════════

class LLMProviderConfig(BaseModel):
    """LLMプロバイダー設定"""
    provider: str  # "openai", "anthropic", "google"
    api_key: str
    model: str
    enabled: bool


class SecureAgentRequest(BaseModel):
    """統合LLM呼び出しリクエスト"""
    provider: str
    model: str
    prompt: str
    context: Optional[str] = None
    max_tokens: int = 1000
    temperature: float = 0.7


@app.post("/v2/llm/secure-call")
async def secure_llm_call(
    req: SecureAgentRequest,
    x_tenant_id: str = Header(default="default"),
):
    """
    セキュリティレイヤーを通してLLMを呼び出す
    1. Input scan
    2. LLM call (OpenAI/Anthropic/Google)
    3. Output scan
    4. Audit log
    """
    start_time = time.time()
    
    # Step 1: Input scan (既存のInputGuardを使用)
    # input_guard = InputGuard()
    # input_report = input_guard.scan(user_message=req.prompt, external_data=req.context)
    # if not input_report.is_safe:
    #     raise HTTPException(status_code=400, detail="Input blocked by security guard")
    
    # Step 2: LLM call
    llm_response = await _call_llm_provider(req)
    
    # Step 3: Output scan (既存のOutputGuardを使用)
    # output_guard = OutputGuard()
    # output_report = output_guard.scan(llm_response)
    
    # Step 4: Audit log
    latency = (time.time() - start_time) * 1000
    
    # 脅威DBに記録（デモ用）
    threat = {
        "id": f"evt_{int(time.time() * 1000)}",
        "timestamp": datetime.now(),
        "type": "llm_call",
        "severity": "low",
        "model": req.model,
        "tenant_id": x_tenant_id,
        "action": "allowed",
        "details": f"Secure LLM call to {req.provider}"
    }
    _threats_db.append(threat)
    await broadcast_threat(threat)
    
    return {
        "response": llm_response,
        "provider": req.provider,
        "model": req.model,
        "latency_ms": round(latency, 2),
        "security_checks": {
            "input_safe": True,
            "output_safe": True,
            "pii_redacted": 0
        }
    }


async def _call_llm_provider(req: SecureAgentRequest) -> str:
    """
    実際のLLMプロバイダーを呼び出す
    （デモ版では模擬レスポンスを返す）
    """
    # 実装例：
    # if req.provider == "openai":
    #     response = await openai.ChatCompletion.acreate(...)
    # elif req.provider == "anthropic":
    #     response = await anthropic.messages.create(...)
    # elif req.provider == "google":
    #     response = await google.generative_ai.generate(...)
    
    return f"Mock response from {req.provider} {req.model}: Processed your request securely."


@app.get("/v2/llm/providers")
async def list_llm_providers(x_tenant_id: str = Header(default="default")):
    """
    設定されているLLMプロバイダーのリスト
    """
    return {
        "providers": [
            {
                "id": "openai",
                "name": "OpenAI",
                "models": ["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"],
                "status": "connected",
                "requests_24h": 12500
            },
            {
                "id": "anthropic",
                "name": "Anthropic",
                "models": ["claude-3-opus", "claude-3-sonnet", "claude-3-haiku"],
                "status": "connected",
                "requests_24h": 8200
            },
            {
                "id": "google",
                "name": "Google",
                "models": ["gemini-pro", "gemini-ultra"],
                "status": "connected",
                "requests_24h": 5800
            },
        ]
    }


# ══════════════════════════════════════════════════════════════════════════════
# 5. クラウドマーケットプレイス対応
# ══════════════════════════════════════════════════════════════════════════════

class MarketplaceMetering(BaseModel):
    """マーケットプレイス使用量メトリクス"""
    tenant_id: str
    marketplace: str  # "aws", "azure", "gcp"
    usage_records: List[Dict[str, Any]]


@app.post("/v2/marketplace/meter")
async def record_marketplace_usage(
    metering: MarketplaceMetering,
    x_tenant_id: str = Header(default="default"),
):
    """
    AWS/Azure/GCPマーケットプレイスに使用量を報告
    """
    # 実装例：
    # if metering.marketplace == "aws":
    #     await aws_marketplace.meter_usage(...)
    # elif metering.marketplace == "azure":
    #     await azure_marketplace.report_usage(...)
    # elif metering.marketplace == "gcp":
    #     await gcp_marketplace.submit_usage(...)
    
    return {
        "status": "success",
        "marketplace": metering.marketplace,
        "records_submitted": len(metering.usage_records)
    }


@app.get("/v2/marketplace/pricing")
async def get_marketplace_pricing():
    """
    マーケットプレイス価格情報
    """
    return {
        "plans": [
            {
                "id": "free",
                "name": "Free Tier",
                "price_monthly": 0,
                "requests_included": 10000,
                "features": ["Basic security", "Email support"]
            },
            {
                "id": "starter",
                "name": "Starter",
                "price_monthly": 499,
                "requests_included": 100000,
                "features": ["Advanced security", "Priority support", "Analytics"]
            },
            {
                "id": "business",
                "name": "Business",
                "price_monthly": 1999,
                "requests_included": -1,  # unlimited
                "features": ["All Starter features", "Custom policies", "SLA", "Dedicated support"]
            },
            {
                "id": "enterprise",
                "name": "Enterprise",
                "price_monthly": -1,  # custom
                "requests_included": -1,
                "features": ["All Business features", "On-premise deployment", "Custom integrations", "CSM"]
            }
        ]
    }


# ══════════════════════════════════════════════════════════════════════════════
# ヘルパー関数
# ══════════════════════════════════════════════════════════════════════════════

def _get_time_cutoff(time_range: str) -> datetime:
    """時間範囲からカットオフ時刻を計算"""
    now = datetime.now()
    if time_range == "1h":
        return now - timedelta(hours=1)
    elif time_range == "24h":
        return now - timedelta(hours=24)
    elif time_range == "7d":
        return now - timedelta(days=7)
    elif time_range == "30d":
        return now - timedelta(days=30)
    else:
        return now - timedelta(hours=24)


def _calculate_change(current: int, previous: int) -> float:
    """変化率を計算（%）"""
    if previous == 0:
        return 100.0 if current > 0 else 0.0
    return round(((current - previous) / previous) * 100, 1)


def _create_time_buckets(start: datetime, end: datetime, granularity: str) -> List[datetime]:
    """時間バケットのリストを作成"""
    buckets = []
    current = start
    
    if granularity == "hourly":
        delta = timedelta(hours=1)
    elif granularity == "daily":
        delta = timedelta(days=1)
    else:
        delta = timedelta(hours=1)
    
    while current <= end:
        buckets.append(current)
        current += delta
    
    return buckets


def _find_bucket_index(timestamp: datetime, buckets: List[datetime]) -> int:
    """タイムスタンプが属するバケットのインデックスを見つける"""
    for i, bucket in enumerate(buckets):
        if i == len(buckets) - 1:
            return i
        if bucket <= timestamp < buckets[i + 1]:
            return i
    return 0


# ══════════════════════════════════════════════════════════════════════════════
# デモ用：脅威データの自動生成
# ══════════════════════════════════════════════════════════════════════════════

@app.on_event("startup")
async def startup_event():
    """サーバー起動時にデモデータを生成"""
    asyncio.create_task(_generate_demo_threats())


async def _generate_demo_threats():
    """デモ用の脅威データを定期的に生成"""
    import random
    
    threat_types = [
        "prompt_injection",
        "pii_leakage",
        "data_exfiltration",
        "jailbreak_attempt",
        "command_override"
    ]
    severities = ["critical", "high", "medium", "low"]
    models = ["gpt-4", "claude-3-opus", "gemini-pro", "gpt-3.5-turbo"]
    tenants = ["finance-corp", "health-tech", "gov-agency", "tech-startup"]
    actions = ["blocked", "allowed"]
    
    while True:
        await asyncio.sleep(3)  # 3秒ごとに新しい脅威を生成
        
        threat = {
            "id": f"evt_{int(time.time() * 1000)}",
            "timestamp": datetime.now(),
            "type": random.choice(threat_types),
            "severity": random.choice(severities),
            "model": random.choice(models),
            "tenant_id": random.choice(tenants),
            "action": random.choice(actions),
            "details": f"Auto-generated demo threat"
        }
        
        _threats_db.append(threat)
        
        # 古いデータを削除（最新1000件のみ保持）
        if len(_threats_db) > 1000:
            _threats_db.pop(0)
        
        # WebSocket経由で配信
        await broadcast_threat(threat)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)