# -*- coding: binary -*-
require 'rex/socket'

module Rex
module Proto
module TFTP

#
# Little util function
#
def self.get_string(data)
  idx = data.index("\x00")
  return nil if not idx
  ret = data.slice!(0, idx)
  # Slice off the nul byte.
  data.slice!(0,1)
  ret
end


##
#
# TFTP Server class
#
##
class Server

  def initialize(port = 69, listen_host = '0.0.0.0', context = {})
    self.listen_host = listen_host
    self.listen_port = port
    self.context = context
    self.sock = nil
    @shutting_down = false
    @output_dir = nil
    @tftproot = nil

    self.files = []
    self.uploaded = []
    self.transfers = []
  end


  #
  # Start the TFTP server
  #
  def start
    self.sock = Rex::Socket::Udp.create(
      'LocalHost' => listen_host,
      'LocalPort' => listen_port,
      'Context'   => context
      )

    self.thread = Rex::ThreadFactory.spawn("TFTPServerMonitor", false) {
      monitor_socket
    }
  end


  #
  # Stop the TFTP server
  #
  def stop
    @shutting_down = true

    # Wait a maximum of 30 seconds for all transfers to finish.
    start = ::Time.now
    while (self.transfers.length > 0)
      ::IO.select(nil, nil, nil, 0.5)
      dur = ::Time.now - start
      break if (dur > 30)
    end

    self.files.clear
    self.thread.kill
    self.sock.close rescue nil # might be closed already
  end


  #
  # Register a filename and content for a client to request
  #
  def register_file(fn, content, once = false)
    self.files << {
      :name => fn,
      :data => content,
      :once => once
    }
  end


  #
  # Register an entire directory to serve files from
  #
  def set_tftproot(rootdir)
    @tftproot = rootdir if ::File.directory?(rootdir)
  end


  #
  # Register a directory to write uploaded files to
  #
  def set_output_dir(outdir)
    @output_dir = outdir if ::File.directory?(outdir)
  end


  #
  # Send an error packet w/the specified code and string
  #
  def send_error(from, num)
    if (num < 1 or num >= Constants::ERRCODES.length)
      # ignore..
      return
    end
    pkt = [Constants::OpError, num].pack('nn')
    pkt << Constants::ERRCODES[num]
    pkt << "\x00"
    send_packet(from, pkt)
  end


  #
  # Send a single packet to the specified host
  #
  def send_packet(from, pkt)
    self.sock.sendto(pkt, from[0], from[1])
  end


  #
  # Find the hash entry for a file that may be offered
  #
  def find_file(fname)
    # Files served via register_file() take precedence.
    self.files.each do |f|
      if (fname == f[:name])
        return f
      end
    end

    # Now, if we have a tftproot, see if it can serve from it
    if @tftproot
      return find_file_in_root(fname)
    end

    nil
  end


  #
  # Find the file in the specified tftp root and add a temporary
  # entry to the files hash.
  #
  def find_file_in_root(fname)
    fn = ::File.expand_path(::File.join(@tftproot, fname))

    # Don't allow directory traversal
    return nil if fn.index(@tftproot) != 0

    return nil if not ::File.file?(fn) or not ::File.readable?(fn)

    # Read the file contents, and register it as being served once
    data = data = ::File.open(fn, "rb") { |fd| fd.read(fd.stat.size) }
    register_file(fname, data)

    # Return the last file in the array
    return self.files[-1]
  end


  attr_accessor :listen_host, :listen_port, :context
  attr_accessor :sock, :files, :transfers, :uploaded
  attr_accessor :thread

  attr_accessor :incoming_file_hook

