|
| 1 | +import json |
| 2 | +import os |
| 3 | +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout |
| 4 | +from typing import Any |
| 5 | +from google import genai |
| 6 | +from google.genai import types |
| 7 | +from core.context import PipelineContext, AIOutput, ErrorDetail |
| 8 | +from ai.contract import AIOutputContract |
| 9 | + |
| 10 | + |
| 11 | +def _call_gemini( |
| 12 | + api_key: str, |
| 13 | + model_name: str, |
| 14 | + prompt: str, |
| 15 | + timeout_ms: int, |
| 16 | +) -> dict[str, Any]: |
| 17 | + client = genai.Client( |
| 18 | + api_key=api_key, |
| 19 | + http_options=types.HttpOptions(timeout=timeout_ms), |
| 20 | + ) |
| 21 | + |
| 22 | + contents = [ |
| 23 | + types.Content( |
| 24 | + role="user", |
| 25 | + parts=[ |
| 26 | + types.Part.from_text(text=prompt), |
| 27 | + ], |
| 28 | + ), |
| 29 | + ] |
| 30 | + |
| 31 | + gen_config = types.GenerateContentConfig() |
| 32 | + |
| 33 | + response = client.models.generate_content( |
| 34 | + model=model_name, |
| 35 | + contents=contents, # type: ignore[arg-type] |
| 36 | + config=gen_config, |
| 37 | + ) |
| 38 | + |
| 39 | + text = "" |
| 40 | + if response.text: |
| 41 | + text = response.text.strip() |
| 42 | + |
| 43 | + if not text: |
| 44 | + raise ValueError("AI returned an empty response.") |
| 45 | + |
| 46 | + result: dict[str, Any] = json.loads(text) |
| 47 | + return result |
| 48 | + |
| 49 | + |
| 50 | +def analyze_with_gemini(ctx: PipelineContext) -> PipelineContext: |
| 51 | + api_key = os.getenv("AURCHECKER_AI_API_KEY") |
| 52 | + if not api_key: |
| 53 | + ctx.errors.append( |
| 54 | + ErrorDetail( |
| 55 | + code="AI_DISABLED", stage="AI", message="AURCHECKER_AI_API_KEY env var is not set.", recoverable=True |
| 56 | + ) |
| 57 | + ) |
| 58 | + return ctx |
| 59 | + |
| 60 | + if not ctx.features or not ctx.metadata: |
| 61 | + ctx.errors.append( |
| 62 | + ErrorDetail( |
| 63 | + code="AI_SCHEMA_INVALID", |
| 64 | + stage="AI", |
| 65 | + message="Cannot run AI analysis without features and metadata.", |
| 66 | + recoverable=True, |
| 67 | + ) |
| 68 | + ) |
| 69 | + return ctx |
| 70 | + |
| 71 | + payload = { |
| 72 | + "package": ctx.package, |
| 73 | + "features": ctx.features.dict(), |
| 74 | + "metadata": ctx.metadata.dict(), |
| 75 | + "risk_score": ctx.risk_score, |
| 76 | + "risk_level": ctx.risk_level, |
| 77 | + } |
| 78 | + payload_str = json.dumps(payload) |
| 79 | + |
| 80 | + model_name = os.getenv("AURCHECKER_AI_MODEL", "gemini-3.1-flash-lite") |
| 81 | + timeout_ms = int(os.getenv("AURCHECKER_AI_TIMEOUT", "120000")) |
| 82 | + deadline_s = timeout_ms / 1000 + 5 |
| 83 | + |
| 84 | + prompt = f"""You are a security expert analyzing an Arch AUR PKGBUILD. Output ONLY valid JSON matching this schema exactly: |
| 85 | +{{ |
| 86 | + "summary": "string (max 200 chars) - brief risk overview", |
| 87 | + "suspicious_patterns": ["string"] - list of specific dangerous patterns found (max 5)", |
| 88 | + "attack_surface": "string (max 300 chars) - what an attacker could exploit", |
| 89 | + "rationale": ["string"] - reasons for the risk score (max 3)", |
| 90 | + "ai_risk_score": integer 0-100 - your independent risk score based on PKGBUILD content", |
| 91 | + "ai_risk_level": "LOW" or "MEDIUM" or "HIGH", |
| 92 | + "ai_risk_details": "string (max 500 chars) - what specifically contributed to the risk score" |
| 93 | +}} |
| 94 | +
|
| 95 | +Input: {payload_str} |
| 96 | +
|
| 97 | +Output valid JSON only. No markdown. No prose.""" |
| 98 | + |
| 99 | + try: |
| 100 | + with ThreadPoolExecutor(max_workers=1) as pool: |
| 101 | + fut = pool.submit(_call_gemini, api_key, model_name, prompt, timeout_ms) |
| 102 | + result = fut.result(timeout=deadline_s) |
| 103 | + |
| 104 | + validated = AIOutputContract(**result) |
| 105 | + |
| 106 | + ctx.ai_analysis = AIOutput( |
| 107 | + summary=validated.summary, |
| 108 | + suspicious_patterns=validated.suspicious_patterns, |
| 109 | + attack_surface=validated.attack_surface, |
| 110 | + rationale=validated.rationale, |
| 111 | + ai_risk_score=validated.ai_risk_score, |
| 112 | + ai_risk_level=validated.ai_risk_level, |
| 113 | + ai_risk_details=validated.ai_risk_details, |
| 114 | + ) |
| 115 | + except FuturesTimeout: |
| 116 | + ctx.errors.append( |
| 117 | + ErrorDetail( |
| 118 | + code="AI_TIMEOUT", |
| 119 | + stage="AI", |
| 120 | + message=f"AI did not respond within {int(deadline_s)}s.", |
| 121 | + recoverable=True, |
| 122 | + ) |
| 123 | + ) |
| 124 | + except json.JSONDecodeError as e: |
| 125 | + ctx.errors.append( |
| 126 | + ErrorDetail(code="AI_SCHEMA_INVALID", stage="AI", message=f"JSON Decode Error: {str(e)}", recoverable=True) |
| 127 | + ) |
| 128 | + except Exception as e: |
| 129 | + ctx.errors.append( |
| 130 | + ErrorDetail( |
| 131 | + code="AI_TIMEOUT" if "timeout" in str(e).lower() else "AI_SCHEMA_INVALID", |
| 132 | + stage="AI", |
| 133 | + message=str(e), |
| 134 | + recoverable=True, |
| 135 | + ) |
| 136 | + ) |
| 137 | + |
| 138 | + return ctx |
0 commit comments