local base64 = require "base64"
local match = require "match"
local nmap = require "nmap"
local sasl = require "sasl"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"
_ENV = stdnse.module("xmpp", stdnse.seeall)
XML = {
  parse_tag = function(s)
    local _, _, contents, empty, name = string.find(s, "([^<]*)<(/?)([?:%w-]+)")
    local attrs = {}
    if not name then
      return
    end
    for k, v in string.gmatch(s, "%s([%w:]+)='([^']+)'") do
      attrs[k] = v
    end
    for k, v in string.gmatch(s, "%s([%w:]+)=\"([^\"]+)\"") do
      attrs[k] = v
    end
    local finish = (empty ~= "") or (s:sub(#s-1) == '/>')
    return { name = name,
    attrs = attrs,
    start = (empty == ""),
    contents = contents,
    finish = finish }
  end,
}
TagProcessor = {
  ["failure"] = function(socket, tag)
    return TagProcessor["success"](socket,tag)
  end,
  ["success"] = function(socket, tag)
    if ( tag.finish ) then return true end
    local newtag
    repeat
      local status, data = socket:receive_buf(match.pattern_limit(">", 2048), true)
      if ( not(status) ) then
        return false, ("ERROR: Failed to process %s tag"):format(tag.name)
      end
      newtag = XML.parse_tag(data)
    until( newtag.finish and newtag.name == tag.name )
    if ( newtag.name == tag.name ) then return true, tag end
    return false, ("ERROR: Failed to process %s tag"):format(tag.name)
  end,
  ["challenge"] = function(socket, tag)
    local status, data = socket:receive_buf(match.pattern_limit(">", 2048), true)
    if ( not(status) ) then return false, "ERROR: Failed to read challenge tag" end
    local tag = XML.parse_tag(data)
    if ( not(status) or tag.name ~= "challenge" ) then
      return false, "ERROR: Failed to process challenge"
    end
    return status, (tag.contents and base64.dec(tag.contents))
  end,
}
XMPP = {
  new = function(self, host, port, options)
    local o = { host = host,
    port = port,
    options = options or {},
    auth = { mechs = {} } }
    o.options.timeout = o.options.timeout and o.options.timeout or 10
    o.servername = stdnse.get_hostname(host) or o.options.servername
    setmetatable(o, self)
    self.__index = self
    return o
  end,
  send = function(self, data)
    self.socket:set_timeout(1)
    repeat
      local status = self.socket:receive_buf("\0", false)
    until(not(status))
    self.socket:set_timeout(self.options.timeout * 1000)
    return self.socket:send(data)
  end,
  receive_tag = function(self, tag, close)
    local result
    repeat
      local status, data = self.socket:receive_buf(match.pattern_limit(">", 2048), true)
      if ( not(status) ) then return false, data end
      result = XML.parse_tag(data)
    until( ( not(tag) and (close == nil or result.finish == close ) ) or
      ( tag == result.name and ( close == nil or result.finish == close ) ) )
    return true, result
  end,
  connect = function(self, socket)
    assert(self.servername,
    "Cannot connect to XMPP server without valid server name")
    if ( not(self.socket) ) then
      self.socket = socket or nmap.new_socket()
      self.socket:set_timeout(self.options.timeout * 1000)
      local status, err = self.socket:connect(self.host, self.port)
      if ( not(status) ) then
        return false, err
      end
    end
    local data = ("<?xml version='1.0' ?><stream:stream to='%s' xmlns='jabber:client'" ..
    " xmlns:stream='http://etherx.jabber.org/streams'" ..
    " version='1.0'>"):format(self.servername)
    local status, err = self:send(data)
    if ( not(status) ) then return false, "ERROR: Failed to connect to server" end
    local version, start_tls
    repeat
      local status, tag = self:receive_tag()
      if ( not(status) ) then return false, "ERROR: Failed to connect to server" end
      if ( tag.name == "stream:stream" ) then
        version = tag.attrs and tag.attrs.version
      elseif ( tag.name == "starttls" and tag.start ) then
        status, tag = self:receive_tag()
        if ( not(status) ) then
          return false, "ERROR: Failed to connect to server"
        end
        if ( tag.name ~= "starttls" ) then
          start_tls = tag.name
        else
          start_tls = "optional"
        end
      elseif ( tag.name == "mechanism" and tag.finish ) then
        self.auth.mechs[tag.contents] = true
      end
    until(tag.name == "stream:features" and tag.finish)
    if ( version ~= "1.0" ) then
      return false, "ERROR: Only version 1.0 is supported"
    end
    if ( start_tls == "required" or self.options.starttls) then
      status, err = self:send("<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>")
      if ( not(status) ) then return false, "ERROR: Failed to initiate STARTTLS" end
      local status, tag = self:receive_tag()
      if ( not(status) ) then return false, "ERROR: Failed to receive from server" end
      if ( tag.name == "proceed" ) then
        status, err = self.socket:reconnect_ssl()
        self.options.starttls = false
        return self:connect()
      end
    end
    return true
  end,
  login = function(self, username, password, mech)
    assert(mech == "PLAIN" or
    mech == "DIGEST-MD5" or
    mech == "CRAM-MD5" or
    mech == "LOGIN",
    "Unsupported authentication mechanism")
    local auth = ("<auth xmlns='urn:ietf:params:xml:ns:xmpp-sasl' " ..
    "mechanism='%s'/>"):format(mech)
    local realm
    if ( username:match("@") ) then
      username, realm = username:match("^(.*)@(.*)$")
    end
    local status, result
    if ( mech == "PLAIN" ) then
      local mech_params = { username, password }
      local auth_data = sasl.Helper:new(mech):encode(table.unpack(mech_params))
      auth = ("<auth xmlns='urn:ietf:params:xml:ns:xmpp-sasl' " ..
      "mechanism='%s'>%s</auth>"):format(mech, base64.enc(auth_data))
      status, result = self.socket:send(auth)
      if ( not(status) ) then return false, "ERROR: Failed to send SASL PLAIN authentication" end
      status, result = self:receive_tag()
      if ( not(status) ) then return false, "ERROR: Failed to receive login response" end
      if ( result.name == "failure" ) then
        status = TagProcessor[result.name](self.socket, result)
      end
    else
      local status, err = self.socket:send(auth)
      if(not(status)) then return false, "ERROR: Failed to initiate SASL login" end
      local chall
      status, result = self:receive_tag()
      if ( not(status) ) then return false, "ERROR: Failed to retrieve challenge" end
      status, chall = TagProcessor[result.name](self.socket, result)
      if ( mech == "LOGIN" ) then
        if ( chall ~= "User Name" ) then
          return false, ("ERROR: Login expected 'User Name' received: %s"):format(chall)
        end
        self.socket:send("<response xmlns='urn:ietf:params:xml:ns:xmpp-sasl'>" ..
        base64.enc(username) ..
        "</response>")
        status, result = self:receive_tag()
        if ( not(status) or result.name ~= "challenge") then
          return false, "ERROR: Receiving tag from server"
        end
        status, chall = TagProcessor[result.name](self.socket, result)
        if ( chall ~= "Password" ) then
          return false, ("ERROR: Login expected 'Password' received: %s"):format(chall)
        end
        self.socket:send("<response xmlns='urn:ietf:params:xml:ns:xmpp-sasl'>" ..
        base64.enc(password) ..
        "</response>")
        status, result = self:receive_tag()
        if ( not(status) ) then return false, "ERROR: Failed to receive login challenge" end
        if ( result.name == "failure" ) then
          status = TagProcessor[result.name](self.socket, result)
          return false, "Login failed"
        end
      else
        local mech_params = { username, password, chall, "xmpp", "xmpp/" .. self.servername  }
        local auth_data = sasl.Helper:new(mech):encode(table.unpack(mech_params))
        auth_data = "<response xmlns='urn:ietf:params:xml:ns:xmpp-sasl'>" ..
        base64.enc(auth_data) .. "</response>"
        status, err = self.socket:send(auth_data)
        repeat
          status, result = self:receive_tag()
          if ( not(status) ) then return false, "ERROR: Failed to receive login challenge" end
          if ( result.name == "failure" ) then
            status = TagProcessor[result.name](self.socket, result)
            return false, "Login failed"
          elseif ( result.name == "success" ) then
            status = TagProcessor[result.name](self.socket, result)
            if ( not(status) ) then return false, "Failed to process success message" end
            return true, "Login success"
          elseif ( result.name ~= "challenge" ) then
            return false, "ERROR: Failed to receive login challenge"
          end
        until( result.name == "challenge" and result.finish )
        if ( result.name == "challenge" and mech == "DIGEST-MD5" ) then
          status, result = self.socket:send("<response xmlns='urn:ietf:params:xml:ns:xmpp-sasl'/>")
          if ( not(status) ) then return false, "ERROR: Failed to send DIGEST-MD5 request" end
          status, result = self:receive_tag()
          if ( not(status) ) then return false, "ERROR: Failed to receive DIGEST-MD5 response" end
        end
      end
    end
    if ( result.name == "success" ) then
      return true, "Login success"
    end
    return false, "Login failed"
  end,
  getAuthMechs = function(self) return self.auth.mechs end,
  disconnect = function(self)
    local status, err = self.socket:close()
    self.socket = nil
    return status, err
  end,
}
Helper = {
  new = function(self, host, port, options)
    local o = { host = host,
    port = port,
    options = options or {},
    xmpp = XMPP:new(host, port, options),
    state = "" }
    setmetatable(o, self)
    self.__index = self
    return o
  end,
  connect = function(self, socket)
    if ( not(self.host.targetname) and
      not(self.options.servername) ) then
      return false, "ERROR: Cannot connect to XMPP server without valid server name"
    end
    self.state = "CONNECTED"
    return self.xmpp:connect(socket)
  end,
  login = function(self, username, password, mech)
    return self.xmpp:login(username, password, mech)
  end,
  getAuthMechs = function(self)
    if ( self.state == "CONNECTED" ) then
      return self.xmpp:getAuthMechs()
    end
    return
  end,
  close = function(self)
    self.xmpp:disconnect()
    self.state = "DISCONNECTED"
  end,
}
return _ENV;
