import gradio as gr
import subprocess

# 常见 SQLMap 指令选项
COMMON_OPTIONS = {
    "URL": "",
    "参数名称": "",
    "数据库类型": "",
    "线程数": "1",
    "代理": "",
}

# 构建命令函数
def build_sqlmap_command(url, param, dbms, threads, proxy, extra):
    cmd = ["sqlmap", "-u", url]
    if param:
        cmd += ["--param", param]
    if dbms:
        cmd += ["-p", dbms]
    cmd += ["--threads", threads]
    if proxy:
        cmd += ["--proxy", proxy]
    if extra:
        cmd += extra.split()
    return cmd

# 执行函数

def run_sqlmap(url, param, dbms, threads, proxy, extra):
    cmd = build_sqlmap_command(url, param, dbms, threads, proxy, extra)
    try:
        result = subprocess.check_output(cmd, stderr=subprocess.STDOUT, text=True, timeout=300)
    except subprocess.CalledProcessError as e:
        result = e.output
    except subprocess.TimeoutExpired:
        result = "执行超时，请检查目标或减少负载。"
    return result

# Gradio 界面布局
with gr.Blocks(title="SQLMap Web UI") as demo:
    gr.Markdown("# SQLMap Web UI")
    with gr.Row():
        with gr.Column():
            url = gr.Textbox(label="目标 URL", placeholder="http://example.com/vuln.php?id=1")
            param = gr.Textbox(label="参数名称 (param)", placeholder="id, user 等，可留空")
            dbms = gr.Dropdown(label="数据库类型 (dbms)", choices=["", "MySQL", "PostgreSQL", "MSSQL", "Oracle", "SQLite", "MongoDB"], value="")
            threads = gr.Slider(label="线程数", minimum=1, maximum=10, step=1, value=1)
            proxy = gr.Textbox(label="HTTP 代理 (可选)", placeholder="http://127.0.0.1:8080")
            extra = gr.Textbox(label="额外参数 (如: --os-shell)", placeholder="--dump --batch 等")
            run_btn = gr.Button("运行 SQLMap")
        with gr.Column():
            output = gr.Textbox(label="输出结果", interactive=False, lines=20)

    run_btn.click(fn=run_sqlmap, inputs=[url, param, dbms, threads, proxy, extra], outputs=output)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)