aiscan-security 0.0.2__py3-none-any.whl

aiscan/api/routes.py

@@ -0,0 +1,177 @@
+"""
+REST API — all 5 scan type endpoints.
+POST /scan/prompt
+POST /scan/agent
+POST /scan/api
+POST /scan/mcp
+POST /scan/integration
+POST /scan/all  — runs all 5 scanners, returns combined results
+GET  /health
+"""
+import os
+import tempfile
+
+from fastapi import FastAPI, UploadFile, File
+from pydantic import BaseModel
+
+from aiscan.models import ScanType, Severity
+from aiscan.scanners.agent_scanner import scan_agent
+from aiscan.scanners.api_scanner import scan_api
+from aiscan.scanners.integration_scanner import scan_integration
+from aiscan.scanners.mcp_scanner import scan_mcp
+from aiscan.scanners.prompt_scanner import scan_prompt
+from aiscan.scanners.tm_coding_standards_scanner import tm_coding_standards_scanner
+from aiscan.utils.aggregator import aggregate
+
+app = FastAPI(
+    title="AI Security Scanner",
+    description="Scan prompts, agents, APIs, MCP tools and integrations for security vulnerabilities",
+    version="0.0.2",
+)
+
+
+class ScanRequest(BaseModel):
+    content: str
+    filename: str = "input.txt"
+
+
+class FindingOut(BaseModel):
+    rule_id: str
+    severity: str
+    threat: str
+    title: str
+    detail: str
+    file: str
+    line: int
+    fix: str
+
+
+class ScanResponse(BaseModel):
+    scan_type: str
+    passed: bool
+    critical: int
+    high: int
+    medium: int
+    info: int
+    total: int
+    findings: list[FindingOut]
+
+
+def _run_scan(req: ScanRequest, scan_fn, scan_type: ScanType, suffix: str = ".txt") -> ScanResponse:
+    with tempfile.NamedTemporaryFile(mode="w", suffix=suffix, delete=False, encoding="utf-8") as f:
+        f.write(req.content)
+        tmp = f.name
+    try:
+        result = scan_fn(tmp)
+        result.findings = aggregate(result.findings)
+        return _shape(result, scan_type)
+    finally:
+        os.unlink(tmp)
+
+
+def _shape(result, scan_type: ScanType) -> ScanResponse:
+    findings = [{
+        "rule_id": f.rule_id, "severity": f.severity.value, "threat": f.threat_type.value,
+        "title": f.title, "detail": f.detail, "file": f.file, "line": f.line, "fix": f.fix,
+    } for f in result.findings]
+    return ScanResponse(
+        scan_type=scan_type.value, passed=result.passed,
+        critical=result.critical_count, high=result.high_count,
+        medium=sum(1 for f in result.findings if f.severity == Severity.MEDIUM),
+        info=sum(1 for f in result.findings if f.severity == Severity.INFO),
+        total=len(result.findings), findings=findings,
+    )
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "service": "ai-security-scanner", "version": "0.0.2",
+            "scan_types": ["prompt", "agent", "api", "mcp", "integration"]}
+
+
+@app.post("/scan/prompt", response_model=ScanResponse)
+async def scan_prompt_ep(req: ScanRequest):
+    return _run_scan(req, scan_prompt, ScanType.PROMPT, ".txt")
+
+
+@app.post("/scan/agent", response_model=ScanResponse)
+async def scan_agent_ep(req: ScanRequest):
+    return _run_scan(req, scan_agent, ScanType.AGENT, ".yaml")
+
+
+@app.post("/scan/api", response_model=ScanResponse)
+async def scan_api_ep(req: ScanRequest):
+    suffix = "." + req.filename.rsplit(".", 1)[-1] if "." in req.filename else ".yaml"
+    return _run_scan(req, scan_api, ScanType.API, suffix)
+
+
+@app.post("/scan/mcp", response_model=ScanResponse)
+async def scan_mcp_ep(req: ScanRequest):
+    return _run_scan(req, scan_mcp, ScanType.MCP, ".json")
+
+
+@app.post("/scan/integration", response_model=ScanResponse)
+async def scan_integration_ep(req: ScanRequest):
+    return _run_scan(req, scan_integration, ScanType.INTEGRATION, ".yaml")
+
+
+@app.post("/scan/all")
+async def scan_all(req: ScanRequest):
+    """Run all 5 scanners against the same content. Useful for mixed config files."""
+    results = []
+    for fn, st, sfx in [
+        (scan_prompt, ScanType.PROMPT, ".txt"),
+        (scan_agent, ScanType.AGENT, ".yaml"),
+        (scan_api, ScanType.API, ".yaml"),
+        (scan_mcp, ScanType.MCP, ".json"),
+        (scan_integration, ScanType.INTEGRATION, ".yaml"),
+    ]:
+        results.append(_run_scan(req, fn, st, sfx).model_dump())
+    total_findings = sum(r["total"] for r in results)
+    total_critical = sum(r["critical"] for r in results)
+    return {
+        "passed": total_critical == 0 and sum(r["high"] for r in results) == 0,
+        "total_findings": total_findings,
+        "total_critical": total_critical,
+        "results": results,
+    }
+
+
+# coding standards
+
+@app.post("/tm-coding-standards")
+async def tm_coding_standards(file: UploadFile = File(...)):
+    suffix = os.path.splitext(file.filename)[1]
+
+    with tempfile.NamedTemporaryFile(
+            delete=False,
+            suffix=suffix
+    ) as tmp:
+
+        content = await file.read()
+        tmp.write(content)
+
+        temp_path = tmp.name
+
+    try:
+        result = tm_coding_standards_scanner(temp_path)
+
+        return {
+            "target": result.target,
+            "total_findings": len(result.findings),
+            "findings": [
+                {
+                    "rule_id": f.rule_id,
+                    "severity": f.severity.value if hasattr(f.severity, "value") else str(f.severity),
+                    "title": f.title,
+                    "detail": f.detail,
+                    "file": f.file,
+                    "line": f.line,
+                    "fix": f.fix,
+                }
+                for f in result.findings
+            ]
+        }
+
+    finally:
+        os.unlink(temp_path)

