From 63b94197ec7c63f28ffa46db094fe651c2f13a56 Mon Sep 17 00:00:00 2001
From: olebeck <31539311+olebeck@users.noreply.github.com>
Date: Fri, 21 Apr 2023 00:55:52 +0200
Subject: [PATCH] add -capture for all commands

---
 cmd/bedrocktool/main.go | 4 ++--
 handlers/capture.go     | 4 ++++
 subcommands/capture.go  | 3 +--
 utils/proxy.go          | 4 ++++
 utils/utils.go          | 2 +-
 5 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/cmd/bedrocktool/main.go b/cmd/bedrocktool/main.go
index 561f2cb..56bed96 100644
--- a/cmd/bedrocktool/main.go
+++ b/cmd/bedrocktool/main.go
@@ -99,11 +99,11 @@ func main() {
 	flag.BoolVar(&utils.Options.ExtraDebug, "extra-debug", false, locale.Loc("extra_debug", nil))
 	flag.StringVar(&utils.Options.PathCustomUserData, "userdata", "", locale.Loc("custom_user_data", nil))
 	flag.String("lang", "", "lang")
-	flag.BoolVar(&utils.Options.EnableDNS, "dns", false, locale.Loc("enable_dns", nil))
+	flag.BoolVar(&utils.Options.Capture, "capture", false, "Capture Packet log")
 
 	subcommands.Register(subcommands.HelpCommand(), "")
 	subcommands.ImportantFlag("debug")
-	subcommands.ImportantFlag("dns")
+	subcommands.ImportantFlag("capture")
 	subcommands.ImportantFlag("preload")
 	subcommands.HelpCommand()
 
diff --git a/handlers/capture.go b/handlers/capture.go
index b1e26d4..ad700d3 100644
--- a/handlers/capture.go
+++ b/handlers/capture.go
@@ -65,3 +65,7 @@ func NewPacketCapturer() *utils.ProxyHandler {
 		},
 	}
 }
+
+func init() {
+	utils.NewPacketCapturer = NewPacketCapturer
+}
diff --git a/subcommands/capture.go b/subcommands/capture.go
index d5e4d57..72b581f 100644
--- a/subcommands/capture.go
+++ b/subcommands/capture.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"flag"
 
-	"github.com/bedrock-tool/bedrocktool/handlers"
 	"github.com/bedrock-tool/bedrocktool/locale"
 	"github.com/bedrock-tool/bedrocktool/utils"
 )
@@ -34,6 +33,6 @@ func (c *CaptureCMD) Execute(ctx context.Context, ui utils.UI) error {
 		return err
 	}
 	proxy.AlwaysGetPacks = true
-	proxy.AddHandler(handlers.NewPacketCapturer())
+	utils.Options.Capture = true
 	return proxy.Run(ctx, address, hostname)
 }
diff --git a/utils/proxy.go b/utils/proxy.go
index 3a03fc3..4c9f799 100644
--- a/utils/proxy.go
+++ b/utils/proxy.go
@@ -282,6 +282,7 @@ func (p *ProxyContext) IsClient(addr net.Addr) bool {
 }
 
 var NewDebugLogger func(bool) *ProxyHandler
+var NewPacketCapturer func() *ProxyHandler
 
 func (p *ProxyContext) connectClient(ctx context.Context, serverAddress string, cdpp **login.ClientData) (err error) {
 	GetTokenSource() // ask for login before listening
@@ -337,6 +338,9 @@ func (p *ProxyContext) Run(ctx context.Context, serverAddress, name string) (err
 	if Options.Debug || Options.ExtraDebug {
 		p.AddHandler(NewDebugLogger(Options.ExtraDebug))
 	}
+	if Options.Capture {
+		p.AddHandler(NewPacketCapturer())
+	}
 	p.AddHandler(&ProxyHandler{
 		Name:     "Commands",
 		PacketCB: p.CommandHandlerPacketCB,
diff --git a/utils/utils.go b/utils/utils.go
index c85a0d9..0a8ca7f 100644
--- a/utils/utils.go
+++ b/utils/utils.go
@@ -30,7 +30,7 @@ var Options struct {
 	Preload            bool
 	IsInteractive      bool
 	ExtraDebug         bool
-	EnableDNS          bool
+	Capture            bool
 	PathCustomUserData string
 }