client fix

app.py

@@ -9,11 +9,31 @@ load_dotenv()
 # MCP Server URL for GitHub tools
 MCP_SERVER_URL = "https://himanshugoyal2004-github-mcp-server.hf.space/gradio_api/mcp/"
 
+def parse_github_url(url):
+    """Parse GitHub URL to extract owner, repo, and file path"""
+    import re
+    
+    # Handle repository URLs
+    repo_pattern = r'https://github\.com/([^/]+)/([^/]+)/?$'
+    repo_match = re.match(repo_pattern, url.strip())
+    if repo_match:
+        return repo_match.group(1), repo_match.group(2), None
+    
+    # Handle file URLs
+    file_pattern = r'https://github\.com/([^/]+)/([^/]+)/blob/[^/]+/(.+)$'
+    file_match = re.match(file_pattern, url.strip())
+    if file_match:
+        return file_match.group(1), file_match.group(2), file_match.group(3)
+    
+    return None, None, None
+
 def analyze_vulnerabilities(message, history):
-    """Analyze GitHub repository for vulnerabilities using AI agent"""
+    """Analyze GitHub repository or specific file for vulnerabilities using AI agent"""
     try:
-        # Connect to MCP server and get GitHub tools
-        mcp_client = MCPClient({"url": MCP_SERVER_URL, "timeout": 120})
+        mcp_client = MCPClient({
+            "url": MCP_SERVER_URL,
+            "timeout": 120
+        })
         tools = mcp_client.get_tools()
         
         # Initialize AI model
@@ -23,11 +43,47 @@ def analyze_vulnerabilities(message, history):
         agent = CodeAgent(
             tools=[*tools], 
             model=model, 
-            additional_authorized_imports=["json", "ast", "urllib", "base64", "re"]
+            additional_authorized_imports=["json", "ast", "urllib", "base64", "re"],
+            max_steps=10
         )
         
-        # Enhanced prompt for vulnerability analysis
-        enhanced_prompt = f"""
+        # Parse the GitHub URL
+        owner, repo, file_path = parse_github_url(message)
+        
+        if not owner or not repo:
+            return "❌ Invalid GitHub URL. Please provide a valid GitHub repository or file URL."
+        
+        # Generate different prompts based on whether it's a file or repository
+        if file_path:
+            # Single file analysis
+            enhanced_prompt = f"""
+You are a cybersecurity expert. Analyze the specific GitHub file for security vulnerabilities.
+
+GitHub URL: {message}
+Repository: {owner}/{repo}
+File Path: {file_path}
+
+Please:
+1. First, get repository information to verify it exists
+2. Get the content of the specific file: {file_path}
+3. Analyze the file content line by line for security vulnerabilities
+4. Look for these security issues:
+   - Command injection: os.system, exec, eval calls
+   - Input validation: unvalidated user inputs
+   - Error handling: unhandled exceptions that could leak info
+   - Hardcoded secrets: API keys, passwords, tokens
+   - Unsafe operations: file operations without validation
+
+5. Create a professional security report with:
+   - 🔍 File Overview (path, language, size)
+   - 📊 Vulnerability Summary (counts by severity)
+   - 🚨 Detailed Findings (line numbers, code snippets, impacts, fixes)
+
+Use simple string operations and avoid complex regex patterns. Focus on clear, actionable security findings.
+"""
+        else:
+            # Full repository analysis
+            enhanced_prompt = f"""
 You are a cybersecurity expert. Analyze the GitHub repository for security vulnerabilities.
 
 Repository: {message}
@@ -35,32 +91,21 @@ Repository: {message}
 Please:
 1. First, get repository information to verify it exists
 2. Scan the repository for code files (.py, .js, .ts, .php, .java, .cpp, .c, .cs, .go, .rb, .rs, .swift, .kt, .scala, .sh, .bash, .ps1, .ipynb, .sql, .xml, .yaml, .yml, .json, .config, .ini, .env)
-3. For each code file found, get its content and analyze for security vulnerabilities
-4. Focus on detecting:
-   - SQL injection vulnerabilities
-   - Command injection (os.system, exec, eval)
-   - Cross-site scripting (XSS)
-   - Path traversal attacks
-   - Hardcoded secrets/credentials
-   - Insecure deserialization
-   - Weak cryptography
-   - Authentication/authorization flaws
-   - Input validation issues
-   - Unsafe file operations
+3. For the first 5-10 most important code files, get their content and analyze for security issues
+4. Look for these security vulnerabilities:
+   - Command injection: os.system, exec, eval calls
+   - Input validation: unvalidated user inputs, missing parameter checks
+   - Error handling: unhandled exceptions, information disclosure
+   - Hardcoded secrets: API keys, passwords, database credentials
+   - Unsafe operations: file operations, deserialization without validation
 
 5. Generate a comprehensive security report with:
-   - Repository overview
-   - Total files analyzed
-   - Vulnerability count by severity (Critical/High/Medium/Low)
-   - Detailed findings with:
-     - File path and line number
-     - Vulnerability type
-     - Code snippet
-     - Security impact
-     - Remediation advice
-     - Related CVE examples when applicable
+   - 🔍 Repository Overview
+   - 📁 Files Analyzed
+   - 📊 Vulnerability Summary (counts by severity)
+   - 🚨 Detailed Findings (file paths, line numbers, code snippets, impacts, remediation)
 
-Format the report professionally with emojis and clear sections.
+Use simple string operations and focus on the most critical security issues. Limit analysis to prevent timeouts.
 """
         
         # Run the AI agent analysis
@@ -79,14 +124,10 @@ demo = gr.ChatInterface(
     fn=analyze_vulnerabilities,
     type="messages",
     examples=[
-        "https://github.com/WebGoat/WebGoat",
-        "https://github.com/OWASP/NodeGoat", 
-        "https://github.com/digininja/DVWA",
-        "https://github.com/juice-shop/juice-shop",
-        "https://github.com/vulhub/vulhub"
+        "https://github.com/banno-0720/documentation-agent/blob/main/code.py"       
     ],
     title="🛡️ AI-Powered GitHub Vulnerability Scanner",
-    description="Paste a GitHub repository URL to scan for security vulnerabilities using AI agents with MCP tools. The AI will intelligently analyze code and provide detailed security reports.",
+    description="Paste a GitHub repository URL to scan the entire repo, or paste a specific file URL to analyze just that file for security vulnerabilities using AI agents with MCP tools. The AI will intelligently analyze code and provide detailed security reports.",
 )
 
 if __name__ == "__main__":