子进程主动上传日志

Author: EITS403

lib/utils/api.py

@@ -11,6 +11,7 @@
 import contextlib
 import logging
 import os
+import queue
 import re
 import shlex
 import socket
@@ -174,14 +175,22 @@ def engine_start(self):
         os.close(handle)
         saveConfig(self.options, configFile)
 
+        # 准备子进程环境变量，传递外部日志 URL（如果配置了）
+        env = os.environ.copy()
+        if hasattr(self.options, 'externalLogUrl') and self.options.externalLogUrl:
+            env['SQLMAP_EXTERNAL_LOG_URL'] = self.options.externalLogUrl
+        elif 'SQLMAP_EXTERNAL_LOG_URL' in os.environ:
+            # 如果父进程设置了环境变量，也传递给子进程
+            pass  # 已经在 env.copy() 中包含了
+
         if os.path.exists("sqlmap.py"):
-            self.process = Popen([sys.executable or "python", "sqlmap.py", "--api", "-c", configFile], shell=False, close_fds=not IS_WIN)
+            self.process = Popen([sys.executable or "python", "sqlmap.py", "--api", "-c", configFile], shell=False, close_fds=not IS_WIN, env=env)
         elif os.path.exists(os.path.join(os.getcwd(), "sqlmap.py")):
-            self.process = Popen([sys.executable or "python", "sqlmap.py", "--api", "-c", configFile], shell=False, cwd=os.getcwd(), close_fds=not IS_WIN)
+            self.process = Popen([sys.executable or "python", "sqlmap.py", "--api", "-c", configFile], shell=False, cwd=os.getcwd(), close_fds=not IS_WIN, env=env)
         elif os.path.exists(os.path.join(os.path.abspath(os.path.dirname(sys.argv[0])), "sqlmap.py")):
-            self.process = Popen([sys.executable or "python", "sqlmap.py", "--api", "-c", configFile], shell=False, cwd=os.path.join(os.path.abspath(os.path.dirname(sys.argv[0]))), close_fds=not IS_WIN)
+            self.process = Popen([sys.executable or "python", "sqlmap.py", "--api", "-c", configFile], shell=False, cwd=os.path.join(os.path.abspath(os.path.dirname(sys.argv[0]))), close_fds=not IS_WIN, env=env)
         else:
-            self.process = Popen(["sqlmap", "--api", "-c", configFile], shell=False, close_fds=not IS_WIN)
+            self.process = Popen(["sqlmap", "--api", "-c", configFile], shell=False, close_fds=not IS_WIN, env=env)
 
     def engine_stop(self):
         if self.process:
@@ -271,12 +280,97 @@ def seek(self):
         pass
 
 class LogRecorder(logging.StreamHandler):
