Refactor reverse lookups and connection logging

This commit is contained in:
Yun Zheng Hu
2024-10-03 09:33:53 +00:00
parent 3aa16279a4
commit e846b3833c

View File

@@ -4,7 +4,6 @@ import (
"context" "context"
"errors" "errors"
"flag" "flag"
"fmt"
"net" "net"
"os" "os"
"os/exec" "os/exec"
@@ -27,33 +26,6 @@ type PcapClient struct {
totalBytes uint64 totalBytes uint64
} }
func lookupHostnameWithTimeout(addr net.Addr, timeout time.Duration) (string, string, error) {
// Extract the IP address and port from the Addr object
tcpAddr, ok := addr.(*net.TCPAddr)
if !ok {
return "", "", fmt.Errorf("unsupported address type: %T", addr)
}
ip := tcpAddr.IP.String()
port := fmt.Sprintf("%d", tcpAddr.Port)
// Create a new context with the given timeout
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
// Create a new Resolver and perform the IP lookup with the given context
resolver := net.Resolver{}
names, err := resolver.LookupAddr(ctx, ip)
if err != nil {
return "", "", err
}
if len(names) == 0 {
return "", "", fmt.Errorf("no hostnames found for %s", ip)
}
// Return the first IP address found and the original port
return names[0], port, nil
}
var ( var (
pcapCommand = flag.String("cmd", "", "command to execute for pcap data (eg: tcpdump -i eth0 -n --immediate-mode -s 65535 -U -w -)") pcapCommand = flag.String("cmd", "", "command to execute for pcap data (eg: tcpdump -i eth0 -n --immediate-mode -s 65535 -U -w -)")
listenAddress = flag.String("listen", "", "listen address for pcap-over-ip (eg: localhost:4242)") listenAddress = flag.String("listen", "", "listen address for pcap-over-ip (eg: localhost:4242)")
@@ -174,13 +146,16 @@ func main() {
} }
if *noReverseLookup { if *noReverseLookup {
log.Printf("PCAP-over-IP connection from %v", conn.RemoteAddr()) log.Info().Msgf("PCAP-over-IP connection from %v", conn.RemoteAddr())
} else { } else {
ip, port, err := lookupHostnameWithTimeout(conn.RemoteAddr(), 100*time.Millisecond) ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
if err != nil { defer cancel()
log.Printf("PCAP-over-IP connection from %v", conn.RemoteAddr()) ipAddr := conn.RemoteAddr().(*net.TCPAddr).IP.String()
names, _ := net.DefaultResolver.LookupAddr(ctx, ipAddr)
if len(names) == 0 {
log.Info().Msgf("PCAP-over-IP connection from %v", conn.RemoteAddr())
} else { } else {
log.Printf("PCAP-over-IP connection from %s:%s", ip, port) log.Info().Msgf("PCAP-over-IP connection from %v (%v)", conn.RemoteAddr(), names[0])
} }
} }