staros/network_windows.go

//go:build windows
// +build windows

package staros

import (
	"errors"
	"net"
	"strconv"
	"strings"
	"syscall"
	"time"

	"b612.me/win32api"
)

const windowsErrorNotSupported syscall.Errno = 50

func NetUsage() ([]NetAdapter, error) {
	rows, err := win32api.GetIfTable2()
	if err != nil {
		return nil, err
	}
	res := make([]NetAdapter, 0, len(rows))
	for _, row := range rows {
		name := windowsInterfaceName(row)
		if name == "" {
			continue
		}
		res = append(res, NetAdapter{
			Name:      name,
			RecvBytes: row.InOctets,
			SendBytes: row.OutOctets,
		})
	}
	if len(res) == 0 {
		return nil, errors.New("No Adaptor")
	}
	return res, nil
}

func NetUsageByname(name string) (NetAdapter, error) {
	ada, err := NetUsage()
	if err != nil {
		return NetAdapter{}, err
	}
	for _, v := range ada {
		if v.Name == name {
			return v, nil
		}
	}
	return NetAdapter{}, errors.New("Not Found")
}

func NetSpeeds(duration time.Duration) ([]NetSpeed, error) {
	if duration <= 0 {
		return nil, errors.New("duration must be positive")
	}
	list1, err := NetUsage()
	if err != nil {
		return nil, err
	}
	time.Sleep(duration)
	list2, err := NetUsage()
	if err != nil {
		return nil, err
	}
	byName := make(map[string]NetAdapter, len(list2))
	for _, item := range list2 {
		byName[item.Name] = item
	}
	res := make([]NetSpeed, 0, len(list1))
	for _, v := range list1 {
		next, ok := byName[v.Name]
		if !ok {
			continue
		}
		var recvDelta, sendDelta uint64
		if next.RecvBytes >= v.RecvBytes {
			recvDelta = next.RecvBytes - v.RecvBytes
		}
		if next.SendBytes >= v.SendBytes {
			sendDelta = next.SendBytes - v.SendBytes
		}
		res = append(res, NetSpeed{
			Name:       v.Name,
			RecvSpeeds: float64(recvDelta) / duration.Seconds(),
			SendSpeeds: float64(sendDelta) / duration.Seconds(),
			RecvBytes:  next.RecvBytes,
			SendBytes:  next.SendBytes,
		})
	}
	if len(res) == 0 {
		return nil, errors.New("NetWork Adaptor Num Not ok")
	}
	return res, nil
}

func NetSpeedsByName(duration time.Duration, name string) (NetSpeed, error) {
	ada, err := NetSpeeds(duration)
	if err != nil {
		return NetSpeed{}, err
	}
	for _, v := range ada {
		if v.Name == name {
			return v, nil
		}
	}
	return NetSpeed{}, errors.New("Not Found")
}

// NetConnections return all TCP/UDP/UNIX DOMAIN SOCKET Connections
// if your uid != 0 ,and analysePid==true ,you should have CAP_SYS_PRTACE and CAP_DAC_OVERRIDE/CAP_DAC_READ_SEARCH Caps
func NetConnections(analysePid bool, types string) ([]NetConn, error) {
	wantTCP, wantUDP, err := windowsNetConnectionTypes(types)
	if err != nil {
		return nil, err
	}

	result := make([]NetConn, 0)
	processCache := make(map[int64]*Process)
	if wantTCP {
		result, err = appendWindowsTCPConnections(result, analysePid, processCache)
		if err != nil {
			return result, err
		}
	}
	if wantUDP {
		result, err = appendWindowsUDPConnections(result, analysePid, processCache)
		if err != nil {
			return result, err
		}
	}
	return result, nil
}

func GetInodeMap() (map[string]int64, error) {
	return nil, ERR_UNSUPPORTED
}

func windowsInterfaceName(row win32api.MIB_IF_ROW2) string {
	name := strings.TrimSpace(syscall.UTF16ToString(row.Alias[:]))
	if name != "" {
		return name
	}
	name = strings.TrimSpace(syscall.UTF16ToString(row.Description[:]))
	if name != "" {
		return name
	}
	if row.InterfaceIndex != 0 {
		return "if" + strconv.FormatUint(uint64(row.InterfaceIndex), 10)
	}
	if row.InterfaceLuid != 0 {
		return "luid" + strconv.FormatUint(row.InterfaceLuid, 10)
	}
	return ""
}

func windowsNetConnectionTypes(types string) (wantTCP, wantUDP bool, err error) {
	normalized := strings.ToLower(strings.TrimSpace(types))
	if strings.Contains(normalized, "unix") {
		return false, false, ERR_UNSUPPORTED
	}
	if normalized == "" || strings.Contains(normalized, "all") {
		return true, true, nil
	}
	if strings.Contains(normalized, "tcp") {
		wantTCP = true
	}
	if strings.Contains(normalized, "udp") {
		wantUDP = true
	}
	if !wantTCP && !wantUDP {
		return false, false, errors.New("unsupported net connection type")
	}
	return wantTCP, wantUDP, nil
}

