add new flag: -c

mozillazg · Apr 27, 2024 · 1437442 · 1437442
1 parent e750e34
commit 1437442
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/cmd/options.go b/cmd/options.go
@@ -10,6 +10,7 @@ type Options struct {
 	listInterfaces bool
 	version        bool
 	print          bool
+	maxPacketCount uint
 }
 
 func (o Options) WritePath() string {

diff --git a/cmd/root.go b/cmd/root.go
@@ -53,6 +53,7 @@ func init() {
 		"Print the ptcpdump and libpcap version strings and exit")
 	rootCmd.Flags().BoolVar(&opts.print, "print", false,
 		"Print parsed packet output, even if the raw packets are being saved to a file with the -w flag")
+	rootCmd.Flags().UintVarP(&opts.maxPacketCount, "receive-count", "c", 0, "Exit after receiving count packets")
 }
 
 func Execute() error {
@@ -104,7 +105,10 @@ func run(cmd *cobra.Command, args []string) error {
 	execConsumer := consumer.NewExecEventConsumer(pcache)
 	go execConsumer.Start(ctx, execEventReader)
 	packetConsumer := consumer.NewPacketEventConsumer(writers, devices)
-	go packetConsumer.Start(ctx, packetEventReader)
+	go func() {
+		packetConsumer.Start(ctx, packetEventReader, opts.maxPacketCount)
+		stop()
+	}()
 
 	runtime.Gosched()
 

diff --git a/internal/consumer/net.go b/internal/consumer/net.go
@@ -22,7 +22,8 @@ func NewPacketEventConsumer(writers []writer.PacketWriter, devices map[int]dev.D
 	}
 }
 
-func (c *PacketEventConsumer) Start(ctx context.Context, reader *perf.Reader) {
+func (c *PacketEventConsumer) Start(ctx context.Context, reader *perf.Reader, maxPacketCount uint) {
+	var n uint
 	for {
 		select {
 		case <-ctx.Done():
@@ -43,6 +44,12 @@ func (c *PacketEventConsumer) Start(ctx context.Context, reader *perf.Reader) {
 			log.Printf("[PacketEventConsumer] lost samples: %d", record.LostSamples)
 		}
 		c.parsePacketEvent(record.RawSample)
+
+		n++
+		if maxPacketCount > 0 && n == maxPacketCount {
+			log.Printf("%d packets captured", n)
+			break
+		}
 	}
 }