local bit    = require 'bit'
local smb    = require 'smb'
local stdnse = require 'stdnse'
local tab    = require 'tab'
local table = require "table"
local openssl= stdnse.silent_require 'openssl'

description = [[
Attempts to retrieve useful information about files shared on SMB volumes.
The output is intended to resemble the output of the UNIX <code>ls</code> command.
]]

---
-- @usage
-- nmap -p 445 <ip> --script smb-ls --script-args 'share=c$,path=\temp'
--
-- @output
-- Host script results:
-- | smb-ls:
-- |   Directory of \\192.168.56.101\c$\
-- |     2007-12-02 00:20:09  0      AUTOEXEC.BAT
-- |     2007-12-02 00:20:09  0      CONFIG.SYS
-- |     2007-12-02 00:53:39  <DIR>  Documents and Settings
-- |     2009-09-08 13:26:10  <DIR>  e5a6b742d36facb19c5192852c43
-- |     2008-12-01 02:06:29  <DIR>  Inetpub
-- |     2007-02-18 00:31:38  94720  msizap.exe
-- |     2007-12-02 00:55:01  <DIR>  Program Files
-- |     2008-12-01 02:05:52  <DIR>  temp
-- |     2011-12-16 14:40:18  <DIR>  usr
-- |     2007-12-02 00:42:40  <DIR>  WINDOWS
-- |_    2007-12-02 00:22:38  <DIR>  wmpub
--
-- @args smb-ls.share the share to connect to
-- @args smb-ls.path the path, relative to the share to list the contents from
-- @args smb-ls.pattern [optional] the search pattern to execute (default: *)
-- @args smb-ls.maxdepth [optional] the maximum depth to recurse into a directory
-- @args smb-ls.maxfiles [optional] return only a certain amount of files
-- @args smb-ls.checksum [optional] download each file and calculate a SHA1 checksum
--

author = "Patrik Karlsson"
license = "Same as Nmap--See http://nmap.org/book/man-legal.html"
categories = {"discovery", "safe"}

local arg_share    = stdnse.get_script_args(SCRIPT_NAME .. '.share')
local arg_path     = stdnse.get_script_args(SCRIPT_NAME .. '.path')
local arg_pattern  = stdnse.get_script_args(SCRIPT_NAME .. '.pattern') or '*'
local arg_maxfiles = tonumber(stdnse.get_script_args(SCRIPT_NAME .. '.maxfiles'))
local arg_maxdepth = tonumber(stdnse.get_script_args(SCRIPT_NAME .. '.maxdepth'))
local arg_checksum = stdnse.get_script_args(SCRIPT_NAME .. '.checksum')

hostrule = function(host)
  return ( smb.get_port(host) ~= nil and arg_share and arg_path )
end

-- checks whether the file entry is a directory
local function is_dir(fe)
  return ( bit.band(fe.attrs, 16) == 16 )
end

local function fail(err) return ("\n  ERROR: %s"):format(err or "") end

action = function(host)
  local status, smbstate = smb.start_ex(host, true, true, arg_share, nil, nil, nil)
  if ( not(status) ) then
    return fail("Failed to authenticate to server (" .. smbstate .. ")")
  end

  -- remove leading slash
  arg_path = ( arg_path:sub(1,2) == '\\' and arg_path:sub(2) or arg_path )

  -- fixup checksum argument
  arg_checksum = ( arg_checksum == 'true' or arg_checksum == '1' ) and true or false

  local options = { max_depth = arg_maxdepth, max_files = arg_maxfiles }
  local depth, path, output, dirs = 0, arg_path, {}, {}
  local file_count, dir_count, total_bytes = 0, 0, 0

  repeat
    local lstab = tab.new((arg_checksum and 4 or 3))

    for fe in smb.find_files(smbstate, path .. '\\' .. arg_pattern, options ) do
      if ( arg_checksum and not(is_dir(fe)) ) then
        local status, content = smb.file_read(host, arg_share, path .. '\\' .. fe.fname, nil, {file_create_disposition=1})
        local sha1 = ( status and stdnse.tohex(openssl.sha1(content)) or "" )
        tab.addrow(lstab, fe.created, (is_dir(fe) and '<DIR>' or fe.eof), fe.fname, sha1)
      else
        tab.addrow(lstab, fe.created, (is_dir(fe) and '<DIR>' or fe.eof), fe.fname)
      end

      arg_maxfiles = ( arg_maxfiles and arg_maxfiles - 1 )
      if ( arg_maxfiles == 0 ) then
        break
      end

      if ( is_dir(fe) ) then
        dir_count = dir_count + 1
        if ( fe.fname ~= '.' and fe.fname ~= '..' ) then
          table.insert(dirs, { depth = depth + 1, path = path .. '\\' .. fe.fname } )
        end
      else
        total_bytes = total_bytes + fe.eof
        file_count = file_count + 1
      end
    end
    table.insert(output, { name = ("Directory of %s"):format( '\\\\' .. stdnse.get_hostname(host) .. '\\' .. arg_share .. path), tab.dump(lstab) })

    path = nil
    if ( #dirs ~= 0 ) then
      local dir = table.remove(dirs, 1)
      depth = dir.depth
      if ( not(arg_maxdepth) or ( dir.depth < arg_maxdepth ) ) then
        path = dir.path
        table.insert(output, "")
      end
    end
  until(not(path) or arg_maxfiles == 0)

  smb.stop(smbstate)

  local summary = { name = "Total Files Listed:",
    ("%8d File(s)\t%d bytes"):format(file_count, total_bytes),
    ("%8d Dir(s)"):format(dir_count) }
  table.insert(output, "")
  table.insert(output, summary)

  return stdnse.format_output(true, output)
end