diff --git a/conn_darwin.go b/conn_darwin.go index 8aa7927..208b7f2 100644 --- a/conn_darwin.go +++ b/conn_darwin.go @@ -6,7 +6,6 @@ package main import ( "bytes" "context" - "fmt" "os/exec" "strconv" "strings" @@ -14,19 +13,16 @@ import ( ) type lsofConn struct { - invoker Invoker + invoker LsofInvoker } -type Invoker struct{} +type LsofInvoker struct{} -func (i Invoker) Command(name string, arg ...string) ([]byte, error) { +func (i LsofInvoker) Exec() ([]byte, error) { ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() - return i.CommandWithContext(ctx, name, arg...) -} -func (i Invoker) CommandWithContext(ctx context.Context, name string, arg ...string) ([]byte, error) { - cmd := exec.CommandContext(ctx, name, arg...) + cmd := exec.CommandContext(ctx, "lsof", "-n", "-R", "-P", "-iTCP", "-iUDP", "-s", "TCP:ESTABLISHED", "+c", "0") var buf bytes.Buffer cmd.Stdout = &buf @@ -43,27 +39,47 @@ func (i Invoker) CommandWithContext(ctx context.Context, name string, arg ...str return buf.Bytes(), nil } -func (lc *lsofConn) GetProcSockets(pid int32) (OpenSockets, error) { return nil, nil } +func (lc *lsofConn) GetProcSockets(pids ...int32) (OpenSockets, error) { + return lc.getOpenSockets(pids...) +} func (lc *lsofConn) GetOpenSockets() (OpenSockets, error) { + return lc.getOpenSockets() +} + +func (lc *lsofConn) getOpenSockets(pids ...int32) (OpenSockets, error) { sockets := make(OpenSockets) - output, err := lc.invoker.Command("lsof", "-n", "-P", "-iTCP", "-iUDP", "-s", "TCP:ESTABLISHED", "+c", "0") + output, err := lc.invoker.Exec() if err != nil { - fmt.Println(err) return sockets, err } + set := make(map[int32]bool) + for _, pid := range pids { + set[pid] = true + } + lines := strings.Split(string(output), "\n") for _, line := range lines { fields := strings.Fields(line) - if len(fields) < 9 { + if len(fields) < 10 { continue } procName := strings.ReplaceAll(fields[0], "\\x20", " ") - switch fields[7] { + if len(pids) > 0 { + pid, err := strconv.Atoi(fields[1]) + if err != nil { + continue + } + if !set[int32(pid)] { + continue + } + } + + switch fields[8] { case "TCP": - addr := strings.Split(fields[8], "->") + addr := strings.Split(fields[9], "->") if len(addr) != 2 { continue } @@ -75,10 +91,21 @@ func (lc *lsofConn) GetOpenSockets() (OpenSockets, error) { if err != nil { continue } + + if len(pids) > 0 { + pid, err := strconv.Atoi(fields[1]) + if err != nil { + continue + } + if !set[int32(pid)] { + continue + } + } + sockets[LocalSocket{IP: ipport[0], Port: uint16(port), Protocol: ProtoTCP}] = procName case "UDP": - ipport := strings.Split(fields[8], ":") + ipport := strings.Split(fields[9], ":") if len(ipport) != 2 { continue } diff --git a/conn_linux.go b/conn_linux.go index 05efd0b..d1aaece 100644 --- a/conn_linux.go +++ b/conn_linux.go @@ -300,7 +300,7 @@ func (nl *netlinkConn) getOpenSockets(inodeMap map[uint32]string) (map[LocalSock return sockets, nil } -func (nl *netlinkConn) getAllProcsInodes(pids []int32) map[uint32]string { +func (nl *netlinkConn) getAllProcsInodes(pids ...int32) map[uint32]string { inode2Procs := make(map[uint32]string) for _, pid := range pids { procName, inodes, err := nl.getProcInodes(pid) @@ -382,12 +382,13 @@ func (nl *netlinkConn) GetOpenSockets() (OpenSockets, error) { return nil, err } - inodeMap := nl.getAllProcsInodes(pids) + inodeMap := nl.getAllProcsInodes(pids...) return nl.getOpenSockets(inodeMap) } -func (nl *netlinkConn) GetProcSockets(pid int32) (OpenSockets, error) { - return nil, nil +func (nl *netlinkConn) GetProcSockets(pids ...int32) (OpenSockets, error) { + inodeMap := nl.getAllProcsInodes(pids...) + return nl.getOpenSockets(inodeMap) } func GetSocketFetcher() SocketFetcher { diff --git a/conn_windows.go b/conn_windows.go index 674ea7f..9c0304b 100644 --- a/conn_windows.go +++ b/conn_windows.go @@ -13,19 +13,25 @@ import ( type psutilConn struct{} func (ps *psutilConn) GetOpenSockets() (OpenSockets, error) { + return ps.getOpenSockets() +} + +func (ps *psutilConn) GetProcSockets(pids ...int32) (OpenSockets, error) { + return ps.getOpenSockets(pids...) +} + +func (ps *psutilConn) getOpenSockets(pids ...int32) (OpenSockets, error) { openSockets := make(OpenSockets) - if err := ps.getConnections(ProtoTCP, openSockets); err != nil { + if err := ps.getConnections(ProtoTCP, openSockets, pids...); err != nil { return nil, err } - if err := ps.getConnections(ProtoUDP, openSockets); err != nil { + if err := ps.getConnections(ProtoUDP, openSockets, pids...); err != nil { return nil, err } return openSockets, nil } -func (ps *psutilConn) GetProcSockets(pid int32) (OpenSockets, error) { return nil, nil } - func (ps *psutilConn) getProcName(pid int32) string { proc, err := process.NewProcess(pid) if err != nil { @@ -38,30 +44,32 @@ func (ps *psutilConn) getProcName(pid int32) string { return filepath.Base(exe) } -func (ps *psutilConn) getConnections(proto Protocol, openSockets OpenSockets) error { - protos := []string{"tcp", "tcp6"} - if proto == ProtoUDP { - protos = []string{"udp", "udp6"} +func (ps *psutilConn) getConnections(proto Protocol, openSockets OpenSockets, pids ...int32) error { + connections, err := net.Connections(string(proto)) + if err != nil { + return err } - for _, p := range protos { - connections, err := net.Connections(p) - if err != nil { - return err + set := make(map[int32]bool) + for _, pid := range pids { + set[pid] = true + } + + for _, conn := range connections { + if proto == ProtoTCP && conn.Status != "ESTABLISHED" { + continue } - for _, conn := range connections { - if proto == ProtoTCP && conn.Status != "ESTABLISHED" { - continue - } - - localSocket := LocalSocket{ - IP: conn.Laddr.IP, - Port: uint16(conn.Laddr.Port), - Protocol: proto, - } - openSockets[localSocket] = ps.getProcName(conn.Pid) + if len(pids) > 0 && !set[conn.Pid] { + continue } + + localSocket := LocalSocket{ + IP: conn.Laddr.IP, + Port: uint16(conn.Laddr.Port), + Protocol: proto, + } + openSockets[localSocket] = ps.getProcName(conn.Pid) } return nil } diff --git a/go.mod b/go.mod index 7bc626d..8e4719f 100644 --- a/go.mod +++ b/go.mod @@ -15,4 +15,4 @@ require ( golang.org/x/sys v0.0.0-20210816074244-15123e1e1f71 ) -replace github.com/gizak/termui/v3 v3.1.0 => ../../gizak/termui +replace github.com/gizak/termui/v3 v3.1.0 => ../../chenjiandongx/termui diff --git a/pcap.go b/pcap.go index 177b51e..5ee53de 100644 --- a/pcap.go +++ b/pcap.go @@ -32,7 +32,7 @@ type Utilization map[Connection]*ConnectionInfo type SocketFetcher interface { GetOpenSockets() (OpenSockets, error) - GetProcSockets(pid int32) (OpenSockets, error) + GetProcSockets(pid ...int32) (OpenSockets, error) } type Protocol string diff --git a/sniffer.go b/sniffer.go index ced4532..3dbd077 100644 --- a/sniffer.go +++ b/sniffer.go @@ -70,7 +70,7 @@ func NewSniffer(fn ...OptionsFn) (*Sniffer, error) { opts: &opts, dnsResolver: dnsResolver, pcapClient: pcapClient, - statsManager: NewStatsManager(opts.Interval), + statsManager: NewStatsManager(opts.Interval, opts.RenderMode), ui: NewUIComponent(opts.RenderMode), socketFetcher: GetSocketFetcher(), }, nil diff --git a/stat.go b/stat.go index 89eed39..14ccdd0 100644 --- a/stat.go +++ b/stat.go @@ -145,12 +145,14 @@ type StatsManager struct { mut sync.Mutex ring *deque.Deque ratio int + mode RenderMode } -func NewStatsManager(ratio int) *StatsManager { +func NewStatsManager(ratio int, mode RenderMode) *StatsManager { return &StatsManager{ ring: deque.New(), ratio: ratio, + mode: mode, } } @@ -199,6 +201,12 @@ func (s *StatsManager) GetSnapshot() *Snapshot { stat := s.ring.At(i).(Stat) for conn, info := range stat.Utilization { procName := s.getProcName(stat.OpenSockets, conn.Local) + if s.mode == RModeProcess { + if procName == unknownProcessName { + continue + } + } + if _, ok := connections[conn]; !ok { connections[conn] = &ConnectionData{ InterfaceName: info.Interface, diff --git a/ui.go b/ui.go index da1b622..7158712 100644 --- a/ui.go +++ b/ui.go @@ -6,6 +6,8 @@ import ( "strings" "time" + "github.com/dustin/go-humanize" + "github.com/gizak/termui/v3" "github.com/gizak/termui/v3/widgets" ) @@ -15,15 +17,21 @@ const ( ) type UIComponent struct { - header *widgets.Paragraph - footer *widgets.Paragraph + header *widgets.Paragraph + footer *widgets.Paragraph + processes *widgets.Table remoteAddrs *widgets.Table connections *widgets.Table - tableRef []*widgets.Table - grid *termui.Grid - shiftIdx int - mode RenderMode + + packetsPlot *widgets.Plot + bytesPlot *widgets.Plot + connsPlot *widgets.Plot + + tableRef []*widgets.Table + grid *termui.Grid + shiftIdx int + mode RenderMode } type RenderMode uint8 @@ -31,7 +39,7 @@ type RenderMode uint8 const ( RModeBytes RenderMode = iota RModePackets - RModePlot + RModeProcess ) func NewUIComponent(mode RenderMode) *UIComponent { @@ -110,6 +118,7 @@ func newTable(title string) *widgets.Table { table.TextAlignment = termui.AlignLeft table.TextStyle = termui.NewStyle(termui.ColorClear) table.BorderStyle = termui.NewStyle(termui.ColorClear) + table.VerticalLine = ' ' table.RowStyles = map[int]termui.Style{0: termui.NewStyle(termui.ColorCyan)} return table }