zutils/cmd/ports.go

package cmd

import (
	"fmt"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/zemenawi/zutils/pkg/colors"
	"github.com/zemenawi/zutils/pkg/formatter"
)

type PortInfo struct {
	Protocol  string
	LocalAddr string
	Port      int
	Process   string
	PID       int
	State     string
}

func PortsCommand(args []string) error {
	showAll := false
	portFilter := 0

	for _, arg := range args {
		if arg == "-a" || arg == "--all" {
			showAll = true
		} else if strings.HasPrefix(arg, "-") {
			continue
		} else {
			if port, err := strconv.Atoi(arg); err == nil {
				portFilter = port
			}
		}
	}

	ports, err := getPortList()
	if err != nil {
		return fmt.Errorf("error getting ports: %w", err)
	}

	if portFilter > 0 {
		var filtered []PortInfo
		for _, p := range ports {
			if p.Port == portFilter {
				filtered = append(filtered, p)
			}
		}
		ports = filtered
	}

	if len(ports) == 0 {
		if portFilter > 0 {
			return fmt.Errorf("no process found using port %d", portFilter)
		}
		fmt.Printf("%sNo listening ports found%s\n", colors.Yellow, colors.Reset)
		return nil
	}

	sort.Slice(ports, func(i, j int) bool {
		if ports[i].Port != ports[j].Port {
			return ports[i].Port < ports[j].Port
		}
		return ports[i].Protocol < ports[j].Protocol
	})

	PrintBoxHeader("PORT LISTENING", colors.Cyan)

	if portFilter > 0 {
		fmt.Printf("%s%s🔍 Filtering by port: %d%s\n", colors.Yellow, colors.Bold, portFilter, colors.Reset)
		fmt.Println()
	}

	if !showAll && len(ports) > 50 {
		fmt.Printf("%s%s📊 Showing %d of %d ports (use -a for all)%s\n\n",
			colors.Yellow, colors.Bold, 50, len(ports), colors.Reset)
		ports = ports[:50]
	}

	headers := []string{"#", "Protocol", "Port", "Process", "PID", "State"}
	data := make([][]string, 0, len(ports))

	for i, p := range ports {
		state := p.State
		if state == "LISTEN" {
			state = colors.Green + "LISTEN" + colors.Reset
		} else if state == "ESTABLISHED" {
			state = colors.Cyan + "ESTABLISHED" + colors.Reset
		} else {
			state = colors.Gray + state + colors.Reset
		}

		row := []string{
			fmt.Sprintf("%d", i+1),
			strings.ToUpper(p.Protocol),
			fmt.Sprintf("%d", p.Port),
			p.Process,
			fmt.Sprintf("%d", p.PID),
			state,
		}
		data = append(data, row)
	}

	formatter.PrintTable(headers, data, colors.Cyan)
	fmt.Println()

	fmt.Printf("%s💡 Tip:%s Use %s z ports <port>%s to filter by specific port\n",
		colors.Cyan, colors.Reset, colors.Green, colors.Reset)

	return nil
}

func getPortList() ([]PortInfo, error) {
	var ports []PortInfo

	cmd := exec.Command("ss", "-tulpn")
	output, err := cmd.CombinedOutput()
	if err != nil {
		cmd = exec.Command("netstat", "-tulpn")
		output, err = cmd.CombinedOutput()
		if err != nil {
			return ports, fmt.Errorf("ss and netstat both failed")
		}
	}

	lines := strings.Split(string(output), "\n")
	for _, line := range lines[1:] {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "Netid") || strings.HasPrefix(line, "State") {
			continue
		}

		port := parseSSLine(line)
		if port != nil && port.Port > 0 {
			ports = append(ports, *port)
		}
	}

	return ports, nil
}