aiscan/cli.py

@@ -0,0 +1,387 @@
+"""
+aiscan — AI Security Scanner CLI
+
+A generic, installable command-line tool to scan AI system components
+for security vulnerabilities.
+
+Usage after install:
+  aiscan scan prompt.txt
+  aiscan scan --type api openapi.yaml
+  aiscan scan --type all config.yaml --output json
+  aiscan init
+  aiscan serve
+"""
+import importlib
+from pathlib import Path
+from typing import Optional
+
+import typer
+from rich.console import Console
+from rich.panel import Panel
+
+from aiscan.rules.registry import PYTHON_RULES
+
+app = typer.Typer(
+    name="aiscan",
+    help="AI Security Scanner — scan prompts, agents, APIs, MCP tools and integrations.",
+    add_completion=True,
+    no_args_is_help=True,
+)
+console = Console()
+
+# ── scan type → (module, function, default suffix) ────────────────────────────
+SCAN_REGISTRY = {
+    "prompt": ("aiscan.scanners.prompt_scanner", "scan_prompt", ".txt"),
+    "agent": ("aiscan.scanners.agent_scanner", "scan_agent", ".yaml"),
+    "api": ("aiscan.scanners.api_scanner", "scan_api", ".yaml"),
+    "mcp": ("aiscan.scanners.mcp_scanner", "scan_mcp", ".json"),
+    "integration": ("aiscan.scanners.integration_scanner", "scan_integration", ".yaml"),
+    "code": ("aiscan.scanners.api_scanner", "scan_code_api", ".py"),
+}
+
+# Auto-detect scan type from file extension
+AUTO_TYPE = {
+    # Prompt files
+    ".py": "prompt", ".txt": "prompt", ".md": "prompt", ".j2": "prompt", ".jinja": "prompt",
+    # Code files — scanned with full rule set
+    ".js": "prompt", ".ts": "prompt", ".jsx": "prompt",
+    ".tsx": "prompt", ".go": "prompt", ".java": "prompt", ".rb": "prompt",
+    ".php": "prompt", ".cs": "prompt", ".cpp": "prompt", ".c": "prompt",
+    ".rs": "prompt", ".swift": "prompt", ".kt": "prompt",
+    # API / config files
+    ".yaml": "api", ".yml": "api", ".json": "mcp",
+    # Integration / environment files
+    ".env": "integration", ".toml": "integration",
+    ".ini": "integration", ".cfg": "integration", ".conf": "integration",
+    ".properties": "integration",
+    # code file
+    ".py": "code",
+}
+
+# All file extensions aiscan will pick up when scanning a directory
+SCANNABLE_EXTENSIONS = set(AUTO_TYPE.keys())
+
+
+def _resolve_type(file: Path, declared: str) -> list[str]:
+    if declared == "all":
+        return list(SCAN_REGISTRY.keys())
+
+    if declared == "auto":
+        ext = file.suffix.lower()
+
+        # Python files → run BOTH scanners
+        if ext == ".py":
+            return ["prompt", "code"]
+
+        return [AUTO_TYPE.get(ext, "prompt")]
+
+    return [declared]
+
+
+def _run_one(file: Path, scan_type: str):
+    """Run a single scanner and return ScanResult."""
+    mod_name, fn_name, _ = SCAN_REGISTRY[scan_type]
+    mod = importlib.import_module(mod_name)
+    return getattr(mod, fn_name)(str(file))
+
+
+# ── scan command ─────────────────────────────────────────────────────────────
+@app.command()
+def scan(
+        files: list[Path] = typer.Argument(
+            ..., help="One or more files to scan. Globs are expanded by the shell.",
+            exists=True, readable=True,
+        ),
+        type: str = typer.Option(
+            "auto",
+            "--type", "-t",
+            help="Scan type: prompt | agent | api | mcp | integration |code | all | auto",
+        ),
+        output: str = typer.Option(
+            "console",
+            "--output", "-o",
+            help="Output format: console | json | sarif | pdf",
+        ),
+        out: Optional[Path] = typer.Option(
+            None,
+            "--out", "-O",
+            help="Write output to file instead of stdout",
+        ),
+        fail_on: str = typer.Option(
+            "high",
+            "--fail-on", "-f",
+            help="Exit code 1 if any finding reaches this severity: critical | high | medium | none",
+        ),
+        no_ai: bool = typer.Option(
+            False,
+            "--no-ai",
+            help="Skip AI deep scan — run static rules only (faster, no API key needed)",
+        ),
+        quiet: bool = typer.Option(
+            False,
+            "--quiet", "-q",
+            help="Suppress all output except errors. Useful in CI.",
+        ),
+):
+    """
+   Analyze one or more files for AI security threats, code quality issues, and compliance with TMotions Python coding standards.
+
+    Examples:
+
+      aiscan scan system_prompt.txt
+
+      aiscan scan --type api openapi.yaml
+
+      aiscan scan --type all config.yaml --output json
+
+      aiscan scan --type prompt prompts/ --output sarif --out results.sarif
+
+      aiscan scan --no-ai --fail-on critical agent.yaml
+    """
+    from aiscan.utils.aggregator import aggregate
+    from aiscan.reporters.console_reporter import print_results
+    from aiscan.reporters.json_reporter import to_json, to_sarif
+
+    if no_ai:
+        import os
+        os.environ["AI_PROVIDER"] = ""
+
+    results = []
+    for file in files:
+        if file.is_dir():
+            # Expand directory — scan all recognised file types inside
+            all_files = [
+                f for f in file.rglob("*")
+                if f.is_file() and f.suffix.lower() in SCANNABLE_EXTENSIONS
+            ]
+            if not all_files and not quiet:
+                console.print(f"[yellow]No scannable files found in {file}[/yellow]")
+            for f in all_files:
+                for st in _resolve_type(f, type):
+                    r = _run_one(f, st)
+                    r.findings = aggregate(r.findings)
+                    results.append(r)
+        else:
+            for st in _resolve_type(file, type):
+                r = _run_one(file, st)
+                r.findings = aggregate(r.findings)
+                results.append(r)
+
+    if not results:
+        if not quiet:
+            console.print("[yellow]No files scanned.[/yellow]")
+        raise typer.Exit(0)
+
+    # Output
+    if output == "json":
+        rendered = to_json(results)
+        if out:
+            out.write_text(rendered)
+            if not quiet:
+                console.print(f"[green]JSON written to {out}[/green]")
+        else:
+            print(rendered)
+
+    elif output == "sarif":
+        rendered = to_sarif(results)
+        if out:
+            out.write_text(rendered)
+            if not quiet:
+                console.print(f"[green]SARIF written to {out}[/green]")
+        else:
+            print(rendered)
+
+    elif output == "pdf":
+        from aiscan.reporters.pdf_reporter import generate_pdf
+        pdf_path = str(out) if out else "aiscan-report.pdf"
+        generate_pdf(results, pdf_path)
+        if not quiet:
+            console.print(f"[green]PDF report saved to {pdf_path}[/green]")
+
+    else:
+        if not quiet:
+            print_results(results)
+
+    # CI exit code
+    FAIL_RANK = {"critical": 4, "high": 3, "medium": 2, "none": 0}
+    from aiscan.models import Severity
+    SEV_RANK = {Severity.CRITICAL: 4, Severity.HIGH: 3, Severity.MEDIUM: 2, Severity.INFO: 1}
+    threshold = FAIL_RANK.get(fail_on, 3)
+    max_sev = max((SEV_RANK.get(f.severity, 0) for r in results for f in r.findings), default=0)
+    if max_sev >= threshold:
+        raise typer.Exit(code=1)
+
+
+# ── serve command ─────────────────────────────────────────────────────────────
+@app.command()
+def serve(
+        host: str = typer.Option("0.0.0.0", "--host", help="Host to bind"),
+        port: int = typer.Option(8000, "--port", help="Port to listen on"),
+        reload: bool = typer.Option(False, "--reload", help="Auto-reload on code changes (dev mode)"),
+):
+    """
+    Start the AI Security Scanner REST API server.
+
+    Endpoints:
+      GET  /health
+      POST /scan/prompt
+      POST /scan/agent
+      POST /scan/api
+      POST /scan/mcp
+      POST /scan/integration
+      POST /scan/all
+      GET  /docs  (Swagger UI)
+
+    Example:
+      aiscan serve --port 8000 --reload
+    """
+    try:
+        import uvicorn
+    except ImportError:
+        console.print("[red]Install the api extra to use serve: pip install aiscan[api][/red]")
+        raise typer.Exit(1)
+
+    from aiscan.api import routes  # noqa: F401 — triggers app creation
+    console.print(Panel(
+        f"[bold]AI Security Scanner API[/bold]\n"
+        f"http://{host}:{port}\n"
+        f"Swagger UI: http://{host}:{port}/docs",
+        title="aiscan serve",
+    ))
+    uvicorn.run("aiscan.api.routes:app", host=host, port=port, reload=reload)
+
+
+# ── rules command ─────────────────────────────────────────────────────────────
+@app.command()
+def rules(
+        type: Optional[str] = typer.Argument(None,
+                                             help="Filter by scan type: prompt | agent | api | mcp | integration"),
+        severity: Optional[str] = typer.Option(None, "--severity", "-s",
+                                               help="Filter by severity: critical | high | medium | info"),
+):
+    """
+    List all security rules in the registry.
+
+    Examples:
+      aiscan rules
+      aiscan rules prompt
+      aiscan rules api --severity critical
+    """
+    from rich.table import Table
+    from rich import box
+    from aiscan.rules.registry import RULES
+    from aiscan.models import ScanType
+
+    SEV_STYLE = {
+        "critical": "bold red",
+        "high": "red",
+        "medium": "yellow",
+        "info": "blue",
+        "low": "dim",
+    }
+
+    count = 0
+
+    table = Table(box=box.SIMPLE, show_header=True, header_style="bold")
+    table.add_column("ID", width=9)
+    table.add_column("Target", width=12)
+    table.add_column("Threat", width=20)
+    table.add_column("Severity", width=10)
+    table.add_column("Title", width=48)
+
+    filtered = RULES
+    if type:
+        try:
+            st = ScanType(type)
+            filtered = [r for r in filtered if st in r["scan_types"]]
+        except ValueError:
+            console.print(f"[red]Unknown type '{type}'. Use: {', '.join(SCAN_REGISTRY)}[/red]")
+            raise typer.Exit(1)
+
+    if severity:
+        filtered = [r for r in filtered if r["severity"].value == severity]
+
+    for r in filtered:
+        sev = r["severity"].value
+        targets = ", ".join(s.value for s in r["scan_types"])
+        table.add_row(
+            r["id"], targets, r["threat"].value, sev, r["title"],
+            style=SEV_STYLE.get(sev, ""),
+        )
+
+    if not type or type.lower() == "code":
+
+        for r in PYTHON_RULES:
+            sev = r["severity"].value.lower()
+
+            if severity and sev != severity.lower():
+                continue
+
+            table.add_row(
+                f"[{SEV_STYLE.get(sev, '')}]{r['id']}[/]",
+                f"[{SEV_STYLE.get(sev, '')}]code[/]",
+                f"[{SEV_STYLE.get(sev, '')}]TM_Coding_Standards[/]",
+                f"[{SEV_STYLE.get(sev, '')}]{sev.upper()}[/]",
+                f"[{SEV_STYLE.get(sev, '')}]{r['title']}[/]",
+            )
+            count += 1
+
+    console.print(table)
+    console.print(f"\n[dim]{len(filtered)} rule(s) shown[/dim]")
+
+
+# ── init command ──────────────────────────────────────────────────────────────
+@app.command()
+def init(
+        provider: str = typer.Option("claude", "--provider", "-p", help="AI provider: claude | openai | vertex"),
+        force: bool = typer.Option(False, "--force", help="Overwrite existing .env file"),
+):
+    """
+    Initialise a new .env config file in the current directory.
+
+    Example:
+      aiscan init
+      aiscan init --provider openai
+    """
+    env_path = Path(".env")
+    if env_path.exists() and not force:
+        console.print("[yellow].env already exists. Use --force to overwrite.[/yellow]")
+        raise typer.Exit(0)
+
+    content = f"""# aiscan configuration
+# Generated by: aiscan init --provider {provider}
+
+# Active provider: claude | openai | vertex
+AI_PROVIDER={provider}
+
+# Anthropic (Claude)
+ANTHROPIC_API_KEY=sk-ant-...
+SCANNER_MODEL=claude-sonnet-4-20250514
+
+# OpenAI (GPT-4o)
+OPENAI_API_KEY=sk-...
+OPENAI_MODEL=gpt-4o
+
+# Vertex AI (Gemini)
+VERTEX_PROJECT=your-gcp-project-id
+VERTEX_LOCATION=us-central1
+VERTEX_MODEL=gemini-2.5-flash-001
+
+# Limits
+MAX_TOKENS=2000
+LOG_LEVEL=INFO
+"""
+    env_path.write_text(content)
+    console.print(f"[green]Created .env — set your {provider.upper()}_API_KEY and run:[/green]")
+    console.print(f"  aiscan scan --type prompt your_prompt.txt")
+
+
+# ── version command ───────────────────────────────────────────────────────────
+@app.command()
+def version():
+    """Show aiscan version."""
+    console.print("aiscan 0.0.2")
+
+
+if __name__ == "__main__":
+    app()

