# Sqlmap工具类
import shlex
from tools.ToolBase import ToolBase

class SqlmapTool(ToolBase):
    def validate_instruction(self, instruction):
        timeout = 0
        # 检查sqlmap高风险参数
        high_risk_params = [
            "--os-shell",
            "--os-cmd",
            "--os-pwn",
            "--os-sql-shell",
            "--file-read",
            "--file-write",
            "--reg-add",
            "--reg-del",
            "--eval"
        ]
        # 将命令转换为小写，确保判断不区分大小写
        cmd_lower = instruction.lower()
        for param in high_risk_params:
            if param in cmd_lower:
                return ""
        #检查--batch
        parts = shlex.split(cmd_lower)
        if "--batch" not in parts:
            parts.append("--batch")

        return " ".join(shlex.quote(part) for part in parts),timeout

    def analyze_result(self, result,instruction,stderr,stdout):
        # 检查结果中是否包含"vulnerable"，表示SQL注入漏洞
        return "发现SQL注入漏洞" if "vulnerable" in result else "未发现SQL注入漏洞"

if __name__ == "__main__":
    ST = SqlmapTool()
    strcmd = "sqlmap -u \"http://haitutech.cn/news?id=1\" --os-shell --reg-add --reg-key=\"Software\\Microsoft\\Windows\\CurrentVersion\\Internet Settings\" --reg-value=ProxyEnable --reg-data=0 --reg-type=REG_DWORD"
    res,time_out = ST.validate_instruction(strcmd)
    print("11")
    print(res)