local nmap = require "nmap"
local shortport = require "shortport"
local stdnse = require "stdnse"
local string = require "string"
local http = require "http"
local bin = require "bin"

description = [[
Gathers info from the Metasploit rpc service.  It requires a valid login pair.
After authentication it tries to determine Metasploit version and deduce the OS
type.  Then it creates a new console and executes few commands to get
additional info.

References:
* http://wiki.msgpack.org/display/MSGPACK/Format+specification
* https://community.rapid7.com/docs/DOC-1516 Metasploit RPC API Guide
]]

---
--@usage
-- nmap <target> --script=metasploit-info --script-args username=root,password=root
--@output
-- 55553/tcp open  metasploit-msgrpc syn-ack
-- | metasploit-info:
-- |   Metasploit version: 4.4.0-dev Ruby version: 1.9.3 i386-mingw32 2012-02-16 API version: 1.0
-- |   Additional info:
-- |   Host Name:                 WIN
-- |   OS Name:                   Microsoft Windows XP Professional
-- |   OS Version:                5.1.2600 Service Pack 3 Build 2600
-- |   OS Manufacturer:           Microsoft Corporation
-- |   OS Configuration:          Standalone Workstation
-- |   OS Build Type:             Uniprocessor Free
-- |  ..... lots of other info ....
-- |   Domain:                    WORKGROUP
-- |_  Logon Server:              \\BLABLA
--
-- @args metasploit-info.username  Valid metasploit rpc username (required)
-- @args metasploit-info.password  Valid metasploit rpc password (required)
-- @args metasploit-info.command   Custom command to run on the server (optional)



author = "Aleksandar Nikolic"
license = "Same as Nmap--See https://nmap.org/book/man-legal.html"
categories = {"intrusive","safe"}

portrule = shortport.port_or_service(55553,"metasploit-msgrpc")
local arg_username = stdnse.get_script_args(SCRIPT_NAME .. ".username")
local arg_password = stdnse.get_script_args(SCRIPT_NAME .. ".password")
local arg_command  = stdnse.get_script_args(SCRIPT_NAME .. ".command")
local os_type

-- returns a "prefix" that msgpack uses for strings
local get_prefix = function(data)
  if string.len(data) <= 31 then
    return bin.pack("C",0xa0 + string.len(data))
  else
    return "\xda"  .. bin.pack("s",string.len(data))
  end

end

-- returns a msgpacked data for console.read
local encode_console_read = function(method,token, console_id)
  return "\x93" .. get_prefix(method) .. method .. "\xda\x00\x20" .. token .. get_prefix(console_id) .. console_id
end

-- returns a msgpacked data for console.write
local encode_console_write = function(method, token, console_id, command)
  return "\x94" .. get_prefix(method) .. method .. "\xda\x00\x20" .. token .. get_prefix(console_id) .. console_id .. get_prefix(command) .. command
end

-- returns a msgpacked data for auth.login
local encode_auth = function(username, password)
  local method = "auth.login"
  return "\x93\xaa" .. method .. get_prefix(username) .. username .. get_prefix(password) .. password
end

-- returns a msgpacked data for any method without extra parameters
local encode_noparam = function(token,method)
  -- token is always the same length
  return "\x92" .. get_prefix(method) .. method .. "\xda\x00\x20" .. token
end

-- does the actual call with specified, pre-packed data
-- and returns the response
local msgrpc_call = function(host, port, msg)
  local data
  local options = {
    header = {
      ["Content-Type"] = "binary/message-pack"
    }
  }
  data = http.post(host,port, "/api/",options, nil , msg)
  if data and data.status and tostring( data.status ):match( "200" )  then
    return data.body
  end
  return nil
end

-- auth.login wrapper, returns the auth token
local login = function(username, password,host,port)

  local data  = msgrpc_call(host, port, encode_auth(username,password))

  if data then
    local start = string.find(data,"success")
    if  start > -1 then
      -- get token
      local token = string.sub(string.sub(data,start),17) -- "manually" unpack token
      return true, token
    else
      return false, nil
    end
  end
  stdnse.debug1("something is wrong:" .. data )
  return false, nil
end

