local nmap = require "nmap"
local pgsql = require "pgsql"
local shortport = require "shortport"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"
local unpwdb = require "unpwdb"

local openssl = stdnse.silent_require "openssl"

description = [[
Performs password guessing against PostgreSQL.
]]

---
-- @usage
-- nmap -p 5432 --script pgsql-brute <host>
--
-- @output
-- 5432/tcp open  pgsql
-- | pgsql-brute:
-- |   root:<empty> => Valid credentials
-- |_  test:test => Valid credentials
--
-- @args pgsql.nossl If set to <code>1</code> or <code>true</code>, disables SSL.
-- @args pgsql.version Force protocol version 2 or 3.

-- SSL Encryption
-- --------------
-- We need to handle several cases of SSL support
--  o SSL can be supported on a server level
--  o SSL can be enforced per host or network level
--  o SSL can be denied per host or network level

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


-- Version 0.4
-- Created 01/15/2010 - v0.1 - created by Patrik Karlsson <patrik@cqure.net>
-- Revised 02/20/2010 - v0.2 - moved version detection to pgsql library
-- Revised 03/04/2010 - v0.3 - added code from ssh-hostkey.nse to check for SSL support
--                           - added support for trusted authentication method
-- Revised 09/10/2011 - v0.4 - changed account status text to be more consistent with other *-brute scripts

portrule = shortport.port_or_service(5432, "postgresql")


Driver = {

  new = function(self, host, port)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    o.host = host
    o.port = port
    o.ssl = true

    -- If the user explicitly does not disable SSL, enforce it
    if ((nmap.registry.args['pgsql.nossl'] == 'true') or
        (nmap.registry.args['pgsql.nossl'] == '1' )) then
      o.ssl = false
    end

    -- TODO(claudiu) Make sure we can do this in the constructor.
    if (nmap.registry.args['pgsql.version']) then
      if (tonumber(nmap.registry.args['pgsql.version']) == 2) then
        o.pg = pgsql.v2
      elseif (tonumber(nmap.registry.args['pgsql.version']) == 3) then
        o.pg = pgsql.v3
      else
        stdnse.print_debug("pgsql-brute: Unsupported version %s",
          nmap.registry.args['pgsql.version'])
        return brute.Error:new("Unsupported version %s",
          nmap.registry.args['pgsql.version'])
      end
    else
      o.pg = pgsql.detectVersion(host, port)
    end

    return o
  end,

  connect = function(self)
    self.socket = nmap.new_socket()
    self.socket:set_timeout(5000)

    local status, err = self.socket:connect(self.host, self.port)
    if (not(status)) then
      return false, brute.Error:new("Couldn't connect to host: " .. err)
    end

    if (self.ssl) then
      status = pqsql.requestSSL(self.socket)
      if (status) then
        self.socket:reconnect_ssl()
      end
    end
  end,

  login = function(self, username, password)
    local status, response = self.pg.sendStartup(self.socket,
      username, password)

    if (not(status) and not self.ssl) then
      return false, brute.Error:new("No SSL set")
    end

    -- SSL failed, this can occur due to:
    -- 1. The server does not do SSL
    -- 2. SSL was denied on a per host or network level
    --
    -- Attempt SSL connection
    if (not(status)) then
      self.socket:close()
      self.ssl = true
      self.connect()

      status, response = self.pg.sendStartup(
        self.socket, username, username)
      if (not(status)) then
        stdnse.print_debug(
          "pgsql-brute: Error, sendStartup returned: %s",
          response
        )
        return false, brute.Error:new(
          "Error when accessing db \"%s\" as user \"%s\"",
          username, password
        )
      end
    end

    -- Do not attempt to authenticate if authentication type is trusted
    if (response.authtype ~= pgsql.AuthenticationType.Success) then
      status, response = self.pg.loginRequest(
        socket, response, username, password, response.salt)
    end

    if (status) then
      if not nmap.registry.pgsqlusers then
        nmap.registry.pgsqlusers = {}
      end
      nmap.registry.pgsqlusers[username] = password

      if (response.authtype ~= pgsql.AuthenticationType.Success) then
        return true, brute.Account:new(username, password, creds.State.VALID)
      else
        -- Trusted authentication, no password is needed
        -- TODO(claudiu) Investigate how to handle properly this case,
        -- there is no specific state in creds for this.
        return true, brute.Account:new(username, "", creds.State.VALID)
      end
    else
      return false, brute.Error:new("Bad credentials")
    end
  end,

  disconnect = function(self)
    self.socket:close()
    return true
  end,

  check = function(self)
    return true
  end
}

action = function(host, port)
  local status, result
  local engine = brute.Engine:new(Driver, host, port)

  engine.options.script_name = SCRIPT_NAME
  status, result = engine:start()

  return result
end