--- Library implementing a minimal TFTP server
--
-- Currently only write-operations are supported so that script can trigger
-- TFTP transfers and receive the files and return them as result.
--
-- The library contains the following classes
-- * <code>Packet</code>
-- ** The <code>Packet</code> classes contain one class for each TFTP operation.
-- * <code>File</code>
-- ** The <code>File</code> class holds a received file including the name and contents
-- * <code>ConnHandler</code>
-- ** The <code>ConnHandler</code> class handles and processes incoming connections.
--
-- The following code snippet starts the TFTP server and waits for the file incoming.txt
-- to be uploaded for 10 seconds:
-- <code>
--   tftp.start()
--   local status, f = tftp.waitFile("incoming.txt", 10)
--   if ( status ) then return f:getContent() end
-- </code>
--
-- @author Patrik Karlsson <patrik@cqure.net>
-- @copyright Same as Nmap--See https://nmap.org/book/man-legal.html
--

-- version 0.2
--
-- 2011-01-22 - re-wrote library to use coroutines instead of new_thread code.

local coroutine = require "coroutine"
local nmap = require "nmap"
local os = require "os"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"
_ENV = stdnse.module("tftp", stdnse.seeall)

threads, infiles, running = {}, {}, {}
state = "STOPPED"
srvthread = {}

-- All opcodes supported by TFTP
OpCode = {
  RRQ = 1,
  WRQ = 2,
  DATA = 3,
  ACK = 4,
  ERROR = 5,
}


--- A minimal packet implementation
--
-- The current code only implements the ACK and ERROR packets
-- As the server is write-only the other packet types are not needed
Packet = {

  -- Implements the ACK packet
  ACK = {

    new = function( self, block )
      local o = {}
      setmetatable(o, self)
      self.__index = self
      o.block = block
      return o
    end,

    __tostring = function( self )
      return string.pack(">I2I2", OpCode.ACK, self.block)
    end,

  },

  -- Implements the error packet
  ERROR = {

    new = function( self, code, msg )
      local o = {}
      setmetatable(o, self)
      self.__index = self
      o.msg = msg
      o.code = code
      return o
    end,

    __tostring = function( self )
      return string.pack(">I2I2z", OpCode.ERROR, self.code, self.msg)
    end,
  }

}

--- The File class holds files received by the TFTP server
File = {

  --- Creates a new file object
  --
  -- @param filename string containing the filename
  -- @param content string containing the file content
  -- @return o new class instance
  new = function(self, filename, content, sender)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    o.name = filename
    o.content = content
    o.sender = sender
    return o
  end,

  getContent = function(self) return self.content end,
  setContent = function(self, content) self.content = content end,

  getName = function(self) return self.name end,
  setName = function(self, name) self.name = name end,

  setSender = function(self, sender) self.sender = sender end,
  getSender = function(self) return self.sender end,
}


-- The thread dispatcher is called by the start function once
local function dispatcher()

  local last = os.time()
  local f_condvar = nmap.condvar(infiles)
  local s_condvar = nmap.condvar(state)

  while(true) do

    -- check if other scripts are active
    local counter = 0
    for t in pairs(running) do
      counter = counter + 1
    end
    if ( counter == 0 ) then
      state = "STOPPING"
      s_condvar "broadcast"
    end

    if #threads == 0 then break end
    for i, thread in ipairs(threads) do
      local status, res = coroutine.resume(thread)
      if ( not(res) ) then    -- thread finished its task?
        table.remove(threads, i)
        break
      end
    end

    -- Make sure to process waitFile atleast every 2 seconds
    -- in case no files have arrived
    if ( os.time() - last >= 2 ) then
      last = os.time()
      f_condvar "broadcast"
    end

  end
  state = "STOPPED"
  s_condvar "broadcast"
  stdnse.debug1("Exiting _dispatcher")
end