protected

  def find_transfer(type, from, block)
    self.transfers.each do |tr|
      if (tr[:type] == type and tr[:from] == from and tr[:block] == block)
        return tr
      end
    end
    nil
  end

  def save_output(tr)
    self.uploaded << tr[:file]

    return incoming_file_hook.call(tr) if incoming_file_hook

    if @output_dir
      fn = tr[:file][:name].split(File::SEPARATOR)[-1]
      if fn
        fn = ::File.join(@output_dir, Rex::FileUtils.clean_path(fn))
        ::File.open(fn, "wb") { |fd|
          fd.write(tr[:file][:data])
        }
      end
    end
  end


  def check_retransmission(tr)
    elapsed = ::Time.now - tr[:last_sent]
    if (elapsed >= tr[:timeout])
      # max retries reached?
      if (tr[:retries] < 3)
        #if (tr[:type] == OpRead)
        #	puts "[-] ack timed out, resending block"
        #else
        #	puts "[-] block timed out, resending ack"
        #end
        tr[:last_sent] = nil
        tr[:retries] += 1
      else
        #puts "[-] maximum tries reached, terminating transfer"
        self.transfers.delete(tr)
      end
    end
  end


  #
  # See if there is anything to do.. If so, dispatch it.
  #
  def monitor_socket
    while true
      rds = [@sock]
      wds = []
      self.transfers.each do |tr|
        if (not tr[:last_sent])
          wds << @sock
          break
        end
      end
      eds = [@sock]

      r,w,e = ::IO.select(rds,wds,eds,1)

      if (r != nil and r[0] == self.sock)
        buf,host,port = self.sock.recvfrom(65535)
        # Lame compatabilitiy :-/
        from = [host, port]
        dispatch_request(from, buf)
      end

      #
      # Check to see if transfers need maintenance
      #
      self.transfers.each do |tr|
        # We handle RRQ and WRQ separately
        #
        if (tr[:type] == Constants::OpRead)
          # Are we awaiting an ack?
          if (tr[:last_sent])
            check_retransmission(tr)
          elsif (w != nil and w[0] == self.sock)
            # No ack waiting, send next block..
            chunk = tr[:file][:data].slice(tr[:offset], tr[:blksize])
            if (chunk and chunk.length >= 0)
              pkt = [Constants::OpData, tr[:block]].pack('nn')
              pkt << chunk

              send_packet(tr[:from], pkt)
              tr[:last_sent] = ::Time.now

              # If the file is a one-serve, mark it as started
              tr[:file][:started] = true if (tr[:file][:once])
            else
              # No more chunks.. transfer is most likely done.
              # However, we can only delete it once the last chunk has been
              # acked.
            end
          end
        else
          # Are we awaiting data?
          if (tr[:last_sent])
            check_retransmission(tr)
          elsif (w != nil and w[0] == self.sock)
            # Not waiting for data, send an ack..
            #puts "[*] sending ack for block %d" % [tr[:block]]
            pkt = [Constants::OpAck, tr[:block]].pack('nn')

            send_packet(tr[:from], pkt)
            tr[:last_sent] = ::Time.now

            # If we had a 0-511 byte chunk, we're done.
            if (tr[:last_size] and tr[:last_size] < tr[:blksize])
              #puts "[*] Transfer complete, saving output"
              save_output(tr)
              self.transfers.delete(tr)
            end
          end
        end
      end
    end
  end


  def next_block(tr)
    tr[:block] += 1
    tr[:last_sent] = nil
    tr[:retries] = 0
  end


  #
  # Dispatch a packet that we received
  #
  def dispatch_request(from, buf)

    op = buf.unpack('n')[0]
    buf.slice!(0,2)

    #XXX: todo - create call backs for status
    #start = "[*] TFTP - %s:%u - %s" % [from[0], from[1], OPCODES[op]]

    case op
    when Constants::OpRead
      # Process RRQ packets
      fn = TFTP::get_string(buf)
      mode = TFTP::get_string(buf).downcase

      #puts "%s %s %s" % [start, fn, mode]

      if (not @shutting_down) and (file = self.find_file(fn))
        if (file[:once] and file[:started])
          send_error(from, Constants::ErrFileNotFound)
        else
          transfer = {
            :type => Constants::OpRead,
            :from => from,
            :file => file,
            :block => 1,
            :blksize => 512,
            :offset => 0,
            :timeout => 3,
            :last_sent => nil,
            :retries => 0
          }

          process_options(from, buf, transfer)

          self.transfers << transfer
        end
      else
        #puts "[-] file not found!"
        send_error(from, Constants::ErrFileNotFound)
      end

    when Constants::OpWrite
      # Process WRQ packets
      fn = TFTP::get_string(buf)
      mode = TFTP::get_string(buf).downcase

      #puts "%s %s %s" % [start, fn, mode]

      if not @shutting_down
        transfer = {
          :type => Constants::OpWrite,
          :from => from,
          :file => { :name => fn, :data => '' },
          :block => 0, # WRQ starts at 0
          :blksize => 512,
          :timeout => 3,
          :last_sent => nil,
          :retries => 0
        }

        process_options(from, buf, transfer)

        self.transfers << transfer
      else
        send_error(from, Constants::ErrIllegalOperation)
      end

    when Constants::OpAck
      # Process ACK packets
      block = buf.unpack('n')[0]

      #puts "%s %d" % [start, block]

      tr = find_transfer(Constants::OpRead, from, block)
      if not tr
        # NOTE: some clients, such as pxelinux, send an ack for block 0.
        # To deal with this, we simply ignore it as we start with block 1.
        return if block == 0

        # If we didn't find it, send an error.
        send_error(from, Constants::ErrUnknownTransferId)
      else
        # acked! send the next block
        tr[:offset] += tr[:blksize]
        next_block(tr)

        # If the transfer is finished, delete it
        if (tr[:offset] > tr[:file][:data].length)
          #puts "[*] Transfer complete"
          self.transfers.delete(tr)

          # if the file is a one-serve, delete it from the files array
          if tr[:file][:once]
            #puts "[*] Removed one-serve file: #{tr[:file][:name]}"
            self.files.delete(tr[:file])
          end
        end
      end

    when Constants::OpData
      # Process Data packets
      block = buf.unpack('n')[0]
      data = buf.slice(2, buf.length)

      #puts "%s %d %d bytes" % [start, block, data.length]

      tr = find_transfer(Constants::OpWrite, from, (block-1))
      if not tr
        # If we didn't find it, send an error.
        send_error(from, Constants::ErrUnknownTransferId)
      else
        tr[:file][:data] << data
        tr[:last_size] = data.length
        next_block(tr)

        # Similar to RRQ transfers, we cannot detect that the
        # transfer finished here. We must do so after transmitting
        # the final ACK.
      end

    else
      # Other packets are unsupported
      #puts start
      send_error(from, Constants::ErrAccessViolation)

    end
  end

  def process_options(from, buf, tr)
    found = 0
    to_ack = []
    while buf.length >= 4
      opt = TFTP::get_string(buf)
      break if not opt
      val = TFTP::get_string(buf)
      break if not val

      found += 1

      # Is it one we support?
      opt.downcase!

      case opt
      when "blksize"
        val = val.to_i
        if val > 0
          tr[:blksize] = val
          to_ack << [ opt, val.to_s ]
        end

      when "timeout"
        val = val.to_i
        if val >= 1 and val <= 255
          tr[:timeout] = val
          to_ack << [ opt, val.to_s ]
        end

      when "tsize"
        if tr[:type] == Constants::OpRead
          len = tr[:file][:data].length
        else
          val = val.to_i
          len = val
        end
        to_ack << [ opt, len.to_s ]

      end
    end

    return if to_ack.length < 1

    # if we have anything to ack, do it
    data = [Constants::OpOptAck].pack('n')
    to_ack.each { |el|
      data << el[0] << "\x00" << el[1] << "\x00"
    }

    send_packet(from, data)
  end

end

end
end
end