aiscan/config.py

@@ -0,0 +1,23 @@
+from typing import Literal
+
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    log_level: str = "INFO"
+    ai_provider: Literal["claude", "openai", "vertex", ""] = "claude"
+    anthropic_api_key: str = ""
+    scanner_model: str = "claude-sonnet-4-20250514"
+    openai_api_key: str = ""
+    openai_model: str = "gpt-4o"
+    vertex_project: str = ""
+    vertex_location: str = "us-central1"
+    vertex_model: str = "gemini-2.5-flash-001"
+    max_tokens: int = 2000
+
+    class Config:
+        env_file = ".env"
+        extra = "ignore"
+
+
+settings = Settings()

aiscan/models.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class Severity(str, Enum):
+    CRITICAL = "critical"
+    HIGH = "high"
+    MEDIUM = "medium"
+    INFO = "info"
+
+
+class ScanType(str, Enum):
+    PROMPT = "prompt"
+    AGENT = "agent"
+    API = "api"
+    MCP = "mcp"
+    INTEGRATION = "integration"
+
+
+class ThreatType(str, Enum):
+    PROMPT_INJECTION = "prompt_injection"
+    DATA_LEAKAGE = "data_leakage"
+    TOOL_ABUSE = "tool_abuse"
+    SECURITY_RISK = "security_risk"
+
+
+@dataclass
+class Finding:
+    scan_type: ScanType
+    threat_type: ThreatType
+    severity: Severity
+    rule_id: str
+    title: str
+    detail: str
+    file: str
+    line: int
+    fix: str
+
+
+@dataclass
+class ScanResult:
+    target: str
+    scan_type: ScanType
+    findings: list[Finding] = field(default_factory=list)
+    error: str = ""
+
+    @property
+    def critical_count(self): return sum(1 for f in self.findings if f.severity == Severity.CRITICAL)
+
+    @property
+    def high_count(self):     return sum(1 for f in self.findings if f.severity == Severity.HIGH)
+
+    @property
+    def passed(self):         return self.critical_count == 0 and self.high_count == 0

