Fixed some whitespace|formatting and made a small improvement to the iterator.

nse_main.lua

@@ -714,11 +714,11 @@ local function run (threads_iter, scantype)
     while threads_iter and num_threads < CONCURRENCY_LIMIT do
       local thread = threads_iter()
       if not thread then
-	threads_iter = nil
+        threads_iter = nil;
-	break
+        break;
       end
       all[thread.co], running[thread.co], total = thread, thread, total+1;
-      num_threads = num_threads + 1
+      num_threads = num_threads + 1;
       thread:start(timeouts);
     end
 
@@ -748,13 +748,12 @@ local function run (threads_iter, scantype)
     -- Checked for timed-out scripts and hosts.
     for co, thread in pairs(waiting) do
       if thread:timed_out() then
-        waiting[co], all[co] = nil, nil;
+        waiting[co], all[co], num_threads = nil, nil, num_threads-1;
         thread:d("%THREAD %stimed out", thread.host
             and format("%s%s ", thread.host.ip,
                     thread.port and ":"..thread.port.number or "")
             or "");
         thread:close(timeouts, "timed out");
-        num_threads = num_threads - 1;
       end
     end
 
@@ -764,21 +763,20 @@ local function run (threads_iter, scantype)
 
       local s, result = resume(co, unpack(thread.args, 1, thread.args.n));
       if not s then -- script error...
-        all[co] = nil;
+        all[co], num_threads = nil, num_threads-1;
         thread:d("%THREAD_AGAINST threw an error!\n%s\n",
             traceback(co, tostring(result)));
         thread:close(timeouts, result);
-        num_threads = num_threads - 1;
       elseif status(co) == "suspended" then
         if result == NSE_YIELD_VALUE then
           waiting[co] = thread;
         else
+          all[co], num_threads = nil, num_threads-1;
           thread:d("%THREAD yielded unexpectedly and cannot be resumed.");
           thread:close();
-          num_threads = num_threads - 1;
         end
       elseif status(co) == "dead" then
-        all[co] = nil;
+        all[co], num_threads = nil, num_threads-1;
         if type(result) == "string" then
           -- Escape any character outside the range 32-126 except for tab,
           -- carriage return, and line feed. This makes the string safe for
@@ -790,7 +788,6 @@ local function run (threads_iter, scantype)
         end
         thread:d("Finished %THREAD_AGAINST.");
         thread:close(timeouts);
-        num_threads = num_threads - 1;
       end
       current = nil;
     end
@@ -952,77 +949,73 @@ local function main (hosts, scantype)
 
   local runlevels = {};
   for i, script in ipairs(chosen_scripts) do
-    if not runlevels[script.runlevel] then
+    runlevels[script.runlevel] = runlevels[script.runlevel] or {};
-      runlevels[script.runlevel] = script.runlevel;
+    insert(runlevels[script.runlevel], script);
-    end
   end
   sort(runlevels);
 
   -- Yield only scripts in the given runlevel.
-  local function runlevel_scripts(chosen_scripts, runlevel)
+  local function runlevel_scripts(scripts)
     return wrap(function ()
-      for i, script in ipairs(chosen_scripts) do
+      for i, script in ipairs(scripts) do
-        if script.runlevel == runlevel then
+        yield(script);
-          yield(script)
-        end
       end
-    end)
+    end);
   end
 
   -- This iterator is passed to the run function. It returns one new script
   -- thread on demand until exhausted.
-  local function threads_iter()
+  local function threads_iter ()
     -- activate prerule scripts
-    for _, runlevel in ipairs(runlevels) do
+    for runlevel, scripts in ipairs(runlevels) do
       print_verbose(1, "Starting runlevel %u (of %u) scan.", runlevel, #runlevels);
       if (scantype == NSE_PRE_SCAN) then
-	print_verbose(1, "Script Pre-scanning.");
+        print_verbose(1, "Script Pre-scanning.");
-	for script in runlevel_scripts(chosen_scripts, runlevel) do
+        for script in runlevel_scripts(scripts) do
-	  local thread = script:new_thread("prerule");
+           local thread = script:new_thread("prerule");
-	  if thread then
+           if thread then
-	    thread.args = {n = 0};
+             thread.args = {n = 0};
-	    yield(thread);
+             yield(thread);
-	  end
+           end
-	end
+        end
       -- activate hostrule and portrule scripts
       elseif (scantype == NSE_SCAN) then
-	if #hosts > 1 then
+        if #hosts > 1 then
-	  print_verbose(1, "Script scanning %d hosts.", #hosts);
+          print_verbose(1, "Script scanning %d hosts.", #hosts);
-	elseif #hosts == 1 then
+        elseif #hosts == 1 then
-	  print_verbose(1, "Script scanning %s.", hosts[1].ip);
+          print_verbose(1, "Script scanning %s.", hosts[1].ip);
-	end
+        end
 
-	-- Check hostrules for this host.
+        -- Check hostrules for this host.
-	for j, host in ipairs(hosts) do
+        for j, host in ipairs(hosts) do
-	  for script in runlevel_scripts(chosen_scripts, runlevel) do
+          for script in runlevel_scripts(scripts) do
-	    local thread = script:new_thread("hostrule", tcopy(host));
+            local thread = script:new_thread("hostrule", tcopy(host));
-	    if thread then
+            if thread then
-	      thread.args, thread.host = {n = 1, tcopy(host)}, host;
+              thread.args, thread.host = {n = 1, tcopy(host)}, host;
-	      yield(thread);
+              yield(thread);
-	    end
+            end
-	  end
+          end
-	  -- Check portrules for this host.
+          -- Check portrules for this host.
-	  for port in cnse.ports(host) do
+          for port in cnse.ports(host) do
-	    for script in runlevel_scripts(chosen_scripts, runlevel) do
+            for script in runlevel_scripts(scripts) do
-	      local thread = script:new_thread("portrule", tcopy(host), tcopy(port));
+              local thread = script:new_thread("portrule", tcopy(host), tcopy(port));
-	      if thread then
+              if thread then
-		thread.args, thread.host, thread.port =
+                thread.args, thread.host, thread.port = {n = 2, tcopy(host), tcopy(port)}, host, port;
-		    {n = 2, tcopy(host), tcopy(port)}, host, port;
+                yield(thread);
-		yield(thread);
+              end
-	      end
+            end
-	    end
+          end
-	  end
+        end
-	end
+        -- activate postrule scripts
-      -- activate postrule scripts
       elseif (scantype == NSE_POST_SCAN) then
-	print_verbose(1, "Script Post-scanning.");
+        print_verbose(1, "Script Post-scanning.");
-	for script in runlevel_scripts(chosen_scripts, runlevel) do
+        for script in runlevel_scripts(scripts) do
-	  local thread = script:new_thread("postrule");
+          local thread = script:new_thread("postrule");
-	  if thread then
+          if thread then
-	    thread.args = {n = 0};
+            thread.args = {n = 0};
-	    yield(thread);
+            yield(thread);
-	  end
+          end
-	end
+        end
       end
     end
   end