func appendWindowsTCPConnections(result []NetConn, analysePid bool, processCache map[int64]*Process) ([]NetConn, error) {
	rows4, err := win32api.GetExtendedTcp4Table(false, win32api.TCP_TABLE_OWNER_PID_ALL)
	if err != nil {
		return result, err
	}
	for _, row := range rows4 {
		conn := NetConn{
			LocalAddr:  windowsIPv4FromDWORD(row.LocalAddr),
			LocalPort:  int(row.LocalPortHost()),
			RemoteAddr: windowsIPv4FromDWORD(row.RemoteAddr),
			RemotePort: int(row.RemotePortHost()),
			Status:     windowsTCPState(row.State),
			Typed:      "tcp",
		}
		attachWindowsProcess(&conn, row.OwningPid, analysePid, processCache)
		result = append(result, conn)
	}

	rows6, err := win32api.GetExtendedTcp6Table(false, win32api.TCP_TABLE_OWNER_PID_ALL)
	if err != nil {
		if isOptionalWindowsNetTableError(err) {
			return result, nil
		}
		return result, err
	}
	for _, row := range rows6 {
		conn := NetConn{
			LocalAddr:  net.IP(row.LocalAddr[:]).String(),
			LocalPort:  int(row.LocalPortHost()),
			RemoteAddr: net.IP(row.RemoteAddr[:]).String(),
			RemotePort: int(row.RemotePortHost()),
			Status:     windowsTCPState(row.State),
			Typed:      "tcp6",
		}
		attachWindowsProcess(&conn, row.OwningPid, analysePid, processCache)
		result = append(result, conn)
	}
	return result, nil
}

func appendWindowsUDPConnections(result []NetConn, analysePid bool, processCache map[int64]*Process) ([]NetConn, error) {
	rows4, err := win32api.GetExtendedUdp4Table(false, win32api.UDP_TABLE_OWNER_PID)
	if err != nil {
		return result, err
	}
	for _, row := range rows4 {
		conn := NetConn{
			LocalAddr: windowsIPv4FromDWORD(row.LocalAddr),
			LocalPort: int(row.LocalPortHost()),
			Typed:     "udp",
		}
		attachWindowsProcess(&conn, row.OwningPid, analysePid, processCache)
		result = append(result, conn)
	}

	rows6, err := win32api.GetExtendedUdp6Table(false, win32api.UDP_TABLE_OWNER_PID)
	if err != nil {
		if isOptionalWindowsNetTableError(err) {
			return result, nil
		}
		return result, err
	}
	for _, row := range rows6 {
		conn := NetConn{
			LocalAddr: net.IP(row.LocalAddr[:]).String(),
			LocalPort: int(row.LocalPortHost()),
			Typed:     "udp6",
		}
		attachWindowsProcess(&conn, row.OwningPid, analysePid, processCache)
		result = append(result, conn)
	}
	return result, nil
}

func attachWindowsProcess(conn *NetConn, pid uint32, analysePid bool, processCache map[int64]*Process) {
	if conn == nil || !analysePid {
		return
	}
	conn.Pid = int64(pid)
	if conn.Pid <= 0 {
		return
	}
	if proc, ok := processCache[conn.Pid]; ok {
		conn.Process = proc
		return
	}
	proc, err := FindProcessByPid(conn.Pid)
	if err != nil {
		processCache[conn.Pid] = nil
		return
	}
	processCache[conn.Pid] = &proc
	conn.Process = &proc
}

func windowsIPv4FromDWORD(addr uint32) string {
	return net.IPv4(byte(addr), byte(addr>>8), byte(addr>>16), byte(addr>>24)).String()
}

func windowsTCPState(state win32api.MIB_TCP_STATE) string {
	switch state {
	case win32api.MIB_TCP_STATE_CLOSED:
		return TCP_STATE[TCP_CLOSE]
	case win32api.MIB_TCP_STATE_LISTEN:
		return TCP_STATE[TCP_LISTEN]
	case win32api.MIB_TCP_STATE_SYN_SENT:
		return TCP_STATE[TCP_SYN_SENT]
	case win32api.MIB_TCP_STATE_SYN_RCVD:
		return TCP_STATE[TCP_SYN_RECV]
	case win32api.MIB_TCP_STATE_ESTAB:
		return TCP_STATE[TCP_ESTABLISHED]
	case win32api.MIB_TCP_STATE_FIN_WAIT1:
		return TCP_STATE[TCP_FIN_WAIT1]
	case win32api.MIB_TCP_STATE_FIN_WAIT2:
		return TCP_STATE[TCP_FIN_WAIT2]
	case win32api.MIB_TCP_STATE_CLOSE_WAIT:
		return TCP_STATE[TCP_CLOSE_WAIT]
	case win32api.MIB_TCP_STATE_CLOSING:
		return TCP_STATE[TCP_CLOSING]
	case win32api.MIB_TCP_STATE_LAST_ACK:
		return TCP_STATE[TCP_LAST_ACK]
	case win32api.MIB_TCP_STATE_TIME_WAIT:
		return TCP_STATE[TCP_TIME_WAIT]
	case win32api.MIB_TCP_STATE_DELETE_TCB:
		return "TCP_DELETE_TCB"
	default:
		return TCP_STATE[TCP_UNKNOWN]
	}
}

func isOptionalWindowsNetTableError(err error) bool {
	if errno, ok := err.(syscall.Errno); ok {
		return errno == windowsErrorNotSupported
	}
	return false
}