#!/usr/bin/env python

"""
Copyright (c) 2006-2024 sqlmap developers (https://sqlmap.org/)
See the file 'LICENSE' for copying permission
"""

# 导入需要的模块
from __future__ import print_function

import difflib  # 用于比较文本差异
import sqlite3  # SQLite数据库支持
import threading  # 多线程支持
import time  # 时间相关功能
import traceback  # 异常追踪

# 导入自定义模块
from lib.core.compat import WichmannHill  # 随机数生成器
from lib.core.compat import xrange  # 兼容Python2和3的range函数
from lib.core.data import conf  # 配置数据
from lib.core.data import kb  # 知识库数据
from lib.core.data import logger  # 日志记录器
from lib.core.datatype import AttribDict  # 属性字典类型
from lib.core.enums import PAYLOAD  # 载荷类型枚举
from lib.core.exception import SqlmapBaseException  # 基础异常类
from lib.core.exception import SqlmapConnectionException  # 连接异常
from lib.core.exception import SqlmapSkipTargetException  # 跳过目标异常
from lib.core.exception import SqlmapThreadException  # 线程异常
from lib.core.exception import SqlmapUserQuitException  # 用户退出异常
from lib.core.exception import SqlmapValueException  # 值错误异常
from lib.core.settings import MAX_NUMBER_OF_THREADS  # 最大线程数
from lib.core.settings import PYVERSION  # Python版本信息

# 创建共享数据对象
shared = AttribDict()

class _ThreadData(threading.local):
    """
    表示线程独立的数据
    每个线程都有自己独立的数据副本
    """

    def __init__(self):
        self.reset()

    def reset(self):
        """
        重置线程数据模型
        初始化所有线程局部变量
        """

        self.disableStdOut = False  # 是否禁用标准输出
        self.hashDBCursor = None  # 哈希数据库游标
        self.inTransaction = False  # 是否在事务中
        self.lastCode = None  # 最后的HTTP状态码
        self.lastComparisonPage = None  # 最后比较的页面内容
        self.lastComparisonHeaders = None  # 最后比较的HTTP头
        self.lastComparisonCode = None  # 最后比较的状态码
        self.lastComparisonRatio = None  # 最后比较的相似度
        self.lastErrorPage = tuple()  # 最后的错误页面
        self.lastHTTPError = None  # 最后的HTTP错误
        self.lastRedirectMsg = None  # 最后的重定向消息
        self.lastQueryDuration = 0  # 最后查询持续时间
        self.lastPage = None  # 最后的页面内容
        self.lastRequestMsg = None  # 最后的请求消息
        self.lastRequestUID = 0  # 最后请求的唯一ID
        self.lastRedirectURL = tuple()  # 最后重定向的URL
        self.random = WichmannHill()  # 随机数生成器
        self.resumed = False  # 是否已恢复
        self.retriesCount = 0  # 重试次数
        self.seqMatcher = difflib.SequenceMatcher(None)  # 序列匹配器
        self.shared = shared  # 共享数据引用
        self.technique = None  # 当前使用的技术
        self.validationRun = 0  # 验证运行次数
        self.valueStack = []  # 值栈

# 创建线程数据实例
ThreadData = _ThreadData()

def readInput(message, default=None, checkBatch=True, boolean=False):
    # 将被lib.core.common中的原始函数覆盖
    pass

def isDigit(value):
    # 将被lib.core.common中的原始函数覆盖
    pass

def getCurrentThreadData():
    """
    返回当前线程的局部数据
    """
    return ThreadData

def getCurrentThreadName():
    """
    返回当前线程的名称
    """
    return threading.current_thread().getName()

def exceptionHandledFunction(threadFunction, silent=False):
    """
    异常处理包装函数
    用于包装线程函数并处理可能发生的异常
    """
    try:
        threadFunction()
    except KeyboardInterrupt:  # 处理键盘中断
        kb.threadContinue = False
        kb.threadException = True
        raise
    except Exception as ex:
        from lib.core.common import getSafeExString

        # 如果不是静默模式且线程应继续运行,且不是多次Ctrl+C,且不是用户退出或跳过目标异常
        if not silent and kb.get("threadContinue") and not kb.get("multipleCtrlC") and not isinstance(ex, (SqlmapUserQuitException, SqlmapSkipTargetException)):
            errMsg = getSafeExString(ex) if isinstance(ex, SqlmapBaseException) else "%s: %s" % (type(ex).__name__, getSafeExString(ex))
            logger.error("thread %s: '%s'" % (threading.currentThread().getName(), errMsg))

            # 在详细模式下打印完整堆栈跟踪
            if conf.get("verbose") > 1 and not isinstance(ex, SqlmapConnectionException):
                traceback.print_exc()

def setDaemon(thread):
    """
    设置线程为守护线程
    守护线程会在主程序退出时自动终止
    """
    if PYVERSION >= "2.6":
        thread.daemon = True
    else:
        thread.setDaemon(True)

