#!/usr/bin/env python3
"""
Vulnerable Mock Server for GHSA-7ppg-37fh-vcr6
Simulates the vulnerable Milvus port 9091 behavior
"""

import http.server
import socketserver
import json
import urllib.parse
import sys
import socket

PORT = 9091
# Default etcd.rootPath - the weak authentication token
DEFAULT_AUTH_TOKEN = "by-dev"

class ReusableTCPServer(socketserver.TCPServer):
    """TCP server that allows port reuse"""
    allow_reuse_address = True
    allow_reuse_port = True

class VulnerableHandler(http.server.BaseHTTPRequestHandler):
    """Simulates vulnerable Milvus port 9091 behavior"""
    
    def log_message(self, format, *args):
        # Suppress request logging
        pass
    
    def do_GET(self):
        parsed = urllib.parse.urlparse(self.path)
        path = parsed.path
        params = urllib.parse.parse_qs(parsed.query)
        
        # VULNERABILITY 1: /expr endpoint with weak auth
        if path == "/expr":
            self.handle_expr(params)
        # VULNERABILITY 2: REST API without auth
        elif path == "/api/v1/credential/users":
            self.handle_users_unauthenticated()
        elif path == "/healthz":
            self.send_json({"status": "OK"})
        else:
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b"404 page not found")
    
    def do_POST(self):
        parsed = urllib.parse.urlparse(self.path)
        path = parsed.path
        
        # VULNERABILITY 2: REST API without auth
        if path == "/api/v1/credential":
            self.handle_create_user_unauthenticated()
        else:
            self.send_response(404)
            self.end_headers()
    
    def handle_expr(self, params):
        """
        VULNERABILITY: The /expr endpoint accepts a weak 'auth' parameter.
        The default value is 'by-dev' (etcd.rootPath), which is predictable.
        This allows arbitrary Go expression evaluation.
        """
        auth_token = params.get("auth", [""])[0]
        code = params.get("code", [""])[0]
        
        # Check auth against the weak default token
        if auth_token != DEFAULT_AUTH_TOKEN:
            self.send_response(401)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({"error": "unauthorized"}).encode())
            return
        
        # AUTHENTICATED - Execute the expression (simulated)
        # In real Milvus, this evaluates arbitrary Go expressions
        output = self.simulate_expr_execution(code)
        
        self.send_json({
            "output": output,
            "code": code
        })
    
    def simulate_expr_execution(self, code):
        """Simulates what would happen in the real /expr endpoint"""
        if "MinioCfg.SecretAccessKey" in code:
            return "minioadmin123 (simulated secret)"
        elif "GetCredential" in code:
            return "{username: 'root', encrypted_password: '$2a$10$...'}"
        elif "proxy.Stop" in code:
            return "proxy stopping..."
        elif code.startswith("'hello"):
            return "hello from expr"
        else:
            return f"executed: {code}"
    
    def handle_users_unauthenticated(self):
        """
        VULNERABILITY: REST API endpoint returns user list without authentication.
        In patched versions, this would require authentication headers.
        """
        # NO AUTH CHECK - This is the vulnerability
        self.send_json({
            "status": {},
            "usernames": ["root", "admin", "attacker_user"]
        })
    
    def handle_create_user_unauthenticated(self):
        """
        VULNERABILITY: REST API allows user creation without authentication.
        """
        # NO AUTH CHECK - This is the vulnerability
        self.send_json({
            "status": {"code": 0, "message": "User created successfully"},
            "data": {}
        })
    
    def send_json(self, data):
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(json.dumps(data).encode())

if __name__ == "__main__":
    with ReusableTCPServer(("", PORT), VulnerableHandler) as httpd:
        print(f"Vulnerable mock server running on port {PORT}", flush=True)
        httpd.serve_forever()