UCSB-SURFI/VulnLLM-R

1import gradio as gr2import requests3import os4 5# API configuration6API_BASE_URL = os.getenv("API_BASE_URL", "http://localhost:8000")7API_KEY = os.getenv("API_KEY", "")8 9def analyze_code(code: str, language: str, model: str) -> str:10    """Send code to the vulnerability analysis API and return results."""11    if not code.strip():12        return "Please enter some code to analyze."13 14    api_url = f"{API_BASE_URL}/internal/analyze"15 16    headers = {17        "Content-Type": "application/json",18        "Authorization": f"Bearer {API_KEY}"19    }20 21    payload = {22        "code": code,23        "model": model,24        "language": language25    }26 27    try:28        response = requests.post(api_url, json=payload, headers=headers, timeout=60)29 30        if response.status_code == 403:31            error_detail = response.json().get("detail", "Token limit exceeded.")32            return f"Error: {error_detail}"33 34        if response.status_code != 200:35            error_detail = response.json().get("detail", f"API error: {response.status_code}")36            return f"Error: {error_detail}"37 38        result = response.json()39 40        # Format the result41        status = result.get("result", {}).get("status", "unknown")42        cwe_type = result.get("result", {}).get("cweType", "N/A")43        model_used = result.get("result", {}).get("model", model)44        explanation = result.get("result", {}).get("response", "")45        if "## Final Answer" in explanation:46            explanation = explanation.split("## Final Answer")[0].strip()47        if "nopolicy" in model_used:48            model_used = "VirtueGuard Code"49        if status == "yes":50            output = f"⚠️ **Vulnerability Detected**\n\n"51            output += f"**CWE Type:** {cwe_type}\n"52            output += f"**Model:** {model_used}\n\n"53            output += f"**Analysis:**\n{explanation}"54        else:55            output = f"✅ **No Vulnerability Detected**\n\n"56            output += f"**Model:** {model_used}\n\n"57            output += f"**Analysis:**\n{explanation}"58 59        return output60 61    except requests.exceptions.Timeout:62        return "Error: Request timed out. Please try again."63    except requests.exceptions.ConnectionError:64        return f"Error: Could not connect to API at {API_BASE_URL}"65    except Exception as e:66        return f"Error: {str(e)}"67 68# Language options69LANGUAGES = [70    "python", "javascript", "typescript", "java", "c", "cpp",71    "csharp", "go", "rust", "php", "ruby", "swift", "kotlin"72]73 74# Model options75MODELS = ["virtueguard-code", "claude-4-sonnet", "gpt-4.1"]76 77# Create Gradio interface78with gr.Blocks(title="VulnLLM-R Demo") as demo:79    gr.Markdown("# VulnLLM-R Demo")80    gr.Markdown("Analyze your code for potential security vulnerabilities using VulnLLM-R.")81 82    with gr.Row():83        with gr.Column(scale=2):84            code_input = gr.Code(85                label="Code to Analyze",86                language="python",87                lines=1588            )89        with gr.Column(scale=1):90            language_dropdown = gr.Dropdown(91                choices=LANGUAGES,92                value="python",93                label="Programming Language"94            )95            model_dropdown = gr.Dropdown(96                choices=MODELS,97                value="virtueguard-code",98                label="Model"99            )100            analyze_btn = gr.Button("🔍 Analyze Code", variant="primary")101 102    result_output = gr.Markdown(label="Analysis Result")103 104    analyze_btn.click(105        fn=analyze_code,106        inputs=[code_input, language_dropdown, model_dropdown],107        outputs=result_output108    )109 110if __name__ == "__main__":111    demo.launch()