From 7b4fc3d31d984d0e355ed002f3acf60e6d74387f Mon Sep 17 00:00:00 2001
From: Simone Margaritelli <evilsocket@gmail.com>
Date: Sat, 21 Sep 2024 17:52:50 +0200
Subject: [PATCH] misc: small fix or general refactoring i did not bother
 commenting

---
 modules/zerogod/zerogod.go           | 31 +++++++++++++++++++++++++++-
 modules/zerogod/zerogod_advertise.go |  2 ++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/modules/zerogod/zerogod.go b/modules/zerogod/zerogod.go
index 0444a76d..929371bb 100644
--- a/modules/zerogod/zerogod.go
+++ b/modules/zerogod/zerogod.go
@@ -1,6 +1,11 @@
 package zerogod
 
 import (
+	"fmt"
+	"math/rand"
+	"os"
+	"path/filepath"
+
 	"github.com/bettercap/bettercap/v2/session"
 	"github.com/bettercap/bettercap/v2/tls"
 	"github.com/google/gopacket"
@@ -74,8 +79,32 @@ func NewZeroGod(s *session.Session) *ZeroGod {
 			return mod.startAdvertiser(args[0])
 		}))
 
+	// TODO: add autocomplete
+	mod.AddHandler(session.NewModuleHandler("zerogod.impersonate ADDRESS", "zerogod.impersonate (.+)",
+		"Impersonate ADDRESS by advertising the same discovery information.",
+		func(args []string) error {
+			if address := args[0]; address == "off" {
+				return mod.stopAdvertiser()
+			} else {
+				tmpDir := os.TempDir()
+				tmpFileName := filepath.Join(tmpDir, fmt.Sprintf("impersonate_%d.yml", rand.Int()))
+
+				if err := mod.save(address, tmpFileName); err != nil {
+					return err
+				}
+
+				return mod.startAdvertiser(tmpFileName)
+			}
+		}))
+
 	mod.AddHandler(session.NewModuleHandler("zerogod.advertise off", "",
-		"Start a previously started advertiser.",
+		"Stop a previously started advertiser.",
+		func(args []string) error {
+			return mod.stopAdvertiser()
+		}))
+
+	mod.AddHandler(session.NewModuleHandler("zerogod.impersonate off", "",
+		"Stop a previously started impersonation.",
 		func(args []string) error {
 			return mod.stopAdvertiser()
 		}))
diff --git a/modules/zerogod/zerogod_advertise.go b/modules/zerogod/zerogod_advertise.go
index 0876a973..31187d34 100644
--- a/modules/zerogod/zerogod_advertise.go
+++ b/modules/zerogod/zerogod_advertise.go
@@ -131,7 +131,9 @@ func (mod *ZeroGod) startAdvertiser(fileName string) error {
 	for _, svc := range advertiser.Services {
 		// if no external responder has been specified, check if port is available
 		if svc.Responder == "" {
+			// if the port was not set or is not available or is requesdted by another service
 			for svc.Port == 0 || !isPortAvailable(svc.Port) || isPortRequested(svc, services) {
+				// set a new one and try again
 				newPort := (rand.Intn(65535-1024) + 1024)
 				mod.Warning("port %d for service %s is not avaialble, trying %d ...",
 					svc.Port,