#!/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.common import openFile  # 用于打开文件
from lib.core.common import randomStr  # 用于生成随机字符串
from lib.core.data import conf  # 配置信息
from lib.core.data import logger  # 日志记录器
from lib.core.enums import REGISTRY_OPERATION  # 注册表操作类型的枚举

class Registry(object):
    """
    这个类定义了读取和写入Windows注册表键值的方法。
    注册表是Windows系统中的一个核心数据库，用于存储系统和应用程序的配置信息。
    """

    def _initVars(self, regKey, regValue, regType=None, regData=None, parse=False):
        """
        初始化注册表操作所需的变量
        :param regKey: 注册表键路径，例如"HKEY_LOCAL_MACHINE\\SOFTWARE\\Microsoft\\Windows"
        :param regValue: 注册表值名称，即要操作的具体配置项名称
        :param regType: 注册表值类型，如REG_SZ(字符串)、REG_DWORD(数字)等
        :param regData: 要写入的注册表值数据
        :param parse: 是否需要解析输出，True表示需要特殊处理输出格式
        """
        # 保存传入的参数
        self._regKey = regKey
        self._regValue = regValue
        self._regType = regType
        self._regData = regData

        # 生成一个随机字符串，用于创建临时批处理文件名
        self._randStr = randomStr(lowercase=True)
        # 设置远程系统上临时批处理文件的路径
        self._batPathRemote = "%s/tmpr%s.bat" % (conf.tmpPath, self._randStr)
        # 设置本地系统上临时批处理文件的路径
        self._batPathLocal = os.path.join(conf.outputPath, "tmpr%s.bat" % self._randStr)

        # 根据parse参数决定读取命令的格式
        if parse:
            # 使用FOR循环解析REG QUERY的输出，可以获得更规范的格式
            readParse = "FOR /F \"tokens=*\" %%A IN ('REG QUERY \"" + self._regKey + "\" /v \"" + self._regValue + "\"') DO SET value=%%A\r\nECHO %value%\r\n"
        else:
            # 直接使用REG QUERY命令查询注册表
            readParse = "REG QUERY \"" + self._regKey + "\" /v \"" + self._regValue + "\""

        # 定义用于读取注册表的批处理命令
        self._batRead = (
            "@ECHO OFF\r\n",  # 关闭命令回显
            readParse,  # 实际的查询命令
        )

        # 定义用于添加注册表值的批处理命令
        self._batAdd = (
            "@ECHO OFF\r\n",
            "REG ADD \"%s\" /v \"%s\" /t %s /d %s /f" % (self._regKey, self._regValue, self._regType, self._regData),  # /f表示强制覆盖不提示
        )

        # 定义用于删除注册表值的批处理命令
        self._batDel = (
            "@ECHO OFF\r\n",
            "REG DELETE \"%s\" /v \"%s\" /f" % (self._regKey, self._regValue),  # /f表示强制删除不提示
        )

    def _createLocalBatchFile(self):
        """
        在本地系统创建临时批处理文件
        这个批处理文件包含了要在远程系统执行的注册表操作命令
        """
        # 以写入模式打开本地临时文件
        self._batPathFp = openFile(self._batPathLocal, "w")

        # 根据操作类型选择相应的命令集
        if self._operation == REGISTRY_OPERATION.READ:
            lines = self._batRead  # 读取操作的命令
        elif self._operation == REGISTRY_OPERATION.ADD:
            lines = self._batAdd  # 添加操作的命令
        elif self._operation == REGISTRY_OPERATION.DELETE:
            lines = self._batDel  # 删除操作的命令

        # 将命令写入批处理文件
        for line in lines:
            self._batPathFp.write(line)

        # 关闭文件
        self._batPathFp.close()

    def _createRemoteBatchFile(self):
        """
        在远程系统创建批处理文件
        这个方法会先在本地创建文件，然后将其复制到远程系统
        """
        # 记录调试信息
        logger.debug("creating batch file '%s'" % self._batPathRemote)

        # 先在本地创建批处理文件
        self._createLocalBatchFile()
        # 将本地文件复制到远程系统
        self.writeFile(self._batPathLocal, self._batPathRemote, "text", forceCheck=True)

        # 删除本地的临时文件
        os.unlink(self._batPathLocal)

    def readRegKey(self, regKey, regValue, parse=False):
        """
        读取注册表键值的方法
        :param regKey: 要读取的注册表键路径
        :param regValue: 要读取的值名称
        :param parse: 是否需要解析输出
        :return: 返回读取到的注册表值内容
        """
        # 设置操作类型为读取
        self._operation = REGISTRY_OPERATION.READ

        # 初始化变量并创建远程批处理文件
        Registry._initVars(self, regKey, regValue, parse=parse)
        self._createRemoteBatchFile()

        # 记录调试信息
        logger.debug("reading registry key '%s' value '%s'" % (regKey, regValue))

        # 在远程系统执行批处理文件并获取输出
        data = self.evalCmd(self._batPathRemote)

        # 如果不需要解析，则处理输出格式（去除前导空格）
        if data and not parse:
            pattern = '    '
            index = data.find(pattern)
            if index != -1:
                data = data[index + len(pattern):]

        # 清理：删除远程临时文件
        self.delRemoteFile(self._batPathRemote)

        return data

    def addRegKey(self, regKey, regValue, regType, regData):
        """
        添加或修改注册表键值
        :param regKey: 要添加的注册表键路径
        :param regValue: 要添加的值名称
        :param regType: 值的类型（如REG_SZ, REG_DWORD等）
        :param regData: 要写入的数据
        """
        # 设置操作类型为添加
        self._operation = REGISTRY_OPERATION.ADD

        # 初始化变量并创建远程批处理文件
        Registry._initVars(self, regKey, regValue, regType, regData)
        self._createRemoteBatchFile()

        # 构造并记录调试信息
        debugMsg = "adding registry key value '%s' " % self._regValue
        debugMsg += "to registry key '%s'" % self._regKey
        logger.debug(debugMsg)

        # 执行远程批处理文件并清理
        self.execCmd(cmd=self._batPathRemote)
        self.delRemoteFile(self._batPathRemote)

    def delRegKey(self, regKey, regValue):
        """
        删除注册表键值
        :param regKey: 要删除的注册表键路径
        :param regValue: 要删除的值名称
        """
        # 设置操作类型为删除
        self._operation = REGISTRY_OPERATION.DELETE

        # 初始化变量并创建远程批处理文件
        Registry._initVars(self, regKey, regValue)
        self._createRemoteBatchFile()

        # 构造并记录调试信息
        debugMsg = "deleting registry key value '%s' " % self._regValue
        debugMsg += "from registry key '%s'" % self._regKey
        logger.debug(debugMsg)

        # 执行远程批处理文件并清理
        self.execCmd(cmd=self._batPathRemote)
        self.delRemoteFile(self._batPathRemote)