package avdetect

import (
	"bufio"
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"

	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
	"github.com/shadow1ng/fscan/plugins/local"
)

//go:embed auto.json
var embeddedAVDatabase []byte

// AVDetectPlugin AV/EDR检测插件 - 使用简化架构
type AVDetectPlugin struct {
	*local.BaseLocalPlugin
	avDatabase map[string]AVProduct
}

// AVProduct AV/EDR产品信息
type AVProduct struct {
	Processes []string `json:"processes"`
	URL       string   `json:"url"`
}

// ProcessInfo 进程信息
type ProcessInfo struct {
	Name        string
	PID         string
	SessionName string
	SessionID   string
	MemUsage    string
	Services    []string // 服务信息
}

// DetectionResult 检测结果
type DetectionResult struct {
	ProductName     string        `json:"product_name"`
	DetectedProcesses []ProcessInfo `json:"detected_processes"`
	URL             string        `json:"url"`
	RiskLevel       string        `json:"risk_level"`
	Category        string        `json:"category"`
}

// NewAVDetectPlugin 创建AV/EDR检测插件 - 简化版本
func NewAVDetectPlugin() *AVDetectPlugin {
	metadata := &base.PluginMetadata{
		Name:        "avdetect",
		Version:     "1.0.0", 
		Author:      "fscan-team",
		Description: "自动化AV/EDR检测插件，基于嵌入式规则库识别安全软件",
		Category:    "local",
		Tags:        []string{"local", "av", "edr", "detection", "security"},
		Protocols:   []string{"local"},
	}

	plugin := &AVDetectPlugin{
		BaseLocalPlugin: local.NewBaseLocalPlugin(metadata),
		avDatabase:      make(map[string]AVProduct),
	}

	// 设置支持的平台 (仅Windows)
	plugin.SetPlatformSupport([]string{"windows"})
	// 不需要特殊权限
	plugin.SetRequiresPrivileges(false)

	return plugin
}

// Initialize 初始化插件
func (p *AVDetectPlugin) Initialize() error {
	// 先调用基类初始化
	if err := p.BaseLocalPlugin.Initialize(); err != nil {
		return err
	}
	
	// 加载AV数据库
	return p.loadAVDatabase()
}

// getRunningProcesses 获取运行中的进程列表
func (p *AVDetectPlugin) getRunningProcesses() ([]ProcessInfo, error) {
	var processes []ProcessInfo
	var cmd *exec.Cmd
	
	switch runtime.GOOS {
	case "windows":
		// Windows使用PowerShell获取进程信息以避免编码问题
		cmd = exec.Command("powershell", "-Command", "Get-Process | Select-Object Name,Id,ProcessName | ConvertTo-Csv -NoTypeInformation")
	case "linux", "darwin":
		// Unix-like系统使用ps命令
		cmd = exec.Command("ps", "aux")
	default:
		return nil, fmt.Errorf("不支持的操作系统: %s", runtime.GOOS)
	}
	
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("执行命令失败: %v", err)
	}
	
	// 解析命令输出
	processes, err = p.parseProcessOutput(string(output))
	if err != nil {
		return nil, fmt.Errorf("解析进程信息失败: %v", err)
	}
	
	return processes, nil
}

// parseProcessOutput 解析进程命令输出
func (p *AVDetectPlugin) parseProcessOutput(output string) ([]ProcessInfo, error) {
	var processes []ProcessInfo
	scanner := bufio.NewScanner(strings.NewReader(output))
	
	switch runtime.GOOS {
	case "windows":
		// 跳过CSV标题行
		if scanner.Scan() {
			// 标题行，跳过
		}
		
		for scanner.Scan() {
			line := scanner.Text()
			if line == "" {
				continue
			}
			
			// 解析PowerShell CSV格式：Name,Id,ProcessName
			fields := p.parseCSVLine(line)
			if len(fields) >= 3 {
				processName := strings.Trim(fields[0], "\"")
				// 如果进程名不包含.exe，则添加
				if !strings.HasSuffix(strings.ToLower(processName), ".exe") {
					processName += ".exe"
				}
				
				process := ProcessInfo{
					Name: processName,
					PID:  strings.Trim(fields[1], "\""),
				}
				
				processes = append(processes, process)
			}
		}
		
	case "linux", "darwin":
		// 跳过ps命令的标题行
		if scanner.Scan() {
			// 标题行，跳过
		}
		
		for scanner.Scan() {
			line := scanner.Text()
			if line == "" {
				continue
			}
			
			// 解析ps aux输出
			fields := strings.Fields(line)
			if len(fields) >= 11 {
				process := ProcessInfo{
					Name:     fields[10], // 命令名
					PID:      fields[1],  // PID
					MemUsage: fields[5],  // 内存使用
				}
				processes = append(processes, process)
			}
		}
	}
	
	return processes, scanner.Err()
}

