codeguard_tool.py

"""
title: CodeGuard
author: tkalevra
author_url: https://github.com/tkalevra
version: 0.1.0
license: MIT
description: Static analysis and edge case review for code in Open WebUI.
  Runs shellcheck (bash), ruff (python), or falls back to LLM analysis.
  Designed as a companion to SuperPowersWUI or standalone use.
"""

# MIT License
#
# Copyright (c) 2026 tkalevra
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import json
import os
import re
import shutil
import subprocess
import tempfile
from typing import Any, Callable, Optional

from pydantic import BaseModel, Field


class Tools:
    class Valves(BaseModel):
        ENABLE_SHELLCHECK: bool = Field(
            default=True,
            description="Enable shellcheck for bash/sh analysis.",
        )
        ENABLE_RUFF: bool = Field(
            default=True,
            description="Enable ruff for Python analysis.",
        )
        ENABLE_LLM_FALLBACK: bool = Field(
            default=True,
            description=(
                "If linter not available, fall back to LLM-based analysis via "
                "Open WebUI's generate_chat_completion. If False and linter "
                "missing, return a clear error instead."
            ),
        )
        SHELLCHECK_SEVERITY: str = Field(
            default="style",
            description=(
                "Minimum severity to report. Values: error, warning, info, style. "
                "Passed as --severity to shellcheck."
            ),
        )
        ENABLE_EDGE_CASE_REVIEW: bool = Field(
            default=True,
            description=(
                "After linter output, run a second LLM sub-agent pass focused "
                "specifically on edge cases, race conditions, silent failures, "
                "and logic gaps not caught by static analysis."
            ),
        )

    def __init__(self):
        self.valves = self.Valves()

    async def analyze_code(
        self,
        code: str,
        language: str,
        __user__: Optional[dict] = None,
        __metadata__: Optional[dict] = None,
        __model__: Optional[dict] = None,
        __request__: Optional[Any] = None,
        __event_emitter__: Optional[Callable] = None,
        __event_call__: Optional[Callable] = None,
        __chat_id__: Optional[str] = None,
        __message_id__: Optional[str] = None,
    ) -> str:
        """
        Analyze code for bugs, edge cases, and quality issues.
        Runs shellcheck for bash/sh, ruff for Python. Falls back
        to LLM analysis if linters are unavailable. Pass language
        as 'auto' to detect automatically. Returns structured
        findings with severity, line numbers where available, and
        actionable fix suggestions. Call this on any code block
        before approving a spec or plan.
        """

        async def emit(message: str, done: bool = False):
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {"description": message, "done": done},
                    }
                )

        # Try to import OWUI internals once; flag availability
        llm_available = False
        generate_chat_completion = None
        try:
            from open_webui.utils.chat import (
                generate_chat_completion as _gcc,
            )

            generate_chat_completion = _gcc
            llm_available = True
        except ImportError:
            pass

        model_id = __model__.get("id", "") if isinstance(__model__, dict) else (__model__ or "")

        # ── Step 1: Language detection ────────────────────────────────────
        await emit("Detecting language...")

        if language == "auto":
            language = _detect_language(code)

        # ── Step 2: Linter execution ──────────────────────────────────────
        linter_used = "none"
        linter_available = False
        findings: list[dict] | str = []

        if language in ("bash", "sh"):
            if self.valves.ENABLE_SHELLCHECK and shutil.which("shellcheck"):
                await emit("Running shellcheck...")
                linter_available = True
                linter_used = "shellcheck"
                findings = _run_shellcheck(
                    code, language, self.valves.SHELLCHECK_SEVERITY
                )
            else:
                linter_available = False

        elif language == "python":
            if self.valves.ENABLE_RUFF and shutil.which("ruff"):
                await emit("Running ruff...")
                linter_available = True
                linter_used = "ruff"
                findings = _run_ruff(code)
            else:
                linter_available = False

        # ── Step 3: LLM fallback ──────────────────────────────────────────
        if not linter_available:
            if self.valves.ENABLE_LLM_FALLBACK and llm_available:
                await emit("LLM fallback analysis...")
                linter_used = "llm-fallback"
                findings = await _llm_analyze(
                    code,
                    language,
                    generate_chat_completion,
                    __request__,
                    __user__,
                    model_id,
                )
            else:
                linter_used = "none"
                tool_name = (
                    "shellcheck" if language in ("bash", "sh") else "ruff"
                )
                findings = [
                    {
                        "line": None,
                        "severity": "ERROR",
                        "code": "CG001",
                        "message": (
                            f"Linter '{tool_name}' is not installed and "
                            "LLM fallback is disabled. Install the linter "
                            "or enable ENABLE_LLM_FALLBACK."
                        ),
                        "fix": None,
                    }
                ]

        # ── Step 4: Edge case review ──────────────────────────────────────
        edge_findings = ""
        if self.valves.ENABLE_EDGE_CASE_REVIEW:
            await emit("Edge case review...")
            if llm_available and generate_chat_completion:
                edge_findings = await _llm_edge_cases(
                    code,
                    language,
                    findings,
                    generate_chat_completion,
                    __request__,
                    __user__,
                    model_id,
                )
            else:
                edge_findings = (
                    "_LLM unavailable — edge case review skipped._"
                )

        # ── Step 5: Build result ──────────────────────────────────────────
        await emit("Analysis complete.", done=True)

        result = _format_result(
            language=language,
            linter_used=linter_used,
            findings=findings,
            edge_findings=edge_findings,
            edge_enabled=self.valves.ENABLE_EDGE_CASE_REVIEW,
        )
        return result


