#!/usr/bin/env python3
import re
import os

# 需要清理的文件（保留ssh.go redis.go ms17010.go）
FILES_TO_CLEAN = [
    'elasticsearch.go', 'findnet.go', 'ftp.go', 'kafka.go', 
    'ldap.go', 'netbios.go', 'rabbitmq.go', 'rdp.go', 
    'rsync.go', 'smb.go', 'smb2.go', 'smbghost.go', 
    'smbinfo.go', 'smtp.go', 'snmp.go', 'telnet.go', 
    'vnc.go', 'webpoc.go', 'webtitle.go'
]

def clean_exploit_function(file_path):
    """清理单个文件的Exploit函数"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 使用正则表达式匹配并删除Exploit函数
    # 匹配从 "// Exploit" 注释开始到函数结束的整个块
    pattern = r'\/\/\s*Exploit.*?利用操作.*?\nfunc\s+\([^)]*\)\s+Exploit\([^{]*\{[^}]*(?:\{[^}]*\}[^}]*)*\}\n'
    
    # 简化方法：按行处理
    lines = content.split('\n')
    output_lines = []
    skip_mode = False
    brace_count = 0
    
    i = 0
    while i < len(lines):
        line = lines[i]
        
        # 检测Exploit函数开始
        if ('// Exploit' in line and '利用操作' in line) or ('func ' in line and 'Exploit(' in line):
            skip_mode = True
            brace_count = 0
            if 'func ' in line and '{' in line:
                brace_count = line.count('{') - line.count('}')
            i += 1
            continue
        
        if skip_mode:
            # 计算花括号
            brace_count += line.count('{')
            brace_count -= line.count('}')
            
            # 如果花括号归零，函数结束
            if brace_count <= 0:
                skip_mode = False
            i += 1
            continue
        
        output_lines.append(line)
        i += 1
    
    # 写回文件
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(output_lines))

def main():
    os.chdir('plugins/services')
    
    for filename in FILES_TO_CLEAN:
        if os.path.exists(filename):
            print(f"Cleaning {filename}...")
            try:
                clean_exploit_function(filename)
                print(f"✅ {filename}")
            except Exception as e:
                print(f"❌ {filename}: {e}")
        else:
            print(f"⚠️  {filename} not found")

if __name__ == "__main__":
    main()