-- core.version wrapper, returns version info, and sets the OS type
-- so we can decide which commands to send later
local get_version = function(host, port, token)
  local msg = encode_noparam(token,"core.version")

  local data = msgrpc_call(host, port, msg)
  -- unpack data
  if data then
    -- get version, ruby version, api version
    local start = string.find(data,"version")
    local metasploit_version
    local ruby_version
    local api_version
    if start then
      metasploit_version = string.sub(string.sub(data,start),9)
      start = string.find(metasploit_version,"ruby")
      start = start - 2
      metasploit_version = string.sub(metasploit_version,1,start)
      start = string.find(data,"ruby")
      ruby_version = string.sub(string.sub(data,start),6)
      start = string.find(ruby_version,"api")
      start = start - 2
      ruby_version = string.sub(ruby_version,1,start)
      start = string.find(data,"api")
      api_version = string.sub(string.sub(data,start),5)
      -- put info in a table and parse for OS detection and other info
      port.version.name = "metasploit-msgrpc"
      port.version.product = metasploit_version
      port.version.name_confidence = 10
      nmap.set_port_version(host,port)
      local info = "Metasploit version: " .. metasploit_version .. " Ruby version: " .. ruby_version .. " API version: " .. api_version
      if string.find(ruby_version,"mingw") < 0 then
        os_type = "linux" -- assume linux for now
      else -- mingw compiler means it's a windows build
        os_type = "windows"
      end
      stdnse.debug1("%s", info)
      return info
    end
  end
  return nil
end

-- console.create wrapper, returns console_id
-- which we can use to interact with metasploit further
local create_console = function(host,port,token)
  local msg = encode_noparam(token,"console.create")
  local data = msgrpc_call(host, port, msg)
  -- unpack data
  if data then
    --get console id
    local start = string.find(data,"id")
    local console_id
    if start then
      console_id = string.sub(string.sub(data,start),4)
      local next_token = string.find(console_id,"prompt")
      console_id = string.sub(console_id,1,next_token-2)
      return console_id
    end
  end
  return nil

end

-- console.read wrapper
local read_console = function(host,port,token,console_id)
  local msg = encode_console_read("console.read",token,console_id)
  local data = msgrpc_call(host, port, msg)
  -- unpack data
  if data then
    -- check if busy
    while string.byte(data,string.len(data)) == 0xc3 do
      -- console is busy , let's retry in one second
      stdnse.sleep(1)
      data = msgrpc_call(host, port, msg)
    end
    local start = string.find(data,"data")
    local read_data
    if start then
      read_data = string.sub(string.sub(data,start),8)
      local next_token = string.find(read_data,"prompt")
      read_data = string.sub(read_data,1,next_token-2)
      return read_data
    end
  end
end

-- console.write wrapper
local write_console = function(host,port,token,console_id,command)
  local msg = encode_console_write("console.write",token,console_id,command .. "\n")
  local data = msgrpc_call(host, port, msg)
  -- unpack data
  if data then
    return true
  end
  return false
end

-- console.destroy wrapper, just to be nice, we don't want console to hang ...
local destroy_console = function(host,port,token,console_id)
  local msg = encode_console_read("console.destroy",token,console_id)
  local data = msgrpc_call(host, port, msg)
end

-- write command and read result helper
local write_read_console = function(host,port,token, console_id,command)
  if write_console(host,port,token,console_id, command) then
    local read_data = read_console(host,port,token,console_id)
    if read_data then
      read_data = string.sub(read_data,string.find(read_data,"\n")+1) -- skip command echo
      return read_data
    end
  end
  return nil
end

action = function( host, port )
  if not arg_username or not arg_password then
    stdnse.debug1("This script requires username and password supplied as arguments")
    return false
  end

  -- authenticate
  local status, token = login(arg_username,arg_password,host,port)
  if  status then
    -- get version info
    local info = get_version(host,port,token)
    local console_id = create_console(host,port,token)
    if console_id then
      local read_data = read_console(host,port,token,console_id) -- first read the banner/ascii art
      stdnse.debug2("%s", read_data) -- print the nice looking banner if dbg level high enough :)
      if read_data then
        if os_type == "linux" then
          read_data = write_read_console(host,port,token,console_id, "uname -a")
          if read_data then
            info = info .. "\nAdditional info: " ..  read_data
          end
          read_data = write_read_console(host,port,token,console_id, "id")
          if read_data then
            info = info ..  read_data
          end
        elseif os_type == "windows" then
          read_data = write_read_console(host,port,token,console_id, "systeminfo")
          if read_data then
            stdnse.debug2("%s", read_data) -- print whole info if dbg level high enough
            local stop = string.find(read_data,"Hotfix") -- trim data down , systeminfo return A LOT
            read_data = string.sub(read_data,1,stop-2)
            info = info .. "\nAdditional info: \n" ..  read_data
          end
        end
        if arg_command then
          read_data = write_read_console(host,port,token,console_id, arg_command)
          if read_data then
            info = info .. "\nCustom command output: " ..  read_data
          end
        end
        if read_data then
          -- let's be nice and close the console
          destroy_console(host,port,token,console_id)
        end
      end
    end
    if info then
      return stdnse.format_output(true,info)
    end
  end
  return false
end
