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

class MetasploitModule < 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,
      'Notes' => {
        'Stability' => [CRASH_SAFE],
        'SideEffects' => [IOC_IN_LOGS, ARTIFACTS_ON_DISK],
        'Reliability' => []
      }
    )
    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)' ]),
      OptAddressLocal.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)
    ])
  end

  def mode
    datastore['MODE'] || 'octet'
  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 !fdata.empty?
        "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) && ::File.readable?(fname_local)
        return fname_data if ::File.file?(fname_data) && ::File.readable?(fname_data)

        return nil # Couldn't find it, giving up.
      end
    else # "Download"
      begin
        ::File.split(datastore['FILENAME'] || datastore['REMOTE_FILENAME']).last
      rescue StandardError
        nil
      end
    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 || rhost) && rport
      [ip || rhost, rport].map(&:to_s).join(':') << ' '
    elsif ip || 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' => 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

    until @tftp_client.complete
      select(nil, nil, nil, 1)
      print_status [rtarget, 'TFTP transfer operation complete.'].join
      save_downloaded_file if action.name == 'Download'
    end
  end

  # Run in case something untoward happened 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 && @tftp_client.respond_to?(:complete)
      until @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}'"
    @tftp_client.send_write_request { |msg| print_tftp_status(msg) }
  end

  def run_download
    print_status "Receiving '#{remote_file}' from #{rhost}:#{rport} as '#{file}'"
    @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 = begin
      File.open(fh, 'rb') { |f| f.read f.stat.size }
    rescue StandardError
      nil
    end
    if data && !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
    begin
      fh.unlink
    rescue StandardError
      nil
    end
  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