# ── Language detection ────────────────────────────────────────────────────────

def _detect_language(code: str) -> str:
    first_line = code.splitlines()[0].strip() if code.strip() else ""

    shebang_map = {
        "bash": ["#!/bin/bash", "#!/usr/bin/env bash"],
        "sh": ["#!/bin/sh", "#!/usr/bin/env sh"],
        "python": [
            "#!/usr/bin/env python",
            "#!/usr/bin/python",
            "#!/usr/bin/env python3",
            "#!/usr/bin/python3",
        ],
    }
    for lang, shebangs in shebang_map.items():
        if any(first_line.startswith(s) for s in shebangs):
            return lang

    # Keyword heuristics
    python_keywords = {"def ", "import ", "class ", "elif ", "print(", "from "}
    bash_keywords = {"function ", " fi", "esac", "local ", "then", "do\n", "done"}

    py_hits = sum(1 for kw in python_keywords if kw in code)
    sh_hits = sum(1 for kw in bash_keywords if kw in code)

    if py_hits > sh_hits:
        return "python"
    if sh_hits > py_hits:
        return "bash"

    return "unknown"


# ── Linter runners ────────────────────────────────────────────────────────────

def _run_shellcheck(code: str, language: str, severity: str) -> list[dict]:
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(
            suffix=".sh", delete=False, mode="w", encoding="utf-8"
        ) as tmp:
            tmp.write(code)
            tmp_path = tmp.name

        result = subprocess.run(
            [
                "shellcheck",
                "--format=json",
                f"--severity={severity}",
                f"--shell={language}",
                tmp_path,
            ],
            capture_output=True,
            text=True,
            timeout=30,
        )

        raw = result.stdout.strip()
        if not raw:
            return []

        items = json.loads(raw)
        findings = []
        for item in items:
            fix = None
            if item.get("fix") and item["fix"].get("replacements"):
                fix = item["fix"]["replacements"]
            findings.append(
                {
                    "line": item.get("line"),
                    "severity": item.get("level", "").upper(),
                    "code": f"SC{item.get('code', '')}",
                    "message": item.get("message", ""),
                    "fix": fix,
                }
            )
        return findings

    except subprocess.TimeoutExpired:
        return [
            {
                "line": None,
                "severity": "ERROR",
                "code": "CG002",
                "message": "shellcheck timed out after 30 seconds.",
                "fix": None,
            }
        ]
    except Exception as exc:
        return [
            {
                "line": None,
                "severity": "ERROR",
                "code": "CG003",
                "message": f"shellcheck error: {exc}",
                "fix": None,
            }
        ]
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


def _run_ruff(code: str) -> list[dict]:
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(
            suffix=".py", delete=False, mode="w", encoding="utf-8"
        ) as tmp:
            tmp.write(code)
            tmp_path = tmp.name

        result = subprocess.run(
            ["ruff", "check", "--output-format=json", tmp_path],
            capture_output=True,
            text=True,
            timeout=30,
        )

        raw = result.stdout.strip()
        if not raw:
            return []

        items = json.loads(raw)
        findings = []
        for item in items:
            fix = None
            if item.get("fix") and item["fix"].get("edits"):
                fix = item["fix"]["edits"]
            findings.append(
                {
                    "line": item.get("location", {}).get("row"),
                    "severity": "WARNING",
                    "code": item.get("code", ""),
                    "message": item.get("message", ""),
                    "fix": fix,
                }
            )
        return findings

    except subprocess.TimeoutExpired:
        return [
            {
                "line": None,
                "severity": "ERROR",
                "code": "CG002",
                "message": "ruff timed out after 30 seconds.",
                "fix": None,
            }
        ]
    except Exception as exc:
        return [
            {
                "line": None,
                "severity": "ERROR",
                "code": "CG003",
                "message": f"ruff error: {exc}",
                "fix": None,
            }
        ]
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


