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