From b9ef1458c9d204b2465eac96ed37889f66744d94 Mon Sep 17 00:00:00 2001
From: Ullaakut <harold.kutenberg@gmail.com>
Date: Fri, 22 Jul 2022 19:13:58 +0200
Subject: [PATCH] Add --unique and --script-timeout flags

---
 iflist_test.go |  3 ++-
 nmap.go        | 21 +++++++++++++++++++++
 nmap_test.go   | 23 +++++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/iflist_test.go b/iflist_test.go
index 47466a2..5ecbc0f 100644
--- a/iflist_test.go
+++ b/iflist_test.go
@@ -1,9 +1,10 @@
 package nmap
 
 import (
-	"github.com/stretchr/testify/assert"
 	"net"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestScanner_GetInterfaceList(t *testing.T) {
diff --git a/nmap.go b/nmap.go
index 6aa12b8..3732218 100644
--- a/nmap.go
+++ b/nmap.go
@@ -532,6 +532,17 @@ func WithRandomTargets(randomTargets int) Option {
 	}
 }
 
+// WithUnique makes each address be scanned only once.
+// The default behavior is to scan each address as many times
+// as it is specified in the target list, such as when network
+// ranges overlap or different hostnames resolve to the same
+// address.
+func WithUnique() Option {
+	return func(s *Scanner) {
+		s.args = append(s.args, "--unique")
+	}
+}
+
 /*** Host discovery ***/
 
 // WithListScan sets the discovery mode to simply list the targets to scan and not scan them.
@@ -1075,6 +1086,16 @@ func WithScriptUpdateDB() Option {
 	}
 }
 
+// WithScriptTimeout sets the script timeout.
+func WithScriptTimeout(timeout time.Duration) Option {
+	milliseconds := timeout.Round(time.Nanosecond).Nanoseconds() / 1000000
+
+	return func(s *Scanner) {
+		s.args = append(s.args, "--script-timeout")
+		s.args = append(s.args, fmt.Sprintf("%dms", int(milliseconds)))
+	}
+}
+
 /*** OS Detection ***/
 
 // WithOSDetection enables OS detection.
diff --git a/nmap_test.go b/nmap_test.go
index be64b9c..61fd3b6 100644
--- a/nmap_test.go
+++ b/nmap_test.go
@@ -548,6 +548,17 @@ func TestTargetSpecification(t *testing.T) {
 				"4",
 			},
 		},
+		{
+			description: "unique addresses",
+
+			options: []Option{
+				WithUnique(),
+			},
+
+			expectedArgs: []string{
+				"--unique",
+			},
+		},
 		{
 			description: "target exclusion",
 
@@ -1341,6 +1352,18 @@ func TestScriptScan(t *testing.T) {
 				"--script-updatedb",
 			},
 		},
+		{
+			description: "set script timeout",
+
+			options: []Option{
+				WithScriptTimeout(40 * time.Second),
+			},
+
+			expectedArgs: []string{
+				"--script-timeout",
+				"40000ms",
+			},
+		},
 	}
 
 	for _, test := range tests {