func parseSSLine(line string) *PortInfo {
	info := &PortInfo{}

	protoMatch := regexp.MustCompile(`^(tcp6?|udp6?|udplite)\s+`)
	protoMatch2 := protoMatch.FindStringSubmatch(strings.TrimSpace(line))
	if len(protoMatch2) > 1 {
		info.Protocol = protoMatch2[1]
		if info.Protocol == "udp" || info.Protocol == "udplite" {
			info.Protocol = "udp"
		}
	}

	stateMatch := regexp.MustCompile(`(LISTEN|UNCONN|ESTAB|CLOSE-WAIT|CLOSING|SYN-SENT|SYN-RECV|FIN-WAIT-1|FIN-WAIT-2|TIME-WAIT)\s+`)
	stateMatch2 := stateMatch.FindStringSubmatch(line)
	if len(stateMatch2) > 1 {
		info.State = stateMatch2[1]
	} else {
		info.State = "-"
	}

	portStr := ""
	for i := len(line) - 1; i >= 0; i-- {
		if line[i] == ':' {
			start := i + 1
			for start < len(line) && line[start] == ' ' {
				start++
			}
			end := start
			for end < len(line) && line[end] >= '0' && line[end] <= '9' {
				end++
			}
			if end > start {
				portStr = line[start:end]
				break
			}
		}
	}
	if port, err := strconv.Atoi(portStr); err == nil {
		info.Port = port
	}

	pidMatch := regexp.MustCompile(`pid=(\d+)`)
	pidMatch2 := pidMatch.FindStringSubmatch(line)
	if len(pidMatch2) > 1 {
		info.PID, _ = strconv.Atoi(pidMatch2[1])
	}

	processMatch := regexp.MustCompile(`process_name="([^"]+)"`)
	processMatch2 := processMatch.FindStringSubmatch(line)
	if len(processMatch2) > 1 {
		info.Process = processMatch2[1]
	}

	return info
}

func getProcessNameByPID(pid int) string {
	if data, err := os.ReadFile(fmt.Sprintf("/proc/%d/comm", pid)); err == nil {
		return strings.TrimSpace(string(data))
	}
	return "unknown"
}

type ProcessStats struct {
	PID       int
	Name      string
	Command   string
	CPUPercent float64
	MemBytes  int64
	MemPercent float64
	Status    string
	User      string
}

func UsagesCommand(args []string) error {
	sortBy := "cpu"
	limit := 20
	filter := ""

	for _, arg := range args {
		if arg == "-m" || arg == "--memory" {
			sortBy = "memory"
		} else if arg == "-c" || arg == "--cpu" {
			sortBy = "cpu"
		} else if strings.HasPrefix(arg, "-n=") {
			if n, err := strconv.Atoi(strings.TrimPrefix(arg, "-n=")); err == nil && n > 0 {
				limit = n
			}
		} else if strings.HasPrefix(arg, "--limit=") {
			if n, err := strconv.Atoi(strings.TrimPrefix(arg, "--limit=")); err == nil && n > 0 {
				limit = n
			}
		} else if strings.HasPrefix(arg, "-") {
			continue
		} else {
			filter = arg
		}
	}

	processes, err := getAllProcessStats()
	if err != nil {
		return fmt.Errorf("error getting process stats: %w", err)
	}

	if filter != "" {
		var filtered []ProcessStats
		for _, p := range processes {
			if strings.Contains(strings.ToLower(p.Name), strings.ToLower(filter)) ||
				strings.Contains(strings.ToLower(p.Command), strings.ToLower(filter)) {
				filtered = append(filtered, p)
			}
		}
		processes = filtered
		if len(processes) == 0 {
			return fmt.Errorf("no processes found matching '%s'", filter)
		}
	}

	sort.Slice(processes, func(i, j int) bool {
		if sortBy == "memory" {
			return processes[i].MemBytes > processes[j].MemBytes
		}
		return processes[i].CPUPercent > processes[j].CPUPercent
	})

	if limit > 0 && len(processes) > limit {
		processes = processes[:limit]
	}

	PrintBoxHeader(fmt.Sprintf("TOP PROCESSES (by %s)", strings.ToUpper(sortBy)), colors.Cyan)

	if filter != "" {
		fmt.Printf("%s%s🔍 Filtering by: %s%s\n", colors.Yellow, colors.Bold, filter, colors.Reset)
		fmt.Println()
	}

	headers := []string{"#", "PID", "Process", "CPU%", "Memory", "Mem%", "Status"}
	data := make([][]string, 0, len(processes))

	for i, p := range processes {
		status := p.Status
		if status == "R" || status == "running" {
			status = colors.Green + "running" + colors.Reset
		} else if status == "S" || status == "I" || status == "sleeping" {
			status = colors.Gray + "sleeping" + colors.Reset
		} else {
			status = colors.Yellow + status + colors.Reset
		}

		cmd := p.Name
		if len(p.Command) > 25 {
			cmd = p.Command[:25] + "..."
		}

		row := []string{
			fmt.Sprintf("%d", i+1),
			fmt.Sprintf("%d", p.PID),
			cmd,
			fmt.Sprintf("%.1f", p.CPUPercent),
			formatter.FormatSize(p.MemBytes),
			fmt.Sprintf("%.1f%%", p.MemPercent),
			status,
		}
		data = append(data, row)
	}

	formatter.PrintTable(headers, data, colors.Cyan)
	fmt.Println()

	if filter != "" {
		var totalCPU float64
		var totalMem int64
		for _, p := range processes {
			totalCPU += p.CPUPercent
			totalMem += p.MemBytes
		}

		PrintSectionHeader("SUMMARY")
		fmt.Printf("%s%s📊 Total CPU:%s      %.1f%%\n", colors.Green, colors.Bold, colors.Reset, totalCPU)
		fmt.Printf("%s%s💾 Total Memory:%s   %s (%.1f%% of system)%s\n",
			colors.Purple, colors.Bold, colors.Reset,
			formatter.FormatSize(totalMem),
			float64(totalMem)/float64(getSystemMem())*100,
			colors.Reset)
		fmt.Println()
	}

	fmt.Printf("%s💡 Tip:%s Use %s z usages -m%s for memory sorting, %s z usages -n=50%s to show 50 processes\n",
		colors.Cyan, colors.Reset, colors.Green, colors.Reset, colors.Green, colors.Reset)

	return nil
}

