From 2b53b245d10ea916e1a8f1c3e5861937c46d6099 Mon Sep 17 00:00:00 2001
From: Michael Kohl <citizen428@gmail.com>
Date: Fri, 13 Jan 2012 18:36:03 +0100
Subject: [PATCH] Lots of improvements

 + add generator.path argument
 + add generator.redirects argument
 + restructure redirect handling
 + improve redirect pattern
 + update documentation
 + add changelog
---
 nmap/nse/generator.nse |   54 ++++++++++++++++++++++++++++++++++++++---------
 1 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/nmap/nse/generator.nse b/nmap/nse/generator.nse
index 11db97e..6c2824a 100644
--- a/nmap/nse/generator.nse
+++ b/nmap/nse/generator.nse
@@ -8,25 +8,64 @@ categories = {"discovery", "safe"}
 
 ---
 -- @usage
--- nmap -p 80,443 --script generator <host>
+-- nmap --script generator [--script-args generator.path=<path>,generator.redirects=<number>,...] <host>
+--
 -- @output
 -- PORT    STATE SERVICE
 -- 80/tcp  open  http
 -- |_generator: TYPO3 4.2 CMS
 -- 443/tcp open  https
 -- |_generator: TYPO3 4.2 CMS
+--
+-- @args generator.path Specify the path you want to check for a generator meta tag (default to '/').
+-- @args generator.redirects Specify the maximum number of redirects to follow (defaults to 3).
+
+--- Changelog:
+-- 2011-12-23 Michael Kohl <citizen428@gmail.com>:
+--   + Initial version
+-- 2012-01-10 Michael Kohl <citizen428@gmail.com>:
+--   + update documentation
+--   + make pattern case insensitive
+--   + only follow first redirect
+-- 2012-01-11 Michael Kohl <citizen428@gmail.com>:
+--   + more generic pattern
+--   + simplified matching
+-- 2012-01-13 Michael Kohl <citizen428@gmail.com>:
+--   + add generator.path argument
+--   + add generator.redirects argument
+--   + restructure redirect handling
+--   + improve redirect pattern
+--   + update documentation
+--   + add changelog
 
 --- TODO:
--- add arg for web path
--- add arg for maximum number of redirects
+-- more generic generator pattern
 
 require('http')
 require('shortport')
+require('stdnse')
+
+-- helper function
+local follow_redirects = function(host, port, path, n)
+   local pattern = "^[hH][tT][tT][pP]/1.[01] 30[12]"
+   local response = http.get(host, port, path)
+
+   while response['status-line']:match(pattern) and n > 0 do
+      n = n - 1
+      loc = response.header['location']
+      response = http.get_url(loc)
+   end
+
+   return response
+end
 
 portrule = shortport.http
 
 action = function(host, port)
    local response, loc, generator
+   local path = stdnse.get_script_args('generator.path') or '/'
+   local redirects = tonumber(stdnse.get_script_args('generator.redirects')) or 3
+
    -- Worst case: <meta name=Generator content="Microsoft Word 11">
    local pattern = '<meta name="?generator"? content="([^\"]*)" ?/?>'
 
@@ -36,14 +75,7 @@ action = function(host, port)
                                               string.upper(c))
              end)
 
-   response = http.get(host, port, '/')
-
-   -- deals with only one redirect
-   if response['status-line']:lower():match("^http/1.1 30[12]") then
-      loc = response.header['location']
-      response = http.get_url(loc)
-   end
-
+   response = follow_redirects(host, port, path, redirects)
    return response.body:match(pattern)
 
 end
-- 
1.7.8.3