aur-shield v0.1.0: AI-powered AUR firewall
- FastAPI server with scan/build/repo endpoints - LLM scanner (Ollama) with regex pre-scan - makepkg/devtools builder with chroot isolation - Scan cache with TTL + PKGBUILD hash - Client installer + safe-yay wrapper - Docs + config example
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
"""AUR-Shield — AI-powered AUR firewall."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Main entry point for AUR-Shield."""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import uvicorn
|
||||
|
||||
from .config import Config
|
||||
from .server import init_server
|
||||
|
||||
|
||||
def main() -> None:
|
||||
config_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
|
||||
config = Config.load(config_path)
|
||||
|
||||
init_server(config)
|
||||
|
||||
print(f"AUR-Shield v0.1.0")
|
||||
print(f" Model: {config.ollama.model}")
|
||||
print(f" Ollama: {config.ollama.url}")
|
||||
print(f" Repo: {config.server.repo_dir}")
|
||||
print(f" Listening on {config.server.host}:{config.server.port}")
|
||||
|
||||
uvicorn.run(
|
||||
"aur_shield.server:app",
|
||||
host=config.server.host,
|
||||
port=config.server.port,
|
||||
reload=False,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,140 @@
|
||||
"""AUR API client — fetches PKGBUILDs and metadata from the AUR."""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
AUR_RPC = "https://aur.archlinux.org/rpc/v5"
|
||||
AUR_CGIT = "https://aur.archlinux.org/cgit/aur.git/plain"
|
||||
|
||||
|
||||
@dataclass
|
||||
class AURPackage:
|
||||
name: str
|
||||
version: str
|
||||
description: str
|
||||
url: str
|
||||
maintainer: str | None
|
||||
num_votes: int
|
||||
popularity: float
|
||||
last_modified: int
|
||||
pkgbase: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class AURSource:
|
||||
pkgbuild: str
|
||||
srcinfo: str
|
||||
package: AURPackage
|
||||
|
||||
|
||||
async def aur_info(name: str) -> AURPackage | None:
|
||||
"""Fetch package info from AUR RPC API."""
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
resp = await client.get(AUR_RPC, params={
|
||||
"type": "info",
|
||||
"v": "5",
|
||||
"arg[]": name,
|
||||
})
|
||||
resp.raise_for_status()
|
||||
data: dict[str, Any] = resp.json()
|
||||
|
||||
results = data.get("results", [])
|
||||
if not results:
|
||||
return None
|
||||
|
||||
r = results[0]
|
||||
return AURPackage(
|
||||
name=r.get("Name", name),
|
||||
version=r.get("Version", ""),
|
||||
description=r.get("Description", ""),
|
||||
url=r.get("URL", ""),
|
||||
maintainer=r.get("Maintainer"),
|
||||
num_votes=r.get("NumVotes", 0),
|
||||
popularity=r.get("Popularity", 0.0),
|
||||
last_modified=r.get("LastModified", 0),
|
||||
pkgbase=r.get("PackageBaseID", ""),
|
||||
)
|
||||
|
||||
|
||||
async def aur_sources(name: str) -> AURSource | None:
|
||||
"""Fetch PKGBUILD and .SRCINFO for a package."""
|
||||
pkg = await aur_info(name)
|
||||
if pkg is None:
|
||||
return None
|
||||
|
||||
pkgbase = r.get("PackageBase", name) if (r := await _raw_info(name)) else name
|
||||
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
# Fetch PKGBUILD
|
||||
pkgbuild_resp = await client.get(
|
||||
AUR_CGIT + f"/PKGBUILD",
|
||||
params={"h": pkgbase},
|
||||
)
|
||||
pkgbuild_resp.raise_for_status()
|
||||
pkgbuild = pkgbuild_resp.text
|
||||
|
||||
# Fetch .SRCINFO
|
||||
srcinfo_resp = await client.get(
|
||||
AUR_CGIT + f"/.SRCINFO",
|
||||
params={"h": pkgbase},
|
||||
)
|
||||
srcinfo = srcinfo_resp.text if srcinfo_resp.status_code == 200 else ""
|
||||
|
||||
return AURSource(
|
||||
pkgbuild=pkgbuild,
|
||||
srcinfo=srcinfo,
|
||||
package=pkg,
|
||||
)
|
||||
|
||||
|
||||
async def _raw_info(name: str) -> dict[str, Any] | None:
|
||||
"""Get raw RPC result for a package."""
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
resp = await client.get(AUR_RPC, params={
|
||||
"type": "info",
|
||||
"v": "5",
|
||||
"arg[]": name,
|
||||
})
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
results = data.get("results", [])
|
||||
return results[0] if results else None
|
||||
|
||||
|
||||
def extract_sources(pkgbuild: str) -> list[str]:
|
||||
"""Extract source URLs from a PKGBUILD."""
|
||||
sources = []
|
||||
in_array = False
|
||||
for line in pkgbuild.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped.startswith("source="):
|
||||
in_array = True
|
||||
# single-line: source=(url)
|
||||
m = re.findall(r'https?://[^\s)\'"]+|git://[^\s)\'"]+|ftp://[^\s)\'"]+',
|
||||
stripped)
|
||||
sources.extend(m)
|
||||
if ")" in stripped and not stripped.endswith("("):
|
||||
in_array = False
|
||||
elif in_array:
|
||||
m = re.findall(r'https?://[^\s)\'"]+|git://[^\s)\'"]+|ftp://[^\s)\'"]+',
|
||||
stripped)
|
||||
sources.extend(m)
|
||||
if ")" in stripped:
|
||||
in_array = False
|
||||
return sources
|
||||
|
||||
|
||||
def extract_install_hooks(pkgbuild: str) -> list[str]:
|
||||
"""Extract post_install/pre_install hooks from PKGBUILD."""
|
||||
hooks = []
|
||||
funcs = re.findall(
|
||||
r'(?:post_install|pre_install|post_upgrade|pre_upgrade|post_remove|pre_remove)\s*\(\)\s*\{[^}]*\}',
|
||||
pkgbuild,
|
||||
re.DOTALL,
|
||||
)
|
||||
return funcs
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Package builder — runs makepkg in a clean environment."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from .config import ServerConfig, BuildConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class BuildResult:
|
||||
success: bool
|
||||
pkg_files: list[str] # paths to .pkg.tar.zst
|
||||
log: str
|
||||
error: str = ""
|
||||
|
||||
|
||||
def build_package(
|
||||
pkgbuild: str,
|
||||
package_name: str,
|
||||
server_cfg: ServerConfig,
|
||||
build_cfg: BuildConfig,
|
||||
) -> BuildResult:
|
||||
"""Build a package from a PKGBUILD string."""
|
||||
|
||||
work_dir = Path(server_cfg.work_dir) / package_name
|
||||
work_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Write PKGBUILD
|
||||
pkgbuild_path = work_dir / "PKGBUILD"
|
||||
pkgbuild_path.write_text(pkgbuild)
|
||||
|
||||
# Clean previous build artifacts
|
||||
for f in work_dir.glob("*.pkg.tar.*"):
|
||||
f.unlink()
|
||||
|
||||
log_lines: list[str] = []
|
||||
|
||||
try:
|
||||
if build_cfg.use_devtools and shutil.which("extra-x86_64-build"):
|
||||
# Use devtools chroot for isolation
|
||||
cmd = [
|
||||
"extra-x86_64-build",
|
||||
"--", "-cC",
|
||||
]
|
||||
cwd = work_dir
|
||||
else:
|
||||
# Fallback: makepkg directly
|
||||
cmd = [
|
||||
"makepkg", "-sf", "--noconfirm", "--noprogressbar",
|
||||
]
|
||||
cwd = work_dir
|
||||
|
||||
log_lines.append(f"Building {package_name} with: {' '.join(cmd)}")
|
||||
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
cwd=cwd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=build_cfg.timeout,
|
||||
env={**os.environ, "MAKEFLAGS": "-j$(nproc)"},
|
||||
)
|
||||
|
||||
log_lines.append(result.stdout)
|
||||
if result.stderr:
|
||||
log_lines.append("--- stderr ---")
|
||||
log_lines.append(result.stderr)
|
||||
|
||||
if result.returncode != 0:
|
||||
return BuildResult(
|
||||
success=False,
|
||||
pkg_files=[],
|
||||
log="\n".join(log_lines),
|
||||
error=f"Build failed (exit {result.returncode})",
|
||||
)
|
||||
|
||||
# Find built packages
|
||||
pkg_files = [
|
||||
str(f) for f in work_dir.glob("*.pkg.tar.zst")
|
||||
]
|
||||
|
||||
if not pkg_files:
|
||||
return BuildResult(
|
||||
success=False,
|
||||
pkg_files=[],
|
||||
log="\n".join(log_lines),
|
||||
error="Build succeeded but no .pkg.tar.zst found",
|
||||
)
|
||||
|
||||
log_lines.append(f"Built: {pkg_files}")
|
||||
return BuildResult(
|
||||
success=True,
|
||||
pkg_files=pkg_files,
|
||||
log="\n".join(log_lines),
|
||||
)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return BuildResult(
|
||||
success=False,
|
||||
pkg_files=[],
|
||||
log="\n".join(log_lines),
|
||||
error=f"Build timed out after {build_cfg.timeout}s",
|
||||
)
|
||||
except Exception as e:
|
||||
return BuildResult(
|
||||
success=False,
|
||||
pkg_files=[],
|
||||
log="\n".join(log_lines),
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
|
||||
def add_to_repo(
|
||||
pkg_files: list[str],
|
||||
repo_dir: str,
|
||||
) -> str:
|
||||
"""Add built packages to the local pacman repo using repo-add."""
|
||||
|
||||
repo_path = Path(repo_dir)
|
||||
repo_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
db_name = "aur-shield.db"
|
||||
db_file = repo_path / f"{db_name}.tar.gz"
|
||||
|
||||
# Copy packages to repo dir
|
||||
copied = []
|
||||
for pkg in pkg_files:
|
||||
src = Path(pkg)
|
||||
dst = repo_path / src.name
|
||||
shutil.copy2(src, dst)
|
||||
copied.append(str(dst))
|
||||
|
||||
if not copied:
|
||||
return str(db_file)
|
||||
|
||||
# Run repo-add
|
||||
cmd = ["repo-add", str(db_file)] + copied
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"repo-add failed: {result.stderr}")
|
||||
|
||||
return str(db_file)
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Scan cache — stores scan results to avoid re-scanning unchanged packages."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass, asdict
|
||||
from pathlib import Path
|
||||
|
||||
from .scanner import ScanResult, ScanVerdict
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheEntry:
|
||||
package: str
|
||||
verdict: str
|
||||
confidence: float
|
||||
findings: list[str]
|
||||
reasoning: str
|
||||
timestamp: float
|
||||
pkgbuild_hash: str
|
||||
|
||||
|
||||
class ScanCache:
|
||||
"""JSON-file based cache for scan results."""
|
||||
|
||||
def __init__(self, cache_dir: str, ttl_hours: int = 168) -> None:
|
||||
self.cache_dir = Path(cache_dir)
|
||||
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.ttl_seconds = ttl_hours * 3600
|
||||
|
||||
def _cache_path(self, package: str) -> Path:
|
||||
return self.cache_dir / f"{package}.json"
|
||||
|
||||
def get(self, package: str, pkgbuild_hash: str) -> CacheEntry | None:
|
||||
"""Get cached scan if it exists and is fresh."""
|
||||
path = self._cache_path(package)
|
||||
if not path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
entry = CacheEntry(**json.loads(path.read_text()))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
# Check TTL
|
||||
if time.time() - entry.timestamp > self.ttl_seconds:
|
||||
return None
|
||||
|
||||
# Check if PKGBUILD changed
|
||||
if entry.pkgbuild_hash != pkgbuild_hash:
|
||||
return None
|
||||
|
||||
return entry
|
||||
|
||||
def put(self, result: ScanResult, pkgbuild_hash: str) -> None:
|
||||
"""Store a scan result."""
|
||||
entry = CacheEntry(
|
||||
package=result.package,
|
||||
verdict=result.verdict.value,
|
||||
confidence=result.confidence,
|
||||
findings=result.findings,
|
||||
reasoning=result.reasoning,
|
||||
timestamp=time.time(),
|
||||
pkgbuild_hash=pkgbuild_hash,
|
||||
)
|
||||
path = self._cache_path(result.package)
|
||||
path.write_text(json.dumps(asdict(entry), indent=2))
|
||||
|
||||
def list_all(self) -> list[dict]:
|
||||
"""List all cached entries."""
|
||||
entries = []
|
||||
for path in self.cache_dir.glob("*.json"):
|
||||
try:
|
||||
entry = CacheEntry(**json.loads(path.read_text()))
|
||||
entries.append(asdict(entry))
|
||||
except Exception:
|
||||
continue
|
||||
return entries
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Configuration loader for AUR-Shield."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
@dataclass
|
||||
class OllamaConfig:
|
||||
url: str = "http://localhost:11434"
|
||||
model: str = "qwen2.5:latest"
|
||||
timeout: int = 120
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerConfig:
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8443
|
||||
repo_dir: str = "/var/cache/aur-shield/repo"
|
||||
work_dir: str = "/var/cache/aur-shield/build"
|
||||
build_user: str = "nobody"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SecurityConfig:
|
||||
block_patterns: list[str] = field(default_factory=lambda: [
|
||||
r"curl.*\|.*bash",
|
||||
r"wget.*/tmp/.*\|.*sh",
|
||||
r"eval.*base64",
|
||||
r"nc\s+-.*\d+",
|
||||
r"/dev/tcp/",
|
||||
r"systemctl.*enable.*--now",
|
||||
])
|
||||
max_pkg_size_mb: int = 500
|
||||
allowed_source_schemes: list[str] = field(default_factory=lambda: [
|
||||
"https", "http", "git", "ftp", "file",
|
||||
])
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheConfig:
|
||||
ttl_hours: int = 168
|
||||
max_cache_gb: int = 10
|
||||
|
||||
|
||||
@dataclass
|
||||
class BuildConfig:
|
||||
use_devtools: bool = True
|
||||
timeout: int = 600
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
ollama: OllamaConfig = field(default_factory=OllamaConfig)
|
||||
server: ServerConfig = field(default_factory=ServerConfig)
|
||||
security: SecurityConfig = field(default_factory=SecurityConfig)
|
||||
cache: CacheConfig = field(default_factory=CacheConfig)
|
||||
build: BuildConfig = field(default_factory=BuildConfig)
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str | Path | None = None) -> "Config":
|
||||
"""Load config from YAML file, falling back to defaults."""
|
||||
if path is None:
|
||||
path = Path(os.environ.get(
|
||||
"AUR_SHIELD_CONFIG",
|
||||
"config.yaml",
|
||||
))
|
||||
|
||||
path = Path(path)
|
||||
if not path.exists():
|
||||
return cls()
|
||||
|
||||
with open(path) as f:
|
||||
raw: dict[str, Any] = yaml.safe_load(f) or {}
|
||||
|
||||
cfg = cls()
|
||||
if "ollama" in raw:
|
||||
cfg.ollama = OllamaConfig(**raw["ollama"])
|
||||
if "server" in raw:
|
||||
cfg.server = ServerConfig(**raw["server"])
|
||||
if "security" in raw:
|
||||
cfg.security = SecurityConfig(**raw["security"])
|
||||
if "cache" in raw:
|
||||
cfg.cache = CacheConfig(**raw["cache"])
|
||||
if "build" in raw:
|
||||
cfg.build = BuildConfig(**raw["build"])
|
||||
return cfg
|
||||
@@ -0,0 +1,179 @@
|
||||
"""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],
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
"""FastAPI server for AUR-Shield."""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import FileResponse, JSONResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from .config import Config
|
||||
from .aur_client import aur_sources, aur_info
|
||||
from .scanner import scan_pkgbuild, ScanVerdict
|
||||
from .builder import build_package, add_to_repo, BuildResult
|
||||
from .cache import ScanCache
|
||||
|
||||
app = FastAPI(
|
||||
title="AUR-Shield",
|
||||
description="AI-powered AUR firewall",
|
||||
version="0.1.0",
|
||||
)
|
||||
|
||||
# Global config — set by main()
|
||||
_config: Config | None = None
|
||||
_scan_cache: ScanCache | None = None
|
||||
|
||||
|
||||
def get_config() -> Config:
|
||||
if _config is None:
|
||||
raise RuntimeError("Server not initialized")
|
||||
return _config
|
||||
|
||||
|
||||
def get_cache() -> ScanCache:
|
||||
if _scan_cache is None:
|
||||
raise RuntimeError("Cache not initialized")
|
||||
return _scan_cache
|
||||
|
||||
|
||||
def _hash_pkgbuild(pkgbuild: str) -> str:
|
||||
return hashlib.sha256(pkgbuild.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
@app.get("/api/status")
|
||||
async def status() -> dict[str, Any]:
|
||||
"""Server status and cache info."""
|
||||
cfg = get_config()
|
||||
cache = get_cache()
|
||||
return {
|
||||
"status": "running",
|
||||
"model": cfg.ollama.model,
|
||||
"ollama_url": cfg.ollama.url,
|
||||
"repo_dir": cfg.server.repo_dir,
|
||||
"cached_scans": len(cache.list_all()),
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/scan/{package}")
|
||||
async def scan_package(package: str) -> dict[str, Any]:
|
||||
"""Scan a package without building."""
|
||||
cfg = get_config()
|
||||
cache = get_cache()
|
||||
|
||||
# Fetch from AUR
|
||||
source = await aur_sources(package)
|
||||
if source is None:
|
||||
raise HTTPException(404, f"Package '{package}' not found in AUR")
|
||||
|
||||
pkgbuild_hash = _hash_pkgbuild(source.pkgbuild)
|
||||
|
||||
# Check cache
|
||||
cached = cache.get(package, pkgbuild_hash)
|
||||
if cached:
|
||||
return {
|
||||
"package": package,
|
||||
"cached": True,
|
||||
"verdict": cached.verdict,
|
||||
"confidence": cached.confidence,
|
||||
"findings": cached.findings,
|
||||
"reasoning": cached.reasoning,
|
||||
"version": source.package.version,
|
||||
}
|
||||
|
||||
# Scan
|
||||
result = await scan_pkgbuild(
|
||||
package, source.pkgbuild, source.srcinfo,
|
||||
cfg.ollama, cfg.security,
|
||||
)
|
||||
|
||||
# Cache the result
|
||||
cache.put(result, pkgbuild_hash)
|
||||
|
||||
return {
|
||||
"package": package,
|
||||
"cached": False,
|
||||
"verdict": result.verdict.value,
|
||||
"confidence": result.confidence,
|
||||
"findings": result.findings,
|
||||
"reasoning": result.reasoning,
|
||||
"sources": result.sources_checked,
|
||||
"version": source.package.version,
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/build/{package}")
|
||||
async def build_endpoint(package: str) -> dict[str, Any]:
|
||||
"""Scan + build + add to repo."""
|
||||
cfg = get_config()
|
||||
cache = get_cache()
|
||||
|
||||
# 1. Fetch from AUR
|
||||
source = await aur_sources(package)
|
||||
if source is None:
|
||||
raise HTTPException(404, f"Package '{package}' not found in AUR")
|
||||
|
||||
pkgbuild_hash = _hash_pkgbuild(source.pkgbuild)
|
||||
|
||||
# 2. Check cache for scan
|
||||
cached = cache.get(package, pkgbuild_hash)
|
||||
if cached:
|
||||
if cached.verdict == ScanVerdict.MALICIOUS.value:
|
||||
raise HTTPException(403, f"Package '{package}' is cached as MALICIOUS")
|
||||
verdict = cached.verdict
|
||||
findings = cached.findings
|
||||
confidence = cached.confidence
|
||||
else:
|
||||
# 3. Scan
|
||||
result = await scan_pkgbuild(
|
||||
package, source.pkgbuild, source.srcinfo,
|
||||
cfg.ollama, cfg.security,
|
||||
)
|
||||
cache.put(result, pkgbuild_hash)
|
||||
verdict = result.verdict.value
|
||||
findings = result.findings
|
||||
confidence = result.confidence
|
||||
|
||||
# 4. Block if malicious
|
||||
if verdict == ScanVerdict.MALICIOUS.value:
|
||||
raise HTTPException(
|
||||
403,
|
||||
f"Package '{package}' was flagged as MALICIOUS. Findings: {findings}",
|
||||
)
|
||||
|
||||
if verdict == ScanVerdict.SUSPICIOUS.value:
|
||||
# Allow suspicious but warn
|
||||
pass
|
||||
|
||||
# 5. Build
|
||||
build_result = build_package(
|
||||
source.pkgbuild, package,
|
||||
cfg.server, cfg.build,
|
||||
)
|
||||
|
||||
if not build_result.success:
|
||||
return {
|
||||
"package": package,
|
||||
"verdict": verdict,
|
||||
"confidence": confidence,
|
||||
"build_success": False,
|
||||
"build_error": build_result.error,
|
||||
"build_log": build_result.log[-2000:],
|
||||
}
|
||||
|
||||
# 6. Add to repo
|
||||
try:
|
||||
db_path = add_to_repo(build_result.pkg_files, cfg.server.repo_dir)
|
||||
except Exception as e:
|
||||
return {
|
||||
"package": package,
|
||||
"verdict": verdict,
|
||||
"confidence": confidence,
|
||||
"build_success": True,
|
||||
"repo_success": False,
|
||||
"repo_error": str(e),
|
||||
}
|
||||
|
||||
return {
|
||||
"package": package,
|
||||
"verdict": verdict,
|
||||
"confidence": confidence,
|
||||
"findings": findings,
|
||||
"build_success": True,
|
||||
"repo_success": True,
|
||||
"pkg_files": [Path(f).name for f in build_result.pkg_files],
|
||||
"repo_db": Path(db_path).name,
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/report/{package}")
|
||||
async def get_report(package: str) -> dict[str, Any]:
|
||||
"""Get the last scan report for a package."""
|
||||
cache = get_cache()
|
||||
entries = cache.list_all()
|
||||
for entry in entries:
|
||||
if entry["package"] == package:
|
||||
return entry
|
||||
raise HTTPException(404, f"No scan report for '{package}'")
|
||||
|
||||
|
||||
@app.get("/api/cache")
|
||||
async def list_cache() -> dict[str, Any]:
|
||||
"""List all cached scan results."""
|
||||
cache = get_cache()
|
||||
return {"entries": cache.list_all()}
|
||||
|
||||
|
||||
@app.delete("/api/cache/{package}")
|
||||
async def clear_cache_entry(package: str) -> dict[str, str]:
|
||||
"""Clear a cached scan result."""
|
||||
cache = get_cache()
|
||||
path = cache._cache_path(package)
|
||||
if path.exists():
|
||||
path.unlink()
|
||||
return {"status": "deleted", "package": package}
|
||||
raise HTTPException(404, f"No cache for '{package}'")
|
||||
|
||||
|
||||
# Serve the pacman repo files
|
||||
@app.get("/repo/{filename}")
|
||||
async def serve_repo_file(filename: str):
|
||||
"""Serve repo files for pacman."""
|
||||
cfg = get_config()
|
||||
repo_path = Path(cfg.server.repo_dir) / filename
|
||||
if not repo_path.exists():
|
||||
raise HTTPException(404, f"File '{filename}' not in repo")
|
||||
return FileResponse(repo_path)
|
||||
|
||||
|
||||
def init_server(config: Config) -> None:
|
||||
"""Initialize the server with config."""
|
||||
global _config, _scan_cache
|
||||
_config = config
|
||||
_scan_cache = ScanCache(
|
||||
cache_dir=str(Path(config.server.work_dir) / "cache"),
|
||||
ttl_hours=config.cache.ttl_hours,
|
||||
)
|
||||
# Ensure dirs exist
|
||||
Path(config.server.repo_dir).mkdir(parents=True, exist_ok=True)
|
||||
Path(config.server.work_dir).mkdir(parents=True, exist_ok=True)
|
||||
Reference in New Issue
Block a user