// parseCSVLine 解析CSV行，处理引号内的逗号
func (p *AVDetectPlugin) parseCSVLine(line string) []string {
	var fields []string
	var current strings.Builder
	inQuotes := false
	
	for i, char := range line {
		switch char {
		case '"':
			inQuotes = !inQuotes
			current.WriteRune(char)
		case ',':
			if inQuotes {
				current.WriteRune(char)
			} else {
				fields = append(fields, current.String())
				current.Reset()
			}
		default:
			current.WriteRune(char)
		}
		
		// 处理行尾
		if i == len(line)-1 {
			fields = append(fields, current.String())
		}
	}
	
	return fields
}

// Scan 重写扫描方法以确保调用正确的ScanLocal实现
func (p *AVDetectPlugin) Scan(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	return p.ScanLocal(ctx, info)
}

// ScanLocal 执行AV/EDR检测扫描 - 简化版本
func (p *AVDetectPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	common.LogInfo("开始AV/EDR安全软件检测...")

	// 获取运行进程
	processes, err := p.getRunningProcesses()
	if err != nil {
		common.LogError(fmt.Sprintf("获取进程列表失败: %v", err))
		// 不返回错误，继续执行但结果可能不完整
		processes = []ProcessInfo{}
	}

	common.LogDebug(fmt.Sprintf("获取到 %d 个运行进程", len(processes)))

	// 检测AV/EDR产品
	detectionResults := p.detectAVEDR(processes)

	// 获取系统信息
	systemInfo := p.GetSystemInfo()

	// 生成检测报告
	report := p.generateDetectionReport(detectionResults, systemInfo)

	if len(detectionResults) == 0 {
		common.LogInfo("未检测到已知的AV/EDR安全产品")
		return &base.ScanResult{
			Success: true,
			Service: "AVDetect", 
			Banner:  "未检测到已知的AV/EDR安全产品",
			Extra: map[string]interface{}{
				"detected_products": detectionResults,
				"total_processes":   len(processes),
				"detection_report":  report,
				"platform":         runtime.GOOS,
				"database_products": len(p.avDatabase),
			},
		}, nil
	}

	// 输出检测结果
	common.LogInfo(fmt.Sprintf("[+] AV/EDR检测完成: 发现 %d 个安全产品", len(detectionResults)))
	for _, result := range detectionResults {
		common.LogInfo(fmt.Sprintf("[+] 检测到: %s (%d个进程)", result.ProductName, len(result.DetectedProcesses)))
	}

	result := &base.ScanResult{
		Success: true,
		Service: "AVDetect",
		Banner:  fmt.Sprintf("检测完成: 发现 %d 个安全产品", len(detectionResults)),
		Extra: map[string]interface{}{
			"detected_products": detectionResults,
			"total_processes":   len(processes),
			"detection_report":  report,
			"platform":         runtime.GOOS,
			"database_products": len(p.avDatabase),
		},
	}
	
	if len(detectionResults) > 0 {
		common.LogSuccess(fmt.Sprintf("AV/EDR检测完成: 发现 %d 个安全产品", len(detectionResults)))
		for _, detection := range detectionResults {
			common.LogSuccess(fmt.Sprintf("检测到: %s (%d个进程)", detection.ProductName, len(detection.DetectedProcesses)))
		}
	} else {
		common.LogBase("未检测到已知的AV/EDR安全产品")
	}
	
	return result, nil
}

// loadAVDatabase 加载AV/EDR数据库
func (p *AVDetectPlugin) loadAVDatabase() error {
	// 首先尝试使用嵌入的数据库
	common.LogDebug("使用嵌入的AV/EDR规则数据库")
	
	err := json.Unmarshal(embeddedAVDatabase, &p.avDatabase)
	if err != nil {
		return fmt.Errorf("解析嵌入的AV数据库失败: %v", err)
	}
	
	if len(p.avDatabase) == 0 {
		return fmt.Errorf("嵌入的AV数据库为空或格式错误")
	}
	
	return nil
}

