feat: 提高VulnShop靶场鲁棒性，防止SQLMap扫描时崩溃

lanshuizhiyue · lanshuizhiyue · commit 2800afbe2575 · 2025-12-24T10:57:17.000+08:00
diff --git a/src/vulnTestServer/handlers/base.py b/src/vulnTestServer/handlers/base.py
@@ -37,14 +37,20 @@ class BaseHandlerMixin:
     
     def send_json_response(self, data, status=200):
         """发送JSON响应"""
-        self.send_response(status)
-        self.send_header('Content-Type', 'application/json; charset=utf-8')
-        self.send_header('Access-Control-Allow-Origin', '*')
-        self.send_header('Access-Control-Allow-Methods', 'GET, POST, PUT, DELETE, OPTIONS')
-        self.send_header('Access-Control-Allow-Headers', 'Content-Type, Authorization')
-        self.end_headers()
-        response = json.dumps(data, ensure_ascii=False)
-        self.wfile.write(response.encode('utf-8'))
+        try:
+            self.send_response(status)
+            self.send_header('Content-Type', 'application/json; charset=utf-8')
+            self.send_header('Access-Control-Allow-Origin', '*')
+            self.send_header('Access-Control-Allow-Methods', 'GET, POST, PUT, DELETE, OPTIONS')
+            self.send_header('Access-Control-Allow-Headers', 'Content-Type, Authorization')
+            self.end_headers()
+            response = json.dumps(data, ensure_ascii=False)
+            self.wfile.write(response.encode('utf-8'))
+        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+            # 客户端断开连接，静默处理
+            pass
+        except Exception as e:
+            error_logger.debug("Error sending JSON response: %s", str(e))
     
     def send_error_response(self, message, status=400, sql_error=None):
         """发送错误响应（可能包含SQL错误信息用于演示）"""
@@ -63,37 +69,43 @@ def send_error_response(self, message, status=400, sql_error=None):
     
     def send_xml_response(self, data, status=200):
         """发送XML响应"""
-        self.send_response(status)
-        self.send_header('Content-Type', 'application/xml; charset=utf-8')
-        self.send_header('Access-Control-Allow-Origin', '*')
-        self.send_header('Access-Control-Allow-Methods', 'GET, POST, PUT, DELETE, OPTIONS')
-        self.send_header('Access-Control-Allow-Headers', 'Content-Type, Authorization, X-Session-Id, X-Token')
-        self.end_headers()
-        
-        # 将dict转换为XML
-        def dict_to_xml(d, root_name='response'):
-            root = ET.Element(root_name)
-            for key, value in d.items():
-                child = ET.SubElement(root, key)
-                if isinstance(value, dict):
-                    for k, v in value.items():
-                        sub = ET.SubElement(child, k)
-                        sub.text = str(v) if v is not None else ''
-                elif isinstance(value, list):
-                    for item in value:
-                        item_el = ET.SubElement(child, 'item')
-                        if isinstance(item, dict):
-                            for k, v in item.items():
-                                sub = ET.SubElement(item_el, k)
-                                sub.text = str(v) if v is not None else ''
-                        else:
-                            item_el.text = str(item)
-                else:
-                    child.text = str(value) if value is not None else ''
-            return ET.tostring(root, encoding='unicode')
-        
-        xml_str = '<?xml version="1.0" encoding="UTF-8"?>\n' + dict_to_xml(data)
-        self.wfile.write(xml_str.encode('utf-8'))
+        try:
+            self.send_response(status)
+            self.send_header('Content-Type', 'application/xml; charset=utf-8')
+            self.send_header('Access-Control-Allow-Origin', '*')
+            self.send_header('Access-Control-Allow-Methods', 'GET, POST, PUT, DELETE, OPTIONS')
+            self.send_header('Access-Control-Allow-Headers', 'Content-Type, Authorization, X-Session-Id, X-Token')
+            self.end_headers()
+            
+            # 将dict转换为XML
+            def dict_to_xml(d, root_name='response'):
+                root = ET.Element(root_name)
+                for key, value in d.items():
+                    child = ET.SubElement(root, key)
+                    if isinstance(value, dict):
+                        for k, v in value.items():
+                            sub = ET.SubElement(child, k)
+                            sub.text = str(v) if v is not None else ''
+                    elif isinstance(value, list):
+                        for item in value:
+                            item_el = ET.SubElement(child, 'item')
+                            if isinstance(item, dict):
+                                for k, v in item.items():
+                                    sub = ET.SubElement(item_el, k)
+                                    sub.text = str(v) if v is not None else ''
+                            else:
+                                item_el.text = str(item)
+                    else:
+                        child.text = str(value) if value is not None else ''
+                return ET.tostring(root, encoding='unicode')
+            
+            xml_str = '<?xml version="1.0" encoding="UTF-8"?>\n' + dict_to_xml(data)
+            self.wfile.write(xml_str.encode('utf-8'))
+        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+            # 客户端断开连接，静默处理
+            pass
+        except Exception as e:
+            error_logger.debug("Error sending XML response: %s", str(e))
     
     def send_static_file(self, filepath):
         """发送静态文件"""