-- Processes a new incoming file transfer
-- Currently only uploads are supported
--
-- @param host containing the hostname or ip of the initiating host
-- @param port containing the port of the initiating host
-- @param data string containing the initial data passed to the server
local function processConnection( host, port, data )
  local op, pos = string.unpack(">I2", data)
  local socket = nmap.new_socket("udp")

  socket:set_timeout(1000)
  local status, err = socket:connect(host, port)
  if ( not(status) ) then return status, err end

  socket:set_timeout(10)

  -- If we get anything else than a write request, abort the connection
  if ( OpCode.WRQ ~= op ) then
    stdnse.debug1("Unsupported opcode")
    socket:send( tostring(Packet.ERROR:new(0, "TFTP server has write-only support")))
  end

  local filename, enctype, pos = string.unpack("zz", data, pos)
  status, err = socket:send( tostring( Packet.ACK:new(0) ) )

  local blocks = {}
  local lastread = os.time()

  while( true ) do
    local status, pdata = socket:receive()
    if ( not(status) ) then
      -- if we're here and haven't successfully read a packet for 5 seconds, abort
      if ( os.time() - lastread  > 5 ) then
        coroutine.yield(false)
      else
        coroutine.yield(true)
      end
    else
      -- record last time we had a successful read
      lastread = os.time()
      op, pos = string.unpack(">I2", pdata)
      if ( OpCode.DATA ~= op ) then
        stdnse.debug1("Expected a data packet, terminating TFTP transfer")
      end

      local block, data
      block, data, pos = string.unpack(">I2 c" .. #pdata - 4, pdata, pos )

      blocks[block] = data

      -- First block was not 1
      if ( #blocks == 0 ) then
        socket:send( tostring(Packet.ERROR:new(0, "Did not receive block 1")))
        break
      end

      -- for every fifth block check that we've received the preceding four
      if ( ( #blocks % 5 ) == 0 ) then
        for b = #blocks - 4, #blocks do
          if ( not(blocks[b]) ) then
            socket:send( tostring(Packet.ERROR:new(0, "Did not receive block " .. b)))
          end
        end
      end

      -- Ack the data block
      status, err = socket:send( tostring(Packet.ACK:new(block)) )

      if ( ( #blocks % 20 ) == 0 ) then
        -- yield every 5th iteration so other threads may work
        coroutine.yield(true)
      end

      -- If the data length was less than 512, this was our last block
      if ( #data < 512 ) then
        socket:close()
        break
      end
    end
  end

  local filecontent = {}

  -- Make sure we received all the blocks needed to proceed
  for i=1, #blocks do
    if ( not(blocks[i]) ) then
      return false, ("Block #%d was missing in transfer")
    end
    filecontent[#filecontent+1] = blocks[i]
  end
  stdnse.debug1("Finished receiving file \"%s\"", filename)

  -- Add  anew file to the global infiles table
  table.insert( infiles, File:new(filename, table.concat(filecontent), host) )

  local condvar = nmap.condvar(infiles)
  condvar "broadcast"
end

-- Waits for a connection from a client
local function waitForConnection()

  local srvsock = nmap.new_socket("udp")
  local status = srvsock:bind(nil, 69)
  assert(status, "Failed to bind to TFTP server port")

  srvsock:set_timeout(0)

  while( state == "RUNNING" ) do
    local status, data = srvsock:receive()
    if ( not(status) ) then
      coroutine.yield(true)
    else
      local status, _, _, rhost, rport = srvsock:get_info()
      local x = coroutine.create( function() processConnection(rhost, rport, data) end )
      table.insert( threads, x )
      coroutine.yield(true)
    end
  end
end


--- Starts the TFTP server and creates a new thread handing over to the dispatcher
function start()
  local disp = nil
  local mutex = nmap.mutex("srvsocket")

  -- register a running script
  running[coroutine.running()] = true

  mutex "lock"
  if ( state == "STOPPED" ) then
    srvthread = coroutine.running()
    table.insert( threads, coroutine.create( waitForConnection ) )
    stdnse.new_thread( dispatcher )
    state = "RUNNING"
  end
  mutex "done"

end

local function waitLast()
  -- The thread that started the server needs to wait here until the rest
  -- of the scripts finish running. We know we are done once the state
  -- shifts to STOPPED and we get a signal from the condvar in the
  -- dispatcher
  local s_condvar = nmap.condvar(state)
  while( srvthread == coroutine.running() and state ~= "STOPPED" ) do
    s_condvar "wait"
  end
end

--- Waits for a file with a specific filename for at least the number of
-- seconds specified by the timeout parameter.
--
-- If this function is called from the thread that's running the server it will
-- wait until all the other threads have finished executing before returning.
--
-- @param filename string containing the name of the file to receive
-- @param timeout number containing the minimum number of seconds to wait
--        for the file to be received
-- @return status true on success false on failure
-- @return File instance on success, nil on failure
function waitFile( filename, timeout )
  local condvar = nmap.condvar(infiles)
  local t = os.time()
  while(os.time() - t < timeout) do
    for _, f in ipairs(infiles) do
      if (f:getName() == filename) then
        running[coroutine.running()] = nil
        waitLast()
        return true, f
      end
    end
    condvar "wait"
  end
  -- de-register a running script
  running[coroutine.running()] = nil
  waitLast()

  return false
end

return _ENV;