//go:build windows

// Package collector 提供 Windows 监听端口发现能力。 // Logan.Li
package collector

import (
	"encoding/csv"
	"os/exec"
	"strconv"
	"strings"

	"ezops/agent/model"
)

type PortCollector struct {
	Policy *model.PortPolicy
}

func NewPortCollector() *PortCollector {
	return &PortCollector{}
}

func (p *PortCollector) SetPolicy(policy *model.PortPolicy) {
	p.Policy = policy
}

func (p *PortCollector) Collect() ([]model.PortStat, error) {
	pidNames := windowsProcessNames()
	out, err := exec.Command("netstat", "-ano", "-p", "tcp").Output()
	if err != nil {
		return []model.PortStat{}, nil
	}
	seen := map[string]bool{}
	var ports []model.PortStat
	for _, line := range strings.Split(string(out), "\n") {
		fields := strings.Fields(line)
		if len(fields) < 5 || strings.ToUpper(fields[0]) != "TCP" || strings.ToUpper(fields[3]) != "LISTENING" {
			continue
		}
		port := parseWindowsPort(fields[1])
		if port <= 0 {
			continue
		}
		if p.Policy != nil && p.Policy.IsBlacklistMode() && !p.Policy.InWhitelist(port) {
			continue
		}
		pid, _ := strconv.Atoi(fields[4])
		key := strconv.Itoa(port) + "|" + strconv.Itoa(pid)
		if seen[key] {
			continue
		}
		seen[key] = true
		ports = append(ports, model.PortStat{
			Port:     port,
			Protocol: "tcp",
			Process:  pidNames[pid],
			PID:      pid,
		})
	}
	return ports, nil
}

func parseWindowsPort(addr string) int {
	addr = strings.Trim(addr, "[]")
	idx := strings.LastIndex(addr, ":")
	if idx < 0 || idx+1 >= len(addr) {
		return 0
	}
	port, _ := strconv.Atoi(addr[idx+1:])
	return port
}

func windowsProcessNames() map[int]string {
	out, err := exec.Command("tasklist", "/FO", "CSV", "/NH").Output()
	if err != nil {
		return map[int]string{}
	}
	reader := csv.NewReader(strings.NewReader(string(out)))
	rows, err := reader.ReadAll()
	if err != nil {
		return map[int]string{}
	}
	result := map[int]string{}
	for _, row := range rows {
		if len(row) < 2 {
			continue
		}
		pid, err := strconv.Atoi(strings.TrimSpace(row[1]))
		if err == nil {
			result[pid] = row[0]
		}
	}
	return result
}