@@ -120,42 +132,78 @@ def send_static_file(self, filepath):
             self.send_header('Content-Length', len(content))
             self.end_headers()
             self.wfile.write(content)
+        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+            # 客户端断开连接，静默处理
+            pass
+        except FileNotFoundError:
+            self.send_error(404, 'File Not Found')
         except Exception as e:
-            error_logger.exception("Error serving static file: %s", filepath)
-            self.send_error(500, str(e))
+            error_logger.debug("Error serving static file %s: %s", filepath, str(e))
+            try:
+                self.send_error(500, 'Internal Server Error')
+            except:
+                pass
     
     def get_post_data(self):
         """获取POST请求数据，支持JSON、URL-encoded和XML格式"""
-        content_length = int(self.headers.get('Content-Length', 0))
-        if content_length == 0:
-            return {}
-        
-        post_data = self.rfile.read(content_length).decode('utf-8')
-        content_type = self.headers.get('Content-Type', '')
-        
-        if 'application/json' in content_type:
+        try:
+            content_length_str = self.headers.get('Content-Length', '0')
             try:
-                return json.loads(post_data)
-            except:
+                content_length = int(content_length_str)
+            except (ValueError, TypeError):
+                content_length = 0
+            
+            if content_length <= 0:
                 return {}
-        elif 'application/xml' in content_type or 'text/xml' in content_type:
-            # XML解析
+            
+            # 限制最大读取长度，防止内存攻击
+            max_length = 10 * 1024 * 1024  # 10MB
+            if content_length > max_length:
+                content_length = max_length
+            
+            raw_data = self.rfile.read(content_length)
             try:
-                root = ET.fromstring(post_data)
+                post_data = raw_data.decode('utf-8')
+            except UnicodeDecodeError:
+                # 尝试其他编码
+                try:
+                    post_data = raw_data.decode('latin-1')
+                except:
+                    return {}
+            
+            content_type = self.headers.get('Content-Type', '')
+            
+            if 'application/json' in content_type:
+                try:
+                    return json.loads(post_data)
+                except:
+                    return {}
+            elif 'application/xml' in content_type or 'text/xml' in content_type:
+                # XML解析
+                try:
+                    root = ET.fromstring(post_data)
+                    result = {}
+                    for child in root:
+                        result[child.tag] = child.text or ''
+                    return result
+                except:
+                    return {}
+            else:
+                # application/x-www-form-urlencoded
                 result = {}
-                for child in root:
-                    result[child.tag] = child.text or ''
+                for pair in post_data.split('&'):
+                    if '=' in pair:
+                        key, value = pair.split('=', 1)
+                        try:
+                            result[unquote(key)] = unquote(value)
+                        except:
+                            pass
                 return result
-            except:
-                return {}
-        else:
-            # application/x-www-form-urlencoded
-            result = {}
-            for pair in post_data.split('&'):
-                if '=' in pair:
-                    key, value = pair.split('=', 1)
-                    result[unquote(key)] = unquote(value)
-            return result
+        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+            return {}
+        except Exception as e:
+            error_logger.debug("Error reading POST data: %s", str(e))
+            return {}
     
     def get_content_type(self):
         """获取请求的Content-Type类型"""