def runThreads(numThreads, threadFunction, cleanupFunction=None, forwardException=True, threadChoice=False, startThreadMsg=True):
    """
    运行多个线程的主函数
    
    参数:
    numThreads - 要运行的线程数
    threadFunction - 每个线程要执行的函数
    cleanupFunction - 清理函数(可选)
    forwardException - 是否转发异常
    threadChoice - 是否允许用户选择线程数
    startThreadMsg - 是否显示启动线程消息
    """
    threads = []

    def _threadFunction():
        """
        内部线程函数
        包装了原始的threadFunction,并确保正确关闭hashDB
        """
        try:
            threadFunction()
        finally:
            if conf.hashDB:
                conf.hashDB.close()

    # 初始化线程控制变量
    kb.multipleCtrlC = False
    kb.threadContinue = True
    kb.threadException = False
    kb.technique = ThreadData.technique
    kb.multiThreadMode = False

    try:
        # 处理单线程情况下的线程数选择
        if threadChoice and conf.threads == numThreads == 1 and not (kb.injection.data and not any(_ not in (PAYLOAD.TECHNIQUE.TIME, PAYLOAD.TECHNIQUE.STACKED) for _ in kb.injection.data)):
            while True:
                message = "please enter number of threads? [Enter for %d (current)] " % numThreads
                choice = readInput(message, default=str(numThreads))
                if choice:
                    skipThreadCheck = False

                    if choice.endswith('!'):
                        choice = choice[:-1]
                        skipThreadCheck = True

                    if isDigit(choice):
                        if int(choice) > MAX_NUMBER_OF_THREADS and not skipThreadCheck:
                            errMsg = "maximum number of used threads is %d avoiding potential connection issues" % MAX_NUMBER_OF_THREADS
                            logger.critical(errMsg)
                        else:
                            conf.threads = numThreads = int(choice)
                            break

            if numThreads == 1:
                warnMsg = "running in a single-thread mode. This could take a while"
                logger.warning(warnMsg)

        # 处理多线程和单线程的执行
        if numThreads > 1:
            if startThreadMsg:
                infoMsg = "starting %d threads" % numThreads
                logger.info(infoMsg)
        else:
            try:
                _threadFunction()
            except (SqlmapUserQuitException, SqlmapSkipTargetException):
                pass
            finally:
                return

        kb.multiThreadMode = True

        # 启动所有线程
        for numThread in xrange(numThreads):
            thread = threading.Thread(target=exceptionHandledFunction, name=str(numThread), args=[_threadFunction])

            setDaemon(thread)

            try:
                thread.start()
            except Exception as ex:
                errMsg = "error occurred while starting new thread ('%s')" % ex
                logger.critical(errMsg)
                break

            threads.append(thread)

        # 等待所有线程完成
        alive = True
        while alive:
            alive = False
            for thread in threads:
                if thread.is_alive():
                    alive = True
                    time.sleep(0.1)

    except (KeyboardInterrupt, SqlmapUserQuitException) as ex:
        # 处理用户中断
        print()
        kb.prependFlag = False
        kb.threadContinue = False
        kb.threadException = True

        if kb.lastCtrlCTime and (time.time() - kb.lastCtrlCTime < 1):
            kb.multipleCtrlC = True
            raise SqlmapUserQuitException("user aborted (Ctrl+C was pressed multiple times)")

        kb.lastCtrlCTime = time.time()

        if numThreads > 1:
            logger.info("waiting for threads to finish%s" % (" (Ctrl+C was pressed)" if isinstance(ex, KeyboardInterrupt) else ""))
        try:
            while (threading.active_count() > 1):
                pass

        except KeyboardInterrupt:
            kb.multipleCtrlC = True
            raise SqlmapThreadException("user aborted (Ctrl+C was pressed multiple times)")

        if forwardException:
            raise

    except (SqlmapConnectionException, SqlmapValueException) as ex:
        # 处理连接和值错误异常
        print()
        kb.threadException = True
        logger.error("thread %s: '%s'" % (threading.currentThread().getName(), ex))

        if conf.get("verbose") > 1 and isinstance(ex, SqlmapValueException):
            traceback.print_exc()

    except Exception as ex:
        # 处理其他未预期的异常
        print()

        if not kb.multipleCtrlC:
            if isinstance(ex, sqlite3.Error):
                raise
            else:
                from lib.core.common import unhandledExceptionMessage

                kb.threadException = True
                errMsg = unhandledExceptionMessage()
                logger.error("thread %s: %s" % (threading.currentThread().getName(), errMsg))
                traceback.print_exc()

    finally:
        # 清理工作
        kb.multiThreadMode = False
        kb.threadContinue = True
        kb.threadException = False
        kb.technique = None

        # 释放所有锁
        for lock in kb.locks.values():
            if lock.locked():
                try:
                    lock.release()
                except:
                    pass

        # 刷新哈希数据库
        if conf.get("hashDB"):
            conf.hashDB.flush(True)

        # 执行清理函数
        if cleanupFunction:
            cleanupFunction()