// detectAVEDR 检测AV/EDR产品
func (p *AVDetectPlugin) detectAVEDR(processes []ProcessInfo) []DetectionResult {
	var results []DetectionResult
	processMap := make(map[string][]ProcessInfo)
	
	// 构建进程名称映射，忽略大小写
	for _, process := range processes {
		processName := strings.ToLower(process.Name)
		processMap[processName] = append(processMap[processName], process)
	}
	
	// 遍历AV数据库进行匹配
	for productName, avProduct := range p.avDatabase {
		var detectedProcesses []ProcessInfo
		
		// 检查每个已知进程
		for _, targetProcess := range avProduct.Processes {
			targetProcessLower := strings.ToLower(targetProcess)
			
			// 精确匹配
			if matchedProcesses, exists := processMap[targetProcessLower]; exists {
				detectedProcesses = append(detectedProcesses, matchedProcesses...)
			} else {
				// 模糊匹配（去除扩展名）
				targetWithoutExt := strings.TrimSuffix(targetProcessLower, filepath.Ext(targetProcessLower))
				for processName, matchedProcesses := range processMap {
					processWithoutExt := strings.TrimSuffix(processName, filepath.Ext(processName))
					if processWithoutExt == targetWithoutExt {
						detectedProcesses = append(detectedProcesses, matchedProcesses...)
						break
					}
				}
			}
		}
		
		// 如果检测到进程，添加到结果中
		if len(detectedProcesses) > 0 {
			// 去重
			detectedProcesses = p.deduplicateProcesses(detectedProcesses)
			
			result := DetectionResult{
				ProductName:       productName,
				DetectedProcesses: detectedProcesses,
				URL:               avProduct.URL,
				RiskLevel:         p.assessRiskLevel(productName, detectedProcesses),
				Category:          p.categorizeProduct(productName),
			}
			
			results = append(results, result)
		}
	}
	
	// 按检测到的进程数量排序
	sort.Slice(results, func(i, j int) bool {
		return len(results[i].DetectedProcesses) > len(results[j].DetectedProcesses)
	})
	
	return results
}

// deduplicateProcesses 去重进程列表
func (p *AVDetectPlugin) deduplicateProcesses(processes []ProcessInfo) []ProcessInfo {
	seen := make(map[string]bool)
	var result []ProcessInfo
	
	for _, process := range processes {
		key := fmt.Sprintf("%s-%s", process.Name, process.PID)
		if !seen[key] {
			seen[key] = true
			result = append(result, process)
		}
	}
	
	return result
}

// assessRiskLevel 评估风险等级
func (p *AVDetectPlugin) assessRiskLevel(productName string, processes []ProcessInfo) string {
	// 基于产品名称和进程数量评估风险等级
	productLower := strings.ToLower(productName)
	
	// 高风险EDR产品
	highRiskKeywords := []string{"crowdstrike", "sentinelone", "cybereason", "endgame", 
		"fireeye", "trellix", "elastic security", "深信服", "奇安信", "天擎"}
	
	for _, keyword := range highRiskKeywords {
		if strings.Contains(productLower, strings.ToLower(keyword)) {
			return "HIGH"
		}
	}
	
	// 中等风险（企业级AV）
	mediumRiskKeywords := []string{"kaspersky", "symantec", "mcafee", "趋势科技", 
		"bitdefender", "eset", "sophos", "火绒", "360"}
	
	for _, keyword := range mediumRiskKeywords {
		if strings.Contains(productLower, strings.ToLower(keyword)) {
			return "MEDIUM"
		}
	}
	
	// 根据进程数量判断
	if len(processes) >= 3 {
		return "MEDIUM"
	}
	
	return "LOW"
}

// categorizeProduct 产品分类
func (p *AVDetectPlugin) categorizeProduct(productName string) string {
	productLower := strings.ToLower(productName)
	
	// EDR产品
	edrKeywords := []string{"edr", "endpoint", "crowdstrike", "sentinelone", 
		"cybereason", "深信服edr", "天擎", "elastic security"}
	
	for _, keyword := range edrKeywords {
		if strings.Contains(productLower, strings.ToLower(keyword)) {
			return "EDR"
		}
	}
	
	// 企业级防病毒
	enterpriseKeywords := []string{"enterprise", "business", "server", 
		"corporate", "管理版", "企业版"}
	
	for _, keyword := range enterpriseKeywords {
		if strings.Contains(productLower, strings.ToLower(keyword)) {
			return "Enterprise AV"
		}
	}
	
	// 云安全
	cloudKeywords := []string{"cloud", "阿里云", "腾讯云", "云锁", "云安全"}
	
	for _, keyword := range cloudKeywords {
		if strings.Contains(productLower, strings.ToLower(keyword)) {
			return "Cloud Security"
		}
	}
	
	// 主机防护
	hostKeywords := []string{"host", "hips", "主机", "防护", "卫士"}
	
	for _, keyword := range hostKeywords {
		if strings.Contains(productLower, strings.ToLower(keyword)) {
			return "Host Protection"
		}
	}
	
	return "Traditional AV"
}

