Feat: display pid and support customize unit in the table widget

README.md

@@ -6,7 +6,7 @@
 
 > *A modern alternative network traffic sniffer inspired by [bandwhich](https://github.com/imsnif/bandwhich)(Rust) and [nethogs](https://github.com/raboof/nethogs)(C++).*
 
-https://user-images.githubusercontent.com/19553554/142255776-4f01f06f-af79-4fb1-820f-7a00c71416d3.mov
+https://user-images.githubusercontent.com/19553554/142890205-62980e37-5861-4161-9669-737317573aa1.mov
 
 ## Introduction
 

cli.go

@@ -6,7 +6,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-const version = "v0.2.0"
+const version = "v0.3.0"
 
 func NewApp() *cobra.Command {
 	defaultOpts := DefaultOptions()
@@ -58,7 +58,7 @@ func NewApp() *cobra.Command {
 	app.Flags().BoolVarP(&opt.DisableDNSResolve, "no-dns-resolve", "n", defaultOpts.DisableDNSResolve, "disable the DNS resolution")
 	app.Flags().Int32SliceVarP(&opt.Pids, "pids", "p", defaultOpts.Pids, "pids to watch, empty stands for all pids")
 	app.Flags().IntVarP(&mode, "mode", "m", int(defaultOpts.ViewMode), "view mode of sniffer (0: bytes 1: packets 2: processes)")
-	app.Flags().StringVarP(&unit, "unit", "u", defaultOpts.Unit.String(), "unit of traffic stats in processes mode, optional: B, KB, MB, GB")
+	app.Flags().StringVarP(&unit, "unit", "u", defaultOpts.Unit.String(), "unit of traffic stats, optional: B, Kb, KB, Mb, MB, Gb, GB")
 
 	app.Flags().PrintDefaults()
 	return app

conn_darwin.go

@@ -58,18 +58,17 @@ func (lc *lsofConn) GetOpenSockets(pids ...int32) (OpenSockets, error) {
 		if len(fields) < 10 {
 			continue
 		}
-		procName := strings.ReplaceAll(fields[0], "\\x20", " ")
 
+		procName := strings.ReplaceAll(fields[0], "\\x20", " ")
+		pid, _ := strconv.Atoi(fields[1])
 		if len(pids) > 0 {
-			pid, err := strconv.Atoi(fields[1])
-			if err != nil {
-				continue
-			}
 			if !set[int32(pid)] {
 				continue
 			}
 		}
 
+		procInfo := ProcessInfo{Pid: pid, Name: procName}
+
 		switch fields[8] {
 		case "TCP":
 			addr := strings.Split(fields[9], "->")
@@ -84,7 +83,7 @@ func (lc *lsofConn) GetOpenSockets(pids ...int32) (OpenSockets, error) {
 			if err != nil {
 				continue
 			}
-			sockets[LocalSocket{IP: ipport[0], Port: uint16(port), Protocol: ProtoTCP}] = procName
+			sockets[LocalSocket{IP: ipport[0], Port: uint16(port), Protocol: ProtoTCP}] = procInfo
 
 		case "UDP":
 			ipport := strings.Split(fields[9], ":")
@@ -96,7 +95,7 @@ func (lc *lsofConn) GetOpenSockets(pids ...int32) (OpenSockets, error) {
 			if err != nil {
 				continue
 			}
-			sockets[LocalSocket{IP: ipport[0], Port: uint16(port), Protocol: ProtoUDP}] = procName
+			sockets[LocalSocket{IP: ipport[0], Port: uint16(port), Protocol: ProtoUDP}] = procInfo
 		}
 	}
 

conn_linux.go

@@ -215,8 +215,8 @@ func (nl *netlinkConn) sockdiagSend(proto, family uint8, states uint32) (skfd in
 	return skfd, nil
 }
 
-func (nl *netlinkConn) sockdiagRecv(skfd, proto int, inodeMap map[uint32]string) (map[LocalSocket]string, error) {
+func (nl *netlinkConn) sockdiagRecv(skfd, proto int, inodeMap map[uint32]string) (OpenSockets, error) {
-	sockets := make(map[LocalSocket]string)
+	sockets := make(OpenSockets)
 	buffer := make([]byte, os.Getpagesize())
 loop:
 	for {
@@ -242,6 +242,11 @@ loop:
 			m := (*inetDiagMsg)(unsafe.Pointer(&msg.Data[0]))
 			srcIP, _ := nl.ipHex2String(m.IDiagFamily, m.ID.IdiagSrc)
 
+			procInfo := ProcessInfo{
+				Pid:  int(msg.Header.Pid),
+				Name: inodeMap[m.IDiagInode],
+			}
+
 			var p Protocol
 			switch proto {
 			case syscall.IPPROTO_TCP:
@@ -249,15 +254,15 @@ loop:
 			case syscall.IPPROTO_UDP:
 				p = ProtoUDP
 			}
-			sockets[LocalSocket{IP: srcIP, Port: uint16(m.ID.IdiagSport.Int()), Protocol: p}] = inodeMap[m.IDiagInode]
+			sockets[LocalSocket{IP: srcIP, Port: uint16(m.ID.IdiagSport.Int()), Protocol: p}] = procInfo
 		}
 	}
 
 	return sockets, nil
 }
 
-func (nl *netlinkConn) getOpenSockets(inodeMap map[uint32]string) (map[LocalSocket]string, error) {
+func (nl *netlinkConn) getOpenSockets(inodeMap map[uint32]string) (OpenSockets, error) {
-	sockets := make(map[LocalSocket]string)
+	sockets := make(OpenSockets)
 
 	type Req struct {
 		Protocol int

conn_windows.go

@@ -28,16 +28,21 @@ func (ps *psutilConn) getOpenSockets(pids ...int32) (OpenSockets, error) {
 	return openSockets, nil
 }
 
-func (ps *psutilConn) getProcName(pid int32) string {
+func (ps *psutilConn) getProcName(pid int32) ProcessInfo {
+	procInfo := ProcessInfo{Name: unknownProcessName}
+
 	proc, err := process.NewProcess(pid)
 	if err != nil {
-		return unknownProcessName
+		return procInfo
 	}
 	exe, err := proc.Exe()
 	if err != nil {
-		return unknownProcessName
+		return procInfo
 	}
-	return filepath.Base(exe)
+
+	procInfo.Pid = int(pid)
+	procInfo.Name = filepath.Base(exe)
+	return procInfo
 }
 
 func (ps *psutilConn) getConnections(proto Protocol, openSockets OpenSockets, pids ...int32) error {

pcap.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"errors"
+	"fmt"
 	"strconv"
 	"strings"
 	"sync"
@@ -28,7 +29,16 @@ type Connection struct {
 	Remote RemoteSocket
 }
 
-type OpenSockets map[LocalSocket]string
+type ProcessInfo struct {
+	Pid  int
+	Name string
+}
+
+func (p ProcessInfo) String() string {
+	return fmt.Sprintf("<%d>:%s", p.Pid, p.Name)
+}
+
+type OpenSockets map[LocalSocket]ProcessInfo
 type Utilization map[Connection]*ConnectionInfo
 
 type SocketFetcher interface {

sniffer.go

@@ -31,7 +31,7 @@ type Options struct {
 	// Pids to watch in processes mode
 	Pids []int32
 
-	// Unit of stats in processes mode, optional: B, KB, MB, GB
+	// Unit of stats in processes mode, optional: B, Kb, KB, Mb, MB, Gb, GB
 	Unit Unit
 
 	// DisableDNSResolve decides whether if disable the DNS resolution

stat.go

@@ -176,7 +176,7 @@ func (s *StatsManager) getProcName(openSockets OpenSockets, localSocket LocalSoc
 
 		v, ok := openSockets[cloned]
 		if ok {
-			return v
+			return v.String()
 		}
 	}
 	return unknownProcessName

ui.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"fmt"
 	"strconv"
-	"strings"
 	"time"
 
 	"github.com/chenjiandongx/termui/v3"
@@ -44,13 +43,16 @@ type Unit string
 const (
 	UnitB  Unit = "B"
 	UnitKB Unit = "KB"
+	UnitKb Unit = "Kb"
 	UnitMB Unit = "MB"
+	UnitMb Unit = "Mb"
 	UnitGB Unit = "GB"
+	UnitGb Unit = "Gb"
 )
 
 func (u Unit) Validate() error {
 	switch u {
-	case UnitB, UnitKB, UnitMB, UnitGB:
+	case UnitB, UnitKB, UnitKb, UnitMB, UnitMb, UnitGB, UnitGb:
 		return nil
 	}
 	return fmt.Errorf("invalid unit %s", u)
@@ -67,10 +69,16 @@ func (u Unit) Ratio() float64 {
 		ratio = 1
 	case UnitKB:
 		ratio = 1024
+	case UnitKb:
+		ratio = 1024 / 8
 	case UnitMB:
 		ratio = 1024 * 1024
+	case UnitMb:
+		ratio = 1024 * 1024 / 8
 	case UnitGB:
 		ratio = 1024 * 1024 * 1024
+	case UnitGb:
+		ratio = 1024 * 1024 * 1024 / 8
 	}
 	return ratio
 }
@@ -122,6 +130,7 @@ func NewUIComponent(opt Options) *UIComponent {
 			remoteAddrs: newTable("Remote Address"),
 			connections: newTable("Connections"),
 			mode:        opt.ViewMode,
+			unit:        opt.Unit,
 		}
 	default:
 		ui.viewer = &PlotViewer{
@@ -327,6 +336,7 @@ type TableViewer struct {
 	grid        *termui.Grid
 	shiftIdx    int
 	mode        ViewMode
+	unit        Unit
 }
 
 func (tv *TableViewer) Setup() {
@@ -352,7 +362,7 @@ func (tv *TableViewer) humanizeNum(n int) string {
 	var s string
 	switch tv.mode {
 	case ModeTableBytes:
-		s = strings.ReplaceAll(humanize.IBytes(uint64(n)), " ", "")
+		s = fmt.Sprintf("%.1f%s", float64(n)/tv.unit.Ratio(), tv.unit.String())
 	case ModeTablePackets:
 		s = humanize.Comma(int64(n))
 	}
@@ -387,7 +397,7 @@ func (tv *TableViewer) updateProcesses(snapshot *Snapshot) {
 		rows = append(rows, []string{r.ProcessName, strconv.Itoa(r.Data.ConnCount), up + " / " + down})
 	}
 
-	header := []string{"Process", "Connections", "Up / Down"}
+	header := []string{"<Pid>:Process", "Connections", "Up / Down"}
 	tv.processes.Rows = [][]string{header, make([]string, 3)}
 	tv.processes.Rows = append(tv.processes.Rows, rows...)
 }
@@ -435,7 +445,7 @@ func (tv *TableViewer) updateConnections(snapshot *Snapshot) {
 		rows = append(rows, []string{conn, r.Data.ProcessName, up + " / " + down})
 	}
 
-	header := []string{"Connections", "Process", "Up / Down"}
+	header := []string{"Connections", "<Pid>:Process", "Up / Down"}
 	tv.connections.Rows = [][]string{header, make([]string, 3)}
 	tv.connections.Rows = append(tv.connections.Rows, rows...)
 }