aiscan/reporters/console_reporter.py

@@ -0,0 +1,49 @@
+"""Rich console reporter for all 5 scan types."""
+from rich import box
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+from aiscan.models import ScanResult, Severity
+
+console = Console()
+
+ICON = {Severity.CRITICAL: "🔴", Severity.HIGH: "🟠", Severity.MEDIUM: "🟡", Severity.INFO: "🔵"}
+STYLE = {Severity.CRITICAL: "bold red", Severity.HIGH: "bold yellow", Severity.MEDIUM: "yellow", Severity.INFO: "dim"}
+
+
+def print_results(results: list[ScanResult]) -> int:
+    total = 0
+    for r in results:
+        if r.error:
+            console.print(f"[red]ERROR {r.target}: {r.error}[/red]")
+            continue
+        if not r.findings:
+            console.print(f"[green]✅  {r.target} ({r.scan_type.value}) — no issues[/green]")
+            continue
+        t = Table(box=box.SIMPLE, show_header=True, header_style="bold")
+        t.add_column("Sev", width=10)
+        t.add_column("Rule", width=8)
+        t.add_column("Threat", width=20)
+        t.add_column("Title", width=44)
+        t.add_column("Line", width=5, justify="right")
+        for f in r.findings:
+            t.add_row(
+                f"{ICON.get(f.severity, '')} {f.severity.value}",
+                f.rule_id, f.threat_type.value, f.title,
+                str(f.line) if f.line else "—",
+                style=STYLE.get(f.severity, ""),
+            )
+        console.print(Panel(t, title=f"[bold]{r.target}[/bold] [{r.scan_type.value}]",
+                            subtitle=f"{len(r.findings)} finding(s)"))
+        for f in r.findings:
+            if f.severity in (Severity.CRITICAL, Severity.HIGH):
+                console.print(f"  [bold]{f.title}[/bold]")
+                console.print(f"  [dim]{f.detail}[/dim]")
+                console.print(f"  [green]Fix:[/green] {f.fix}\n")
+        total += len(r.findings)
+    if total == 0:
+        console.print("\n[bold green]✅  All scans passed.[/bold green]")
+    else:
+        console.print(f"\n[bold]Total findings: {total}[/bold]")
+    return total