// generateDetectionReport 生成检测报告
func (p *AVDetectPlugin) generateDetectionReport(results []DetectionResult, systemInfo map[string]string) string {
	var report strings.Builder
	
	report.WriteString("=== AV/EDR 检测报告 ===\n")
	report.WriteString(fmt.Sprintf("扫描时间: %s\n", time.Now().Format("2006-01-02 15:04:05")))
	report.WriteString(fmt.Sprintf("系统平台: %s/%s\n", systemInfo["os"], systemInfo["arch"]))
	report.WriteString(fmt.Sprintf("检测产品: %d 个\n\n", len(results)))
	
	if len(results) == 0 {
		report.WriteString("未检测到已知的AV/EDR产品\n")
		report.WriteString("注意: 可能存在未知安全软件或进程伪装\n")
		return report.String()
	}
	
	// 按风险等级分组
	riskGroups := map[string][]DetectionResult{
		"HIGH":   {},
		"MEDIUM": {},
		"LOW":    {},
	}
	
	for _, result := range results {
		riskGroups[result.RiskLevel] = append(riskGroups[result.RiskLevel], result)
	}
	
	// 高风险产品
	if len(riskGroups["HIGH"]) > 0 {
		report.WriteString("🔴 高风险安全产品:\n")
		for _, result := range riskGroups["HIGH"] {
			report.WriteString(fmt.Sprintf("  • %s [%s] - %d 个进程\n", 
				result.ProductName, result.Category, len(result.DetectedProcesses)))
			for _, process := range result.DetectedProcesses {
				report.WriteString(fmt.Sprintf("    - %s (PID: %s)\n", process.Name, process.PID))
			}
		}
		report.WriteString("\n")
	}
	
	// 中等风险产品
	if len(riskGroups["MEDIUM"]) > 0 {
		report.WriteString("🟡 中等风险安全产品:\n")
		for _, result := range riskGroups["MEDIUM"] {
			report.WriteString(fmt.Sprintf("  • %s [%s] - %d 个进程\n", 
				result.ProductName, result.Category, len(result.DetectedProcesses)))
		}
		report.WriteString("\n")
	}
	
	// 低风险产品
	if len(riskGroups["LOW"]) > 0 {
		report.WriteString("🟢 低风险安全产品:\n")
		for _, result := range riskGroups["LOW"] {
			report.WriteString(fmt.Sprintf("  • %s [%s] - %d 个进程\n", 
				result.ProductName, result.Category, len(result.DetectedProcesses)))
		}
		report.WriteString("\n")
	}
	
	// 建议
	report.WriteString("=== 渗透测试建议 ===\n")
	if len(riskGroups["HIGH"]) > 0 {
		report.WriteString("⚠️  检测到高级EDR产品，建议:\n")
		report.WriteString("   - 使用内存加载技术\n")
		report.WriteString("   - 避免落地文件\n")
		report.WriteString("   - 使用白名单绕过技术\n")
		report.WriteString("   - 考虑Living off the Land技术\n\n")
	}
	
	if len(results) > 1 {
		report.WriteString("📊 检测到多个安全产品，环境复杂度较高\n")
	}
	
	return report.String()
}

// GetLocalData 获取AV/EDR检测本地数据
func (p *AVDetectPlugin) GetLocalData(ctx context.Context) (map[string]interface{}, error) {
	data := make(map[string]interface{})
	
	// 获取系统信息
	data["plugin_type"] = "avdetect"
	data["platform"] = runtime.GOOS
	data["arch"] = runtime.GOARCH
	data["database_size"] = len(p.avDatabase)
	
	if homeDir, err := os.UserHomeDir(); err == nil {
		data["home_dir"] = homeDir
	}
	
	if workDir, err := os.Getwd(); err == nil {
		data["work_dir"] = workDir
	}
	
	return data, nil
}

// ExtractData 提取AV/EDR检测数据
func (p *AVDetectPlugin) ExtractData(ctx context.Context, info *common.HostInfo, data map[string]interface{}) (*base.ExploitResult, error) {
	return &base.ExploitResult{
		Success: true,
		Output:  "AV/EDR检测完成",
		Data:    data,
		Extra: map[string]interface{}{
			"detection_type": "automated",
			"database_version": "auto.json",
		},
	}, nil
}

// GetInfo 获取插件信息
func (p *AVDetectPlugin) GetInfo() string {
	var info strings.Builder
	
	info.WriteString(fmt.Sprintf("AV/EDR自动检测插件 - 规则库: %d 个产品\n", len(p.avDatabase)))
	info.WriteString(fmt.Sprintf("支持平台: %s\n", strings.Join(p.GetPlatformSupport(), ", ")))
	info.WriteString("检测方式: tasklist/ps + JSON规则匹配\n")
	info.WriteString("功能: 自动识别常见AV/EDR产品并评估风险等级\n")
	
	return info.String()
}

// RegisterAVDetectPlugin 注册AV/EDR检测插件
func RegisterAVDetectPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "avdetect",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "自动化AV/EDR检测插件，基于auto.json规则库识别安全软件",
			Category:    "local",
			Tags:        []string{"avdetect", "local", "av", "edr", "security"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewAVDetectPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("avdetect", factory)
}

// init 插件注册函数
func init() {
	RegisterAVDetectPlugin()
}