Nmap/scripts: add Hextile support to vnc-screenshot

nmap_scripts/vnc-screenshot.nse

@@ -14,8 +14,10 @@
 -- You should have received a copy of the GNU General Public License
 -- along with IVRE. If not, see <http://www.gnu.org/licenses/>.
 
-local shortport = require "shortport"
 local match = require "match"
+local math = require "math"
+local shortport = require "shortport"
+local stdnse = require "stdnse"
 
 description = [[
 
@@ -49,6 +51,15 @@ local function missing_data(buffer, fb)
    return false
 end
 
+local function has_data(buffer, fb)
+   for i=1, fb.height do
+      for j=1, fb.width do
+         if buffer[i][j] ~= nil then return true end
+      end
+   end
+   return false
+end
+
 local function read_pixel(socket, fb)
   local status, data
   local pix = {}
@@ -137,10 +148,14 @@ action = function(host, port)
   fb.desktop_name_len = (">I2I2BBBBI2I2I2BBBc3I4"):unpack(result)
   fb.bytes_per_pixel = fb.bytes_per_pixel // 8
   status, fb.desktop_name = socket:receive_buf(match.numbytes(fb.desktop_name_len), true)
+  if not status then
+    socket:close()
+    return
+  end
 
   socket:send('\000\000\000\000' .. result:sub(5, 17) .. '\000\000\000' ..
-                '\002\000\000\001\000\000\000\000\003\000\000\000\000\000' ..
-                 result:sub(1, 4))
+		'\002\000\000\002\000\000\000\000\000\000\000\005\005\000' ..
+		'\000\000\000\000\003\000\000\000\000\000' .. result:sub(1, 4))
 
   local buffer = {}
   for i = 1, fb.height do
@@ -152,6 +167,9 @@ action = function(host, port)
 
   while missing_data(buffer, fb) do
     status, result = socket:receive_buf(match.numbytes(4), true)
+    if not status then
+      goto draw
+    end
     if result:sub(1, 1) == '\001' then
       status, result = socket:receive_buf(match.numbytes(2), true)
       if not status then
@@ -173,43 +191,115 @@ action = function(host, port)
         goto draw
       end
       local rect = {}
-      rect.xpos, rect.ypos, rect.width, rect.height, encoding = (">I2I2I2I2I4"):unpack(result)
-      if encoding ~= 0 then
-        goto draw
+      rect.xpos, rect.ypos, rect.width, rect.height, rect.encoding = (">I2I2I2I2I4"):unpack(result)
+      if buffer[rect.ypos + 1][rect.xpos + 1] ~= nil then
+	stdnse.debug1("Overriding data: screenshot complete?")
+	goto draw
       end
-      for ih = 1, rect.height do
-        for iw = 1, rect.width do
-          pixel = read_pixel(socket, fb)
-          if pixel == nil then
-            goto draw
-          end
-          buffer[rect.ypos + ih][rect.xpos + iw] = pixel
-        end
+      if rect.encoding == 0 then
+	local pixel
+	for ih = 1, rect.height do
+	  for iw = 1, rect.width do
+	    pixel = read_pixel(socket, fb)
+	    if pixel == nil then
+	      goto draw
+	    end
+	    buffer[rect.ypos + ih][rect.xpos + iw] = pixel
+	  end
+	end
+      elseif rect.encoding == 5 then
+	local flags, tx, ty, back_col, fore_col, nsubrects, pix, subr_x, subr_y, subr_w, subr_h
+	for ty = 0, math.ceil(rect.height / 16) - 1 do
+	  for tx = 0, math.ceil(rect.width / 16) - 1 do
+	    status, data = socket:receive_buf(match.numbytes(1), true)
+	    if not status then
+	      goto draw
+	    end
+	    flags = data:byte()
+	    if flags & 1 == 1 then
+	      for ih = 1, math.min(rect.height - ty * 16, 16) do
+		for iw = 1, math.min(rect.width - tx * 16, 16) do
+		  buffer[rect.ypos + ty * 16 + ih][rect.xpos + tx * 16 + iw] = read_pixel(socket, fb)
+		end
+	      end
+	      goto nexttile
+	    end
+	    if flags & 2 == 2 then
+	      back_col = read_pixel(socket, fb)
+	    end
+	    if flags & 4 == 4 then
+	      fore_col = read_pixel(socket, fb)
+	    end
+	    for ih = 1, math.min(rect.height - ty * 16, 16) do
+	      for iw = 1, math.min(rect.width - tx * 16, 16) do
+		 if buffer[rect.ypos + ty * 16 + ih][rect.xpos + tx * 16 + iw] == nil then
+		   buffer[rect.ypos + ty * 16 + ih][rect.xpos + tx * 16 + iw] = back_col
+		end
+	      end
+	    end
+	    if flags & 8 == 8 then
+	      status, nsubrects = socket:receive_buf(match.numbytes(1), true)
+	      if not status then
+		goto draw
+	      end
+	      nsubrects = nsubrects:byte()
+	      for _ = 1, nsubrects do
+		pix = (flags & 16) == 16 and read_pixel(socket, fb) or fore_col
+		status, data = socket:receive_buf(match.numbytes(1), true)
+		if not status then
+		  goto draw
+		end
+		subr_x = data:byte()
+		subr_y = subr_x % 16
+		subr_x = subr_x >> 4
+		status, data = socket:receive_buf(match.numbytes(1), true)
+		if not status then
+		  goto draw
+		end
+		subr_w = data:byte()
+		subr_h = (subr_w % 16) + 1
+		subr_w = (subr_w >> 4) + 1
+		for ih = 1, subr_h do
+		  for iw = 1, subr_w do
+		    buffer[rect.ypos + ty * 16 + subr_y + ih][rect.xpos + tx * 16 + subr_x + iw] = pix
+		  end
+		end
+	      end
+	    end
+	    ::nexttile::
+	  end
+	end
+      else
+	stdnse.debug1("Unsupported encoding (%d)", rect.encoding)
+        goto draw
       end
     end
     ::next::
   end
 
   ::draw::
   socket:close()
-  local f = assert(io.popen(
-     ("convert -size %dx%d -depth 8 RGB:- %s"):format(
-       fb.width, fb.height, fname), "w"
-  ))
-  local pixel
-  for i = 1, fb.height do
-    for j = 1, fb.width do
-      pixel = buffer[i][j]
-      if pixel == nil then
-        f:write("\000\000\000")
-      else
-        f:write(buffer[i][j])
+  if has_data(buffer, fb) then
+    local f = assert(io.popen(
+      ("convert -size %dx%d -depth 8 RGB:- %s"):format(
+         fb.width, fb.height, fname), "w"
+    ))
+    local pixel
+    for i = 1, fb.height do
+      for j = 1, fb.width do
+	pixel = buffer[i][j]
+	if pixel == nil then
+	  f:write("\000\000\000")
+	else
+	  f:write(buffer[i][j])
+	end
       end
     end
-  end
-  f:close()
-
-  if os.rename(fname, fname) then
-    return ("Saved to %s"):format(fname)
+    f:close()
+    if os.rename(fname, fname) then
+      return ("Saved to %s"):format(fname)
+    end
+  else
+    stdnse.debug1("No data: empty screenshot discarded")
   end
 end