+    _log_queue = None
+    _log_thread = None
+    _log_thread_lock = threading.Lock()
+    _external_log_url = None
+
+    def __init__(self):
+        super(LogRecorder, self).__init__()
+        # 从环境变量或配置中获取外部日志接口 URL
+        external_log_url = os.environ.get("SQLMAP_EXTERNAL_LOG_URL")
+        
+        # 如果配置对象中有 externalLogUrl 选项，优先使用
+        if hasattr(conf, 'externalLogUrl') and conf.externalLogUrl:
+            external_log_url = conf.externalLogUrl
+        
+        # 如果设置了外部日志 URL，初始化日志发送队列和线程
+        if external_log_url:
+            with LogRecorder._log_thread_lock:
+                if LogRecorder._log_queue is None:
+                    LogRecorder._external_log_url = external_log_url
+                    LogRecorder._log_queue = queue.Queue(maxsize=1000)  # 设置队列大小，避免内存溢出
+                    LogRecorder._log_thread = threading.Thread(target=LogRecorder._log_sender_worker, daemon=True)
+                    LogRecorder._log_thread.start()
+
+    @staticmethod
+    def _log_sender_worker():
+        """
+        后台工作线程，从队列中取出日志并发送到外部接口
+        """
+        while True:
+            try:
+                log_data = LogRecorder._log_queue.get(timeout=1)
+                if log_data is None:  # 退出信号
+                    break
+                
+                taskid, log_time, level, message = log_data
+                LogRecorder._send_log_to_external(taskid, log_time, level, message)
+                LogRecorder._log_queue.task_done()
+            except queue.Empty:
+                continue
+            except Exception:
+                # 静默处理错误，避免影响主程序
+                pass
+
+    @staticmethod
+    def _send_log_to_external(taskid, log_time, level, message):
+        """
+        发送日志到外部接口
+        """
+        if not LogRecorder._external_log_url:
+            return
+        
+        try:
+            log_payload = {
+                "taskid": taskid,
+                "time": log_time,
+                "level": level,
+                "message": message
+            }
+            
+            data = getBytes(jsonize(log_payload))
+            headers = {"Content-Type": "application/json"}
+            
+            req = _urllib.request.Request(LogRecorder._external_log_url, data, headers)
+            
+            # 设置超时，避免阻塞
+            response = _urllib.request.urlopen(req, timeout=5)
+            response.read()  # 读取响应，确保请求完成
+        except Exception:
+            # 静默处理所有异常，确保外部接口失败不影响日志记录
+            pass
+
     def emit(self, record):
         """
         Record emitted events to IPC database for asynchronous I/O
         communication with the parent process
+        Also forward logs to external interface if configured
         """
-        conf.databaseCursor.execute("INSERT INTO logs VALUES(NULL, ?, ?, ?, ?)", (conf.taskid, time.strftime("%X"), record.levelname, str(record.msg % record.args if record.args else record.msg)))
+        log_time = time.strftime("%X")
+        log_level = record.levelname
+        log_message = str(record.msg % record.args if record.args else record.msg)
+        
+        # 写入数据库
+        conf.databaseCursor.execute("INSERT INTO logs VALUES(NULL, ?, ?, ?, ?)", (conf.taskid, log_time, log_level, log_message))
+        
+        # 如果配置了外部日志接口，将日志加入队列异步发送
+        if LogRecorder._external_log_url and LogRecorder._log_queue is not None:
+            try:
+                LogRecorder._log_queue.put_nowait((conf.taskid, log_time, log_level, log_message))
+            except queue.Full:
+                # 队列满时静默丢弃，避免阻塞
+                pass
 
 def setRestAPILog():
     if conf.api:
@@ -680,7 +774,7 @@ def version(token=None):
     logger.debug("Fetched version (%s)" % ("admin" if is_admin(token) else request.remote_addr))
     return jsonize({"success": True, "version": VERSION_STRING.split('/')[-1]})
 
-def server(host=RESTAPI_DEFAULT_ADDRESS, port=RESTAPI_DEFAULT_PORT, adapter=RESTAPI_DEFAULT_ADAPTER, username=None, password=None, database=None):
+def server(host=RESTAPI_DEFAULT_ADDRESS, port=RESTAPI_DEFAULT_PORT, adapter=RESTAPI_DEFAULT_ADAPTER, username=None, password=None, database=None, log_url=None):
     """
     REST-JSON API server
     """
@@ -695,6 +789,13 @@ def server(host=RESTAPI_DEFAULT_ADDRESS, port=RESTAPI_DEFAULT_PORT, adapter=REST
     else:
         Database.filepath = database
 
+    # 如果提供了 log_url 参数，设置环境变量，这样所有子进程都会继承
+    if log_url:
+        os.environ['SQLMAP_EXTERNAL_LOG_URL'] = log_url
+        logger.info("External log URL configured: %s" % log_url)
+    elif 'SQLMAP_EXTERNAL_LOG_URL' in os.environ:
+        logger.info("External log URL from environment: %s" % os.environ['SQLMAP_EXTERNAL_LOG_URL'])
+
     if port == 0:  # random
         with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
             s.bind((host, 0))