##
# This module requires Metasploit: http://metasploit.com/download
# Current source: https://github.com/rapid7/metasploit-framework
##


require 'msf/core'

class Metasploit3 < Msf::Auxiliary

  include Rex::Proto::TFTP
  include Msf::Auxiliary::Report

  def initialize
    super(
      'Name'        => 'TFTP File Transfer Utility',
      'Description' => %q{
          This module will transfer a file to or from a remote TFTP server.
          Note that the target must be able to connect back to the Metasploit system,
          and NAT traversal for TFTP is often unsupported.

          Two actions are supported: "Upload" and "Download," which behave as one might
          expect -- use 'set action Actionname' to use either mode of operation.

          If "Download" is selected, at least one of FILENAME or REMOTE_FILENAME
          must be set. If "Upload" is selected, either FILENAME must be set to a valid path to
          a source file, or FILEDATA must be populated. FILENAME may be a fully qualified path,
          or the name of a file in the Msf::Config.local_directory or Msf::Config.data_directory.
        },
      'Author'      => [ 'todb' ],
      'References'  =>
        [
          ['URL', 'http://www.faqs.org/rfcs/rfc1350.html'],
          ['URL', 'http://www.networksorcery.com/enp/protocol/tftp.htm']
        ],
      'Actions' => [
        [ 'Download', {'Description' => "Download REMOTE_FILENAME as FILENAME from the server."}],
        [ 'Upload',   {'Description' => "Upload FILENAME as REMOTE_FILENAME to the server."}]
        ],
      'DefaultAction' => 'Upload',
      'License'     => MSF_LICENSE
    )
    register_options([
      OptString.new( 'FILENAME', [false, "The local filename" ]),
      OptString.new( 'FILEDATA', [false, "Data to upload in lieu of a real local file." ]),
      OptString.new( 'REMOTE_FILENAME', [false, "The remote filename"]),
      OptAddress.new('RHOST',    [true, "The remote TFTP server"]),
      OptPort.new(   'LPORT',    [false, "The local port the TFTP client should listen on (default is random)" ]),
      OptAddress.new('LHOST',    [false, "The local address the TFTP client should bind to"]),
      OptString.new( 'MODE',     [false, "The TFTP mode; usual choices are netascii and octet.", "octet"]),
      Opt::RPORT(69)
    ], self.class)
  end

  def mode
    datastore['MODE'] || "octect"
  end

  def remote_file
    return datastore['REMOTE_FILENAME'] if datastore['REMOTE_FILENAME']
    return ::File.split(datastore['FILENAME']).last if datastore['FILENAME']
  end

  def rport
    datastore['RPORT'] || 69
  end

  def rhost
    datastore['RHOST']
  end

  # Used only to store loot, doesn't actually have any semantic meaning
  # for the TFTP protocol.
  def datatype
    case datastore['MODE']
    when "netascii"
      "text/plain"
    else
      "application/octet-stream"
    end
  end

  def file
    if action.name == "Upload"
      fdata = datastore['FILEDATA'].to_s
      fname = datastore['FILENAME'].to_s
      if not fdata.empty?
        fdata_decorated = "DATA:#{datastore['FILEDATA']}"
      elsif ::File.readable? fname
        fname
      else
        fname_local = ::File.join(Msf::Config.local_directory,fname)
        fname_data  = ::File.join(Msf::Config.data_directory,fname)
        return fname_local if ::File.file?(fname_local) and ::File.readable?(fname_local)
        return fname_data  if ::File.file?(fname_data)  and ::File.readable?(fname_data)
        return nil # Couldn't find it, giving up.
      end
    else # "Download"
      fname = ::File.split(datastore['FILENAME'] || datastore['REMOTE_FILENAME']).last rescue nil
    end
  end

  # Experimental message prepending thinger. Might make it up into the
  # standard Metasploit lib like vprint_status and friends.
  def rtarget(ip=nil)
    if (ip or rhost) and rport
      [(ip || rhost),rport].map {|x| x.to_s}.join(":") << " "
    elsif (ip or rhost)
      "#{rhost} "
    else
      ""
    end
  end

  # This all happens before run(), and should give an idea on how to use
  # the TFTP client mixin. Essentially, you create an instance of the
  # Rex::Proto::TFTP::Client class, fill it up with the relevant host and
  # file data, set it to either :upload or :download, then kick off the
  # transfer as you like.
  def setup
    @lport = datastore['LPORT'] || (1025 + rand(0xffff-1025))
    @lhost = datastore['LHOST'] || "0.0.0.0"
    @local_file = file
    @remote_file = remote_file

    @tftp_client = Rex::Proto::TFTP::Client.new(
      "LocalHost"  => @lhost,
      "LocalPort"  => @lport,
      "PeerHost"   => rhost,
      "PeerPort"   => rport,
      "LocalFile"  => @local_file,
      "RemoteFile" => @remote_file,
      "Mode"       => mode,
      "Context"    => {'Msf' => self.framework, 'MsfExploit' => self},
      "Action"     => action.name.to_s.downcase.intern
    )
  end

  def run
    case action.name
    when 'Upload'
      if file
        run_upload()
      else
        print_error "Need at least a local file name or file data to upload."
        return
      end
    when 'Download'
      if remote_file
        run_download()
      else
        print_error "Need at least a remote file name to download."
        return
      end
    else
      print_error "Unknown action: '#{action.name}'"
    end
    while not @tftp_client.complete
      select(nil,nil,nil,1)
      print_status [rtarget,"TFTP transfer operation complete."].join
      save_downloaded_file() if action.name == 'Download'
      break
    end
  end

  # Run in case something untoward happend with the connection and the
  # client object didn't get stopped on its own. This can happen with
  # transfers that got interrupted or malformed (like sending a 0 byte
  # file).
  def cleanup
    if @tftp_client and @tftp_client.respond_to? :complete
      while not @tftp_client.complete
        select(nil,nil,nil,1)
        vprint_status "Cleaning up the TFTP client ports and threads."
        @tftp_client.stop
      end
    end
  end

  def run_upload
    print_status "Sending '#{file}' to #{rhost}:#{rport} as '#{remote_file}'"
    ret = @tftp_client.send_write_request { |msg| print_tftp_status(msg) }
  end

  def run_download
    print_status "Receiving '#{remote_file}' from #{rhost}:#{rport} as '#{file}'"
    ret = @tftp_client.send_read_request { |msg| print_tftp_status(msg) }
  end

  def save_downloaded_file
    print_status "Saving #{remote_file} as '#{file}'"
    fh = @tftp_client.recv_tempfile
    data = File.open(fh,"rb") {|f| f.read f.stat.size} rescue nil
    if data and not data.empty?
      unless framework.db.active
        print_status "No database connected, so not actually saving the data:"
        print_line data
      end
      this_service = report_service(
        :host => rhost,
        :port => rport,
        :name => "tftp",
        :proto => "udp"
      )
      store_loot("tftp.file",datatype,rhost,data,file,remote_file,this_service)
    else
      print_status [rtarget,"Did not find any data, so nothing to save."].join
    end
    fh.unlink rescue nil # Windows often complains about unlinking tempfiles
  end

  def print_tftp_status(msg)
    case msg
    when /Aborting/, /errors.$/
      print_error [rtarget,msg].join
    when /^WRQ accepted/, /^Sending/, /complete!$/
      print_good [rtarget,msg].join
    else
      vprint_status [rtarget,msg].join
    end
  end

end