From 9c30fc0819af55d61ccd859d0128b736c936ac5b Mon Sep 17 00:00:00 2001 From: chenjiandongx Date: Mon, 22 Nov 2021 23:40:16 +0800 Subject: [PATCH] Feat: display pid and support customize unit in the table widget --- README.md | 2 +- cli.go | 4 ++-- conn_darwin.go | 13 ++++++------- conn_linux.go | 15 ++++++++++----- conn_windows.go | 13 +++++++++---- pcap.go | 12 +++++++++++- sniffer.go | 2 +- stat.go | 2 +- ui.go | 20 +++++++++++++++----- 9 files changed, 56 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 8c343f8..9d38b1b 100644 --- a/README.md +++ b/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 diff --git a/cli.go b/cli.go index dc6ebf7..75b2931 100644 --- a/cli.go +++ b/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 diff --git a/conn_darwin.go b/conn_darwin.go index a08bac3..7817eea 100644 --- a/conn_darwin.go +++ b/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 } } diff --git a/conn_linux.go b/conn_linux.go index 381ab88..6b1c1e6 100644 --- a/conn_linux.go +++ b/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) { - sockets := make(map[LocalSocket]string) +func (nl *netlinkConn) sockdiagRecv(skfd, proto int, inodeMap map[uint32]string) (OpenSockets, error) { + 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) { - sockets := make(map[LocalSocket]string) +func (nl *netlinkConn) getOpenSockets(inodeMap map[uint32]string) (OpenSockets, error) { + sockets := make(OpenSockets) type Req struct { Protocol int diff --git a/conn_windows.go b/conn_windows.go index 8441034..862df3b 100644 --- a/conn_windows.go +++ b/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 { diff --git a/pcap.go b/pcap.go index ee2c928..364c179 100644 --- a/pcap.go +++ b/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 { diff --git a/sniffer.go b/sniffer.go index f52cb52..2c1d940 100644 --- a/sniffer.go +++ b/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 diff --git a/stat.go b/stat.go index 1213c06..f40f92e 100644 --- a/stat.go +++ b/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 diff --git a/ui.go b/ui.go index 268e287..c5f2e9e 100644 --- a/ui.go +++ b/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{":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", ":Process", "Up / Down"} tv.connections.Rows = [][]string{header, make([]string, 3)} tv.connections.Rows = append(tv.connections.Rows, rows...) }