# ── LLM helpers ───────────────────────────────────────────────────────────────

async def _llm_analyze(
    code: str,
    language: str,
    generate_chat_completion,
    request,
    user,
    model_id: str,
) -> str:
    try:
        from open_webui.models.users import UserModel
        user_obj = UserModel(**user) if isinstance(user, dict) else user
        system_prompt = (
            f"You are a senior code reviewer. Analyze the following "
            f"{language} code for bugs, anti-patterns, security issues, "
            f"and correctness problems. Be specific. Reference line "
            f"numbers where possible. Format findings as a numbered "
            f"list with severity (ERROR/WARNING/INFO) prefixed."
        )
        payload = {
            "model": model_id,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"```{language}\n{code}\n```"},
            ],
            "stream": False,
        }
        response = await generate_chat_completion(
            request=request, form_data=payload, user=user_obj
        )
        return response["choices"][0]["message"]["content"]
    except Exception as exc:
        return f"_LLM fallback error: {exc}_"


async def _llm_edge_cases(
    code: str,
    language: str,
    findings,
    generate_chat_completion,
    request,
    user,
    model_id: str,
) -> str:
    try:
        from open_webui.models.users import UserModel
        user_obj = UserModel(**user) if isinstance(user, dict) else user
        findings_text = _findings_to_text(findings)
        system_prompt = (
            f"You are an adversarial code reviewer focused exclusively "
            f"on edge cases. Given this {language} code and the "
            f"following linter findings, identify: race conditions, "
            f"silent failures, missing error handling, resource leaks, "
            f"incorrect assumptions about input, and logic gaps that "
            f"a linter would miss. Be specific and adversarial. "
            f"Do not repeat findings already listed. Format as a "
            f"numbered list with severity prefixed."
        )
        user_content = (
            f"**Code:**\n```{language}\n{code}\n```\n\n"
            f"**Linter findings:**\n{findings_text}"
        )
        payload = {
            "model": model_id,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_content},
            ],
            "stream": False,
        }
        response = await generate_chat_completion(
            request=request, form_data=payload, user=user_obj
        )
        return response["choices"][0]["message"]["content"]
    except Exception as exc:
        return f"_Edge case review error: {exc}_"


# ── Formatting ────────────────────────────────────────────────────────────────

def _findings_to_text(findings) -> str:
    if isinstance(findings, str):
        return findings
    if not findings:
        return "No issues found."

    lines = []
    for i, f in enumerate(findings, 1):
        loc = f"L{f['line']} " if f.get("line") else ""
        code = f"[{f['code']}] " if f.get("code") else ""
        fix_note = ""
        if f.get("fix"):
            fix_note = " *(auto-fix available)*"
        lines.append(
            f"{i}. **{f['severity']}** {loc}{code}{f['message']}{fix_note}"
        )
    return "\n".join(lines)


def _format_result(
    language: str,
    linter_used: str,
    findings,
    edge_findings: str,
    edge_enabled: bool,
) -> str:
    # Static analysis section
    static_section = _findings_to_text(findings)

    # Count errors across structured or string findings
    error_count = 0
    static_count = 0
    if isinstance(findings, list):
        static_count = len(findings)
        error_count = sum(
            1 for f in findings if f.get("severity") == "ERROR"
        )
    else:
        static_count = findings.count("ERROR:")

    # Edge case section
    if edge_enabled:
        edge_section = edge_findings if edge_findings else "No additional findings."
        edge_line_count = (
            len(re.findall(r"^\s*\d+\.", edge_findings, re.MULTILINE))
            if edge_findings
            else 0
        )
    else:
        edge_section = "Edge case review disabled."
        edge_line_count = 0

    verdict = "BLOCKED" if error_count > 0 else "APPROVED"

    result = (
        f"## CodeGuard Analysis\n"
        f"**Language:** {language}  \n"
        f"**Linter:** {linter_used}\n\n"
        f"### Static Analysis Findings\n"
        f"{static_section}\n\n"
        f"### Edge Case Review\n"
        f"{edge_section}\n\n"
        f"### Summary\n"
        f"{static_count} static finding(s). "
        f"{edge_line_count} edge case finding(s).\n"
        f"**Verdict: {verdict}**\n\n"
        f"[SUPERPOWERS:AUTO-CONTINUE]"
    )
    return result