diff --git a/src/vulnTestServer/server.py b/src/vulnTestServer/server.py
@@ -24,8 +24,9 @@
 
 import os
 import sys
-import traceback
+import socket
 from http.server import HTTPServer, BaseHTTPRequestHandler
+from socketserver import ThreadingMixIn
 from urllib.parse import urlparse, parse_qs
 
 # 添加当前目录到路径
@@ -60,9 +61,42 @@ class VulnShopHandler(
     使用Mixin模式组合各个功能模块，实现模块化设计
     """
     
+    # 设置请求超时（秒）
+    timeout = 30
+    
     def log_message(self, format, *args):
         """重写日志方法，使用自定义日志器"""
-        access_logger.info("%s - %s", self.client_address[0], format % args)
+        try:
+            access_logger.info("%s - %s", self.client_address[0], format % args)
+        except Exception:
+            pass
+    
+    def handle_one_request(self):
+        """重写请求处理，增加异常保护"""
+        try:
+            super().handle_one_request()
+        except socket.timeout:
+            # 超时不是错误
+            self.close_connection = True
+        except ConnectionResetError:
+            # 客户端重置连接
+            self.close_connection = True
+        except BrokenPipeError:
+            # 管道破裂，客户端断开
+            self.close_connection = True
+        except ConnectionAbortedError:
+            # 连接被中止
+            self.close_connection = True
+        except Exception as e:
+            error_logger.debug("Request handling error: %s", str(e))
+            self.close_connection = True
+    
+    def finish(self):
+        """重写finish方法，处理连接关闭时的异常"""
+        try:
+            super().finish()
+        except Exception:
+            pass
     
     def do_OPTIONS(self):
         """处理CORS预检请求"""
@@ -109,10 +143,22 @@ def do_GET(self):
                 self.send_error(404, 'Not Found')
         except WAFBlockedException as e:
             access_logger.warning("%s - WAF Blocked: %s", self.client_address[0], e.reason)
-            self.send_error_response(f'WAF Blocked: {e.reason}', 403)
+            self._safe_send_error('WAF Blocked: ' + e.reason, 403)
+        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+            # 客户端断开连接，无需处理
+            pass
         except Exception as e:
             error_logger.exception("Error processing GET %s: %s", path, str(e))
-            self.send_error_response(str(e), 500, sql_error=e)
+            self._safe_send_error(str(e), 500, sql_error=e)
+    
+    def _safe_send_error(self, message, status=400, sql_error=None):
+        """安全发送错误响应，处理连接断开的情况"""
+        try:
+            self.send_error_response(message, status, sql_error)
+        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+            pass
+        except Exception:
+            pass
     
     def do_POST(self):
         """处理POST请求"""
@@ -145,10 +191,19 @@ def do_POST(self):
                 self.send_error(404, 'Not Found')
         except WAFBlockedException as e:
             access_logger.warning("%s - WAF Blocked: %s", self.client_address[0], e.reason)
-            self.send_error_response(f'WAF Blocked: {e.reason}', 403)
+            self._safe_send_error('WAF Blocked: ' + e.reason, 403)
+        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+            # 客户端断开连接，无需处理
+            pass
         except Exception as e:
             error_logger.exception("Error processing POST %s: %s", path, str(e))
-            self.send_error_response(str(e), 500, sql_error=e)
+            self._safe_send_error(str(e), 500, sql_error=e)
+
+
+class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
+    """支持多线程的HTTP服务器"""
+    daemon_threads = True  # 守护线程，主线程退出时自动结束
+    allow_reuse_address = True  # 允许端口复用
 
 
 def run_server():
@@ -157,9 +212,10 @@ def run_server():
     init_database()
     logger.info("Database initialized")
     
-    # 创建服务器
-    server = HTTPServer((HOST, PORT), VulnShopHandler)
-    logger.info("Server created on %s:%d", HOST, PORT)
+    # 创建多线程服务器
+    server = ThreadedHTTPServer((HOST, PORT), VulnShopHandler)
+    server.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    logger.info("Threaded server created on %s:%d", HOST, PORT)
     
     print(f"""
 ╔═══════════════════════════════════════════════════════════════════════════════════════╗
