aur-shield/aur_shield/scanner.py

"""LLM scanner — sends PKGBUILD to Ollama for security analysis."""
from __future__ import annotations

import json
import re
import subprocess
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

import httpx

from .config import OllamaConfig, SecurityConfig
from .aur_client import extract_sources, extract_install_hooks


class ScanVerdict(str, Enum):
    CLEAN = "clean"
    SUSPICIOUS = "suspicious"
    MALICIOUS = "malicious"
    ERROR = "error"


@dataclass
class ScanResult:
    package: str
    verdict: ScanVerdict
    confidence: float  # 0.0 - 1.0
    findings: list[str] = field(default_factory=list)
    reasoning: str = ""
    sources_checked: list[str] = field(default_factory=list)
    raw_response: str = ""


SYSTEM_PROMPT = """\
You are a security scanner for Arch Linux AUR PKGBUILD files.
Your job is to detect malicious or suspicious code in PKGBUILDs.

Check for these threats:
1. DOWNLOAD-AND-EXEC: curl|bash, wget|sh, downloading and executing scripts
2. OBFUSCATION: eval $(base64 -d), hex-encoded commands, printf-encoded payloads
3. REVERSE SHELLS: nc -, /dev/tcp, bash -i >&, socat
4. SUSPICIOUS SOURCES: npm, tor, raw IP addresses, non-HTTPS with no fallback
5. PERSISTENCE: post_install creating systemd services, cronjobs, autostart entries
6. TYPOSQUATTING: package names mimicking popular packages with slight misspellings
7. DATA EXFIL: curl/wget sending data to external servers with system info
8. UNUSUAL DEPENDS: packages requesting network tools unrelated to their purpose

Respond in JSON ONLY:
{
  "verdict": "clean" | "suspicious" | "malicious",
  "confidence": 0.0-1.0,
  "findings": ["list of specific issues found"],
  "reasoning": "brief explanation"
}

Be conservative: if unsure, mark as suspicious.
A clean PKGBUILD that just downloads a tarball from the official upstream
URL and runs make/cmake is CLEAN.
"""


def _pre_scan(pkgbuild: str, sec_cfg: SecurityConfig) -> list[str]:
    """Fast regex pre-scan before LLM. Returns list of findings."""
    findings = []
    for pattern in sec_cfg.block_patterns:
        if re.search(pattern, pkgbuild, re.IGNORECASE):
            findings.append(f"Pattern match: {pattern}")
    return findings


async def scan_pkgbuild(
    package_name: str,
    pkgbuild: str,
    srcinfo: str,
    ollama_cfg: OllamaConfig,
    sec_cfg: SecurityConfig,
) -> ScanResult:
    """Scan a PKGBUILD with regex pre-check + LLM analysis."""

    # 1. Fast regex pre-scan
    regex_findings = _pre_scan(pkgbuild, sec_cfg)

    # 2. Extract metadata for context
    sources = extract_sources(pkgbuild)
    hooks = extract_install_hooks(pkgbuild)

    # 3. Build the prompt
    user_msg = (
        f"Package: {package_name}\n\n"
        f"=== PKGBUILD ===\n{pkgbuild}\n\n"
    )
    if srcinfo:
        user_msg += f"=== .SRCINFO (excerpt) ===\n{srcinfo[:2000]}\n\n"
    if hooks:
        user_msg += f"=== INSTALL HOOKS ===\n{chr(10).join(hooks)}\n\n"
    if sources:
        user_msg += f"=== SOURCE URLS ===\n{chr(10).join(sources)}\n\n"

    # 4. Call Ollama
    payload = {
        "model": ollama_cfg.model,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_msg},
        ],
        "stream": False,
        "format": "json",
        "options": {
            "temperature": 0.1,
            "num_ctx": 8192,
        },
    }

    try:
        async with httpx.AsyncClient(timeout=ollama_cfg.timeout) as client:
            resp = await client.post(
                f"{ollama_cfg.url}/api/chat",
                json=payload,
            )
            resp.raise_for_status()
            data = resp.json()

        raw = data.get("message", {}).get("content", "")
        result = _parse_llm_response(raw)

        # Combine with regex findings
        all_findings = result.get("findings", []) + regex_findings
        verdict = result.get("verdict", "suspicious")
        confidence = result.get("confidence", 0.5)

        # If regex found blocked patterns, escalate
        if regex_findings and verdict == "clean":
            verdict = "suspicious"
            confidence = max(confidence, 0.7)

        return ScanResult(
            package=package_name,
            verdict=ScanVerdict(verdict),
            confidence=confidence,
            findings=all_findings,
            reasoning=result.get("reasoning", ""),
            sources_checked=sources,
            raw_response=raw,
        )

    except Exception as e:
        return ScanResult(
            package=package_name,
            verdict=ScanVerdict.ERROR,
            confidence=0.0,
            findings=[f"Scanner error: {e}"],
            sources_checked=sources,
        )


def _parse_llm_response(raw: str) -> dict[str, Any]:
    """Parse the LLM JSON response, handling malformed output."""
    # Try direct JSON parse
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        pass

    # Try to extract JSON from text
    match = re.search(r'\{[^{}]*"verdict"[^{}]*\}', raw, re.DOTALL)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass

    # Fallback
    return {
        "verdict": "suspicious",
        "confidence": 0.3,
        "findings": ["Could not parse LLM response"],
        "reasoning": raw[:500],
    }