func getSystemMem() int64 {
	if data, err := os.ReadFile("/proc/meminfo"); err == nil {
		lines := strings.Split(string(data), "\n")
		for _, line := range lines {
			if strings.HasPrefix(line, "MemTotal:") {
				fields := strings.Fields(line)
				if len(fields) >= 2 {
					if memKB, err := strconv.ParseInt(fields[1], 10, 64); err == nil {
						return memKB * 1024
					}
				}
			}
		}
	}
	return 1
}

func getAllProcessStats() ([]ProcessStats, error) {
	var processes []ProcessStats

	dir, err := os.Open("/proc")
	if err != nil {
		return processes, err
	}
	defer dir.Close()

	entries, err := dir.Readdirnames(0)
	if err != nil {
		return processes, err
	}

	for _, entry := range entries {
		pid, err := strconv.Atoi(entry)
		if err != nil {
			continue
		}

		proc, err := getFullProcessInfo(pid)
		if err == nil && proc != nil {
			processes = append(processes, *proc)
		}
	}

	return processes, nil
}

func getFullProcessInfo(pid int) (*ProcessStats, error) {
	comm, err := os.ReadFile(fmt.Sprintf("/proc/%d/comm", pid))
	if err != nil {
		return nil, err
	}
	name := strings.TrimSpace(string(comm))

	cmdline, _ := os.ReadFile(fmt.Sprintf("/proc/%d/cmdline", pid))
	cmdStr := strings.Join(strings.Fields(string(cmdline)), " ")
	if cmdStr == "" {
		cmdStr = name
	}

	stat, err := os.ReadFile(fmt.Sprintf("/proc/%d/stat", pid))
	if err != nil {
		return nil, err
	}

	statParts := strings.Split(string(stat), " ")
	status := "?"
	if len(statParts) > 2 {
		status = strings.Trim(statParts[2], "()")
	}

	utime := int64(0)
	stime := int64(0)
	if len(statParts) >= 15 {
		utime, _ = strconv.ParseInt(strings.TrimSpace(statParts[13]), 10, 64)
		stime, _ = strconv.ParseInt(strings.TrimSpace(statParts[14]), 10, 64)
	}

	memBytes := int64(0)
	if data, err := os.ReadFile(fmt.Sprintf("/proc/%d/status", pid)); err == nil {
		lines := strings.Split(string(data), "\n")
		for _, line := range lines {
			if strings.HasPrefix(line, "VmRSS:") {
				fields := strings.Fields(line)
				if len(fields) >= 2 {
					if memKB, err := strconv.ParseInt(fields[1], 10, 64); err == nil {
						memBytes = memKB * 1024
					}
				}
				break
			}
		}
	}

	totalMem := int64(0)
	if data, err := os.ReadFile("/proc/meminfo"); err == nil {
		lines := strings.Split(string(data), "\n")
		for _, line := range lines {
			if strings.HasPrefix(line, "MemTotal:") {
				fields := strings.Fields(line)
				if len(fields) >= 2 {
					if memKB, err := strconv.ParseInt(fields[1], 10, 64); err == nil {
						totalMem = memKB * 1024
					}
				}
				break
			}
		}
	}

	cpuPercent := 0.0
	if totalMem > 0 {
		cpuPercent = float64(utime+stime) / float64(totalMem) * 100
	}

	memPercent := 0.0
	if totalMem > 0 {
		memPercent = float64(memBytes) / float64(totalMem) * 100
	}

	return &ProcessStats{
		PID:        pid,
		Name:       name,
		Command:    cmdStr,
		CPUPercent: cpuPercent,
		MemBytes:   memBytes,
		MemPercent: memPercent,
		Status:     status,
	}, nil
}

func NetworkCommand(args []string) error {
	return NetworkInfoCommand(args)
}