port-manager/internal/scanner/port_scanner.go

package scanner

import (
	"bufio"
	"fmt"
	"net"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/miaomint/port-manager/pkg/types"
)

// PortScanner 端口扫描器
type PortScanner struct {
	includeListening   bool
	includeEstablished bool
}

// NewPortScanner 创建新的端口扫描器
func NewPortScanner(includeListening, includeEstablished bool) *PortScanner {
	return &PortScanner{
		includeListening:   includeListening,
		includeEstablished: includeEstablished,
	}
}

// ScanPorts 扫描本地端口
func (ps *PortScanner) ScanPorts() (*types.ScanResult, error) {
	var ports []types.PortInfo

	// 获取TCP端口
	tcpPorts, err := ps.scanTCPPorts()
	if err != nil {
		return nil, fmt.Errorf("扫描TCP端口失败: %v", err)
	}
	ports = append(ports, tcpPorts...)

	// 获取UDP端口
	udpPorts, err := ps.scanUDPPorts()
	if err != nil {
		return nil, fmt.Errorf("扫描UDP端口失败: %v", err)
	}
	ports = append(ports, udpPorts...)

	return &types.ScanResult{
		Ports:      ports,
		Timestamp:  time.Now(),
		TotalPorts: len(ports),
	}, nil
}

// scanTCPPorts 扫描TCP端口（使用netstat）
func (ps *PortScanner) scanTCPPorts() ([]types.PortInfo, error) {
	var args []string
	args = append(args, "-an")
	args = append(args, "-p", "tcp")

	cmd := exec.Command("netstat", args...)
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	return ps.parseNetstatOutput(string(output), "tcp")
}

// scanUDPPorts 扫描UDP端口
func (ps *PortScanner) scanUDPPorts() ([]types.PortInfo, error) {
	var args []string
	args = append(args, "-an")
	args = append(args, "-p", "udp")

	cmd := exec.Command("netstat", args...)
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	return ps.parseNetstatOutput(string(output), "udp")
}

// parseNetstatOutput 解析netstat输出
func (ps *PortScanner) parseNetstatOutput(output, protocol string) ([]types.PortInfo, error) {
	var ports []types.PortInfo
	scanner := bufio.NewScanner(strings.NewReader(output))

	// 跳过标题行
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "Local Address") {
			break
		}
	}

	// 解析每一行
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}

		port, err := ps.parseNetstatLine(line, protocol)
		if err != nil {
			continue // 跳过解析失败的行
		}

		if port != nil {
			// 过滤状态
			if ps.shouldIncludePort(port) {
				ports = append(ports, *port)
			}
		}
	}

	return ports, nil
}

// parseNetstatLine 解析netstat的单行输出
func (ps *PortScanner) parseNetstatLine(line, protocol string) (*types.PortInfo, error) {
	fields := regexp.MustCompile(`\s+`).Split(line, -1)
	if len(fields) < 3 {
		return nil, fmt.Errorf("invalid line format")
	}

	// 解析本地地址
	localAddr := fields[0]
	parts := strings.Split(localAddr, ":")
	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid address format")
	}

	portStr := parts[len(parts)-1]
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("invalid port: %s", portStr)
	}

	state := ""
	if protocol == "tcp" && len(fields) >= 4 {
		state = fields[3]
	}

	return &types.PortInfo{
		Port:     port,
		Protocol: protocol,
		State:    state,
		LastSeen: time.Now(),
	}, nil
}

// shouldIncludePort 判断是否应该包含此端口
func (ps *PortScanner) shouldIncludePort(port *types.PortInfo) bool {
	if port.Protocol == "tcp" {
		if port.State == "LISTEN" {
			return ps.includeListening
		}
		if port.State == "ESTABLISHED" {
			return ps.includeEstablished
		}
	}

	// UDP端口默认包含
	if port.Protocol == "udp" {
		return ps.includeListening
	}

	return false
}

// ScanPortRange 扫描指定端口范围
func (ps *PortScanner) ScanPortRange(start, end int) ([]int, error) {
	var openPorts []int

	for port := start; port <= end; port++ {
		// TCP扫描
		conn, err := net.DialTimeout("tcp", fmt.Sprintf("localhost:%d", port), time.Millisecond*100)
		if err == nil {
			conn.Close()
			openPorts = append(openPorts, port)
		}
	}

	return openPorts, nil
}

// IsPortOpen 检查指定端口是否开放
func (ps *PortScanner) IsPortOpen(port int, protocol string) bool {
	address := fmt.Sprintf("localhost:%d", port)
	conn, err := net.DialTimeout(protocol, address, time.Millisecond*100)
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

// GenerateRandomPort 生成一个未被占用的随机端口
func (ps *PortScanner) GenerateRandomPort(minPort, maxPort int) (int, error) {
	if minPort <= 0 {
		minPort = 8000
	}
	if maxPort <= 0 || maxPort > 65535 {
		maxPort = 65535
	}

	for attempts := 0; attempts < 100; attempts++ {
		port := minPort + (int(time.Now().UnixNano()) % (maxPort - minPort + 1))

		if !ps.IsPortOpen(port, "tcp") {
			return port, nil
		}
	}

	return 0, fmt.Errorf("无法找到可用端口")
}