Add TCP support to dns.lua

Slightly modified from patch from John Bond:
http://seclists.org/nmap-dev/2014/q1/118

CHANGELOG

@@ -1,4 +1,7 @@
 # Nmap Changelog ($Id$); -*-text-*-
+
+o Added TCP support to dns.lua. [John Bond]
+
 o Added safe fd_set operations. This makes nmap fail gracefully instead of
   crashing when the number of file descriptors grows over FD_SETSIZE. Jacek
   Wielemborek reported the crash. [Henri Doreau]

nselib/dns.lua

@@ -83,7 +83,7 @@ CLASS = {
 -- @param multiple If true, keep reading multiple responses until timeout.
 -- @return Status (true or false).
 -- @return Response (if status is true).
-local function sendPackets(data, host, port, timeout, cnt, multiple)
+local function sendPacketsUDP(data, host, port, timeout, cnt, multiple)
     local socket = nmap.new_socket("udp")
     local responses = {}
 
@@ -131,6 +131,56 @@ local function sendPackets(data, host, port, timeout, cnt, multiple)
     return false
 end
 
+---
+-- Send TCP DNS query
+-- @param data Data to be sent.
+-- @param host Host to connect to.
+-- @param port Port to connect to.
+-- @param timeout Number of ms to wait for a response.
+-- @return Status (true or false).
+-- @return Response (if status is true).
+local function sendPacketsTCP(data, host, port, timeout)
+    local socket = nmap.new_socket()
+    local response
+    local responses = {}
+    socket:set_timeout(timeout)
+    socket:connect(host, port)
+    -- add payload size we are assuming a minimum size here of 256?
+    local send_data = '\000' .. string.char(#data) .. data
+    socket:send(send_data)
+    local response = ''
+    while true do
+        local status, recv_data = socket:receive_bytes(1)
+        if not status then break end
+        response = response .. recv_data
+    end
+    local status, _, _, ip, _ = socket:get_info()
+    -- remove payload size
+    table.insert(responses, { data = string.sub(response,3), peer = ip } )
+    socket:close()
+    if (#responses>0) then
+        return true, responses
+    end
+    return false
+
+end
+
+---
+-- Call appropriate protocol handeler
+-- @param data Data to be sent.
+-- @param host Host to connect to.
+-- @param port Port to connect to.
+-- @param timeout Number of ms to wait for a response.
+-- @param cnt Number of tries.
+-- @param multiple If true, keep reading multiple responses until timeout.
+-- @return Status (true or false).
+local function sendPackets(data, host, port, timeout, cnt, multiple, proto)
+    if proto == nil or proto == 'udp' then
+        return sendPacketsUDP(data, host, port, timeout, cnt, multiple)
+    else
+        return sendPacketsTCP(data, host, port, timeout)
+    end
+end
 
 ---
 -- Checks if a DNS response packet contains a useful answer.
@@ -265,7 +315,9 @@ end
 function query(dname, options)
     if not options then options = {} end
 
-    local dtype, host, port = options.dtype, options.host, options.port
+    local dtype, host, port, proto = options.dtype, options.host, options.port, options.proto
+    if proto == nil then proto = 'udp' end
+    if port == nil then port = '53' end
 
     local class = options.class or CLASS.IN
     if not options.tries then options.tries = 10 end -- don't get into an infinite loop
@@ -319,7 +371,7 @@ function query(dname, options)
 
     local data = encode(pkt)
 
-    local status, response = sendPackets(data, host, port, options.timeout, options.sendCount, options.multiple)
+    local status, response = sendPackets(data, host, port, options.timeout, options.sendCount, options.multiple, proto)
 
 
     -- if working with know nameservers, try the others

scripts/dns-nsid.nse

@@ -63,7 +63,7 @@ end
 action = function(host, port)
 	local result = stdnse.output_table()
 	local flag = false
-	local status, resp = dns.query("id.server", {host = host.ip, dtype='TXT', class=dns.CLASS.CH, retAll=true, retPkt=true, nsid=true, dnssec=true})
+	local status, resp = dns.query("id.server", {host = host.ip, port=port.number, proto=port.protocol, dtype='TXT', class=dns.CLASS.CH, retAll=true, retPkt=true, nsid=true, dnssec=true})
 	if ( status ) then
 		local status, nsid = rr_filter(resp.add,'OPT')
 		if ( status ) then
@@ -83,7 +83,7 @@ action = function(host, port)
 			result["id.server"] = id_server
 		end
 	end
-	local status, bind_version = dns.query("version.bind", {host = host.ip, dtype='TXT', class=dns.CLASS.CH})
+	local status, bind_version = dns.query("version.bind", {host = host.ip, port=port.number, proto=port.protocol, dtype='TXT', class=dns.CLASS.CH})
 	if ( status ) then
 		flag = true
 		result["bind.version"] = bind_version