Adding download and loot functionality.

Still need to deal with the use case of not passing a block; blocks
should not be required, it should be okay to invoke and just wait for
the complete attribute to be true. You'll miss out on error messages but
eh, maybe those should be return values.

lib/rex/proto/tftp/client.rb

@@ -1,6 +1,7 @@
 # $Id$
 require 'rex/socket'
 require 'rex/proto/tftp'
+require 'tempfile'
 
 module Rex
 module Proto
@@ -15,13 +16,16 @@ class Client
 
 	attr_accessor :local_host, :local_port, :peer_host, :peer_port
 	attr_accessor :threads, :context, :server_sock, :client_sock
-	attr_accessor :local_file, :remote_file, :mode
+	attr_accessor :local_file, :remote_file, :mode, :action
-	attr_accessor :complete
+	attr_accessor :complete, :recv_tempfile
 
-	# Returns an array of [code, type, msg]
+	# Returns an array of [code, type, msg]. Data packets
-	def parse_tftp_msg(str)
+	# should set strip to false, or else trailing spaces and nulls
+	# will be dropped during unpacking.
+	def parse_tftp_msg(str,strip=true)
 		return nil unless str.length >= 4
 		ret = str.unpack("nnA*")
+		ret[2] = str[4,str.size] unless strip
 		return ret
 	end
 
@@ -32,9 +36,142 @@ class Client
 		self.peer_host = params["PeerHost"] || (raise ArgumentError, "Need a peer host.")
 		self.peer_port = params["PeerPort"] || 69
 		self.context = params["Context"] || {}
-		self.local_file = params["LocalFile"] || (raise ArgumentError, "Need a file to send.")
+		self.local_file = params["LocalFile"] || (raise ArgumentError, "Need a local file.")
 		self.remote_file = params["RemoteFile"] || ::File.split(self.local_file).last
 		self.mode = params["Mode"] || "octet"
+		self.action = params["Action"] || (raise ArgumentError, "Need an action.")
+	end
+
+	#
+	# Methods for both upload and download
+	#
+	
+	def start_server_socket
+		self.server_sock = Rex::Socket::Udp.create(
+			'LocalHost' => local_host,
+			'LocalPort' => local_port,
+			'Context'   => context
+		)
+		if self.server_sock and block_given?
+			yield "Started TFTP client listener on #{local_host}:#{local_port}"
+		end
+		self.threads << Rex::ThreadFactory.spawn("TFTPServerMonitor", false) {
+			monitor_server_sock {|msg| yield msg}
+		}
+	end
+
+	def monitor_server_sock
+		yield "Listening for incoming ACKs" if block_given?
+		res = self.server_sock.recvfrom(65535)
+		if res and res[0]
+			code, type, data = parse_tftp_msg(res[0])
+			if code == 4 && self.action == :upload
+				send_data(res[1], res[2]) {|msg| yield msg}
+			elsif code == 3 && self.action == :download
+				recv_data(res[1], res[2], data) {|msg| yield msg}
+			else
+				yield("Aborting, got code:%d, type:%d, message:'%s'" % [code, type, msg]) if block_given?
+				stop
+			end
+		end
+		stop
+	end
+
+	def monitor_client_sock
+		res = self.client_sock.recvfrom(65535)
+		if res[1] # Got a response back, so that's never good; Acks come back on server_sock.
+			code, type, msg = parse_tftp_msg(res[0])
+			yield("Aborting, got code:%d, type:%d, message:'%s'" % [code, type, msg]) if block_given?
+			stop
+		end
+	end
+
+	def stop
+		self.complete = true
+		self.threads.each {|t| t.kill}
+		self.server_sock.close rescue nil # might be closed already
+		self.client_sock.close rescue nil # might be closed already
+	end
+
+	#
+	# Methods for download
+	#
+	
+	def rrq_packet
+		req = [OpRead, self.remote_file, self.mode]
+		packstr = "na#{self.remote_file.length+1}a#{self.mode.length+1}"
+		req.pack(packstr)
+	end
+
+	def ack_packet(blocknum=0)
+		req = [OpAck, blocknum].pack("nn")
+	end
+
+	def send_read_request(&block)
+		if block_given?
+			start_server_socket {|msg| yield msg}
+		else
+			start_server_socket
+		end
+		self.client_sock = Rex::Socket::Udp.create(
+			'PeerHost'  => peer_host,
+			'PeerPort'  => peer_port,
+			'LocalHost' => local_host,
+			'LocalPort' => local_port,
+			'Context'   => context
+		)
+		self.client_sock.sendto(rrq_packet, peer_host, peer_port)
+		self.threads << Rex::ThreadFactory.spawn("TFTPClientMonitor", false) {
+			monitor_client_sock {|msg| yield msg}
+		}
+	end
+
+	def recv_data(host, port, first_block)
+		self.recv_tempfile = Rex::Quickfile.new('msf-tftp')
+		recvd_blocks = 1
+		if block_given?
+			yield "Source file: #{self.remote_file}, destination file: #{self.local_file}"
+			yield "Received and acknowledged #{first_block.size} in block #{recvd_blocks}"
+		end
+		write_and_ack_data(first_block,1,host,port) {|msg| yield msg}
+		current_block = first_block
+		while current_block.size == 512
+			res = self.server_sock.recvfrom(65535)
+			if res and res[0]
+				code, block_num, current_block = parse_tftp_msg(res[0])
+				if code == 3
+					write_and_ack_data(current_block,block_num,host,port) {|msg| yield msg}
+					recvd_blocks += 1
+				else
+					yield("Aborting, got code:%d, type:%d, message:'%s'" % [code, type, msg]) if block_given?
+					stop
+				end
+			end
+		end
+		yield("Transferred #{recvd_blocks} blocks, download complete!")
+		self.recv_tempfile.close
+		stop
+	end
+
+	def write_and_ack_data(data,blocknum,host,port)
+		self.recv_tempfile.write(data)
+		self.recv_tempfile.flush
+		req = ack_packet(blocknum)
+		self.server_sock.sendto(req, host, port)
+		yield "Received and acknowledged #{data.size} in block #{blocknum}"
+	end
+
+
+	#
+	# Methods for upload
+	#
+
+	def wrq_packet
+		req = "\x00\x02"
+		req += self.remote_file
+		req += "\x00"
+		req += self.mode
+		req += "\x00"
 	end
 
 	def blockify_file
@@ -42,13 +179,32 @@ class Client
 		data.scan(/.{1,512}/)
 	end
 
+	def send_write_request(&block)
+		if block_given?
+			start_server_socket {|msg| yield msg}
+		else
+			start_server_socket
+		end
+		self.client_sock = Rex::Socket::Udp.create(
+			'PeerHost'  => peer_host,
+			'PeerPort'  => peer_port,
+			'LocalHost' => local_host,
+			'LocalPort' => local_port,
+			'Context'   => context
+		)
+		self.client_sock.sendto(wrq_packet, peer_host, peer_port)
+		self.threads << Rex::ThreadFactory.spawn("TFTPClientMonitor", false) {
+			monitor_client_sock {|msg| yield msg}
+		}
+	end
+
 	def send_data(host,port)
 		data_blocks = blockify_file()
 		sent_data = 0
 		sent_blocks = 0
 		expected_blocks = data_blocks.size
 		expected_size = data_blocks.join.size
-		if block_given? 
+		if block_given?
 			yield "Source file: #{self.local_file}, destination file: #{self.remote_file}"
 			yield "Sending #{expected_size} bytes (#{expected_blocks} blocks)"
 		end
@@ -67,7 +223,7 @@ class Client
 					if block_given?
 						yield "Got an unexpected response: Code:%d, Type:%d, Message:'%s'. Aborting." % [code, type, msg]
 					end
-					break 
+					break
 				end
 			end
 		end
@@ -80,75 +236,6 @@ class Client
 		end
 	end
 
-	def start_server_socket
-		self.server_sock = Rex::Socket::Udp.create(
-			'LocalHost' => local_host,
-			'LocalPort' => local_port,
-			'Context'   => context
-		)
-		if self.server_sock and block_given?
-			yield "Started TFTP client listener on #{local_host}:#{local_port}"
-		end
-
-		self.threads << Rex::ThreadFactory.spawn("TFTPServerMonitor", false) {
-			monitor_server_sock {|msg| yield msg}
-		}
-
-	end
-
-	def monitor_server_sock
-			yield "Listening for incoming ACKs" if block_given?
-			res = self.server_sock.recvfrom(65535)
-			if res[0] and res[0] =~ /^\x00\x04/
-				send_data(res[1], res[2]) {|msg| yield msg}
-			end
-		stop
-	end
-
-	def wrq_packet
-		req = "\x00\x02"
-		req += self.remote_file
-		req += "\x00"
-		req += self.mode
-		req += "\x00"
-	end
-
-	def send_write_request(&block)
-		if block_given?
-			start_server_socket {|msg| yield msg}
-		else
-			start_server_socket 
-		end
-		self.client_sock = Rex::Socket::Udp.create(
-			'PeerHost'  => peer_host,
-			'PeerPort'  => peer_port,
-			'LocalHost' => local_host,
-			'LocalPort' => local_port,
-			'Context'   => context
-		)
-		self.client_sock.sendto(wrq_packet, peer_host, peer_port)
-
-		self.threads << Rex::ThreadFactory.spawn("TFTPClientMonitor", false) {
-			monitor_client_sock {|msg| yield msg}
-		}
-	end
-
-	def monitor_client_sock
-		res = self.client_sock.recvfrom(65535)
-		if res[1] # Got a response back, so that's never good; Acks come back on server_sock. 
-			code, type, msg = parse_tftp_msg(res[0])
-			yield("Aborting, got code:%d, type:%d, message:'%s'" % [code, type, msg]) if block_given?
-			stop
-		end
-	end
-
-	def stop
-		self.complete = true
-		self.threads.each {|t| t.kill}
-		self.server_sock.close rescue nil # might be closed already
-		self.client_sock.close rescue nil # might be closed already
-	end
-
 end
 
 end

modules/auxiliary/admin/tftp/tftp_upload_file.rb

@@ -11,6 +11,7 @@ require 'msf/core'
 class Metasploit3 < Msf::Auxiliary
 
 	include Rex::Proto::TFTP
+	include Msf::Auxiliary::Report
 
 	def initialize
 		super(
@@ -21,22 +22,36 @@ class Metasploit3 < Msf::Auxiliary
 					for TFTP is often unsupported.
 				},
 			'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."}],
+				[ 'Upload', {'Description' => "Upload FILENAME as REMOTE_FILENAME to the server."}]
+				],
+			'DefaultAction' => 'Upload',
 			'License'     => MSF_LICENSE
 		)
 		register_options([
-			OptPath.new('FILENAME', [true, "The local file to upload" ]),
+			OptPath.new(   'FILENAME', [false, "The local filename" ]),
-			OptString.new('REMOTE_FILENAME', [false, "The filename to provide to the TFTP server" ]),
+			OptString.new( 'REMOTE_FILENAME', [false, "The remote filename"]),
-			OptAddress.new('RHOST', [true, "The remote TFTP server"]),
+			OptAddress.new('RHOST',    [true, "The remote TFTP server"]),
-			OptPort.new('LPORT',    [false, "The local port the TFTP client should listen on" ]),
+			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"]),
+			OptAddress.new('LHOST',    [false, "The local address the TFTP client should bind to"]),
-			OptBool.new('VERBOSE', [false, "Provide more details about the transfer", false]),
+			OptBool.new(   'VERBOSE',  [false, "Display verbose details about the transfer", false]),
-			OptString.new('MODE', [false, "The TFTP mode. Usual choices are netsacii and octet.", "octet"]),
+			OptString.new(  'MODE',    [false, "The TFTP mode; usual choices are netascii and octet.", "octet"]),
 			Opt::RPORT(69)
 		], self.class)
 	end
 
 	def file
-		datastore['FILENAME']
+		if action.name == "Upload"
+			datastore['FILENAME']
+		else # "Download
+			fname = ::File.split(datastore['FILENAME'] || datastore['REMOTE_FILENAME']).last
+		end
 	end
 
 	def mode
@@ -44,17 +59,26 @@ class Metasploit3 < Msf::Auxiliary
 	end
 
 	def remote_file
-		datastore['REMOTE_FILENAME'] || ::File.split(file).last
+		datastore['REMOTE_FILENAME'] || ::File.split(datastore['FILENAME']).last
 	end
 
 	def rport
-		datastore['RPORT']
+		datastore['RPORT'] || 69
 	end
 
 	def rhost
 		datastore['RHOST']
 	end
 
+	def datatype
+		case datastore['MODE']
+		when "netascii"
+			"text/plain"
+		else
+			"application/octet-stream"
+		end
+	end
+
 	def rtarget(ip=nil)
 		if (ip or rhost) and rport
 			[(ip || rhost),rport].map {|x| x.to_s}.join(":") << " "
@@ -65,34 +89,44 @@ class Metasploit3 < Msf::Auxiliary
 		end
 	end
 
+	def check_valid_filename
+		not (datastore['FILENAME'].to_s.empty? and datastore['REMOTE_FILENAME'].to_s.empty?)
+	end
+
 	#
 	# TFTP is a funny service and needs to kind of be a server on our side, too.
 	def setup
-		@rport = datastore['RPORT'] || 69
+		unless check_valid_filename()
+			print_error "Need at least one valid filename."
+			return
+		end
 		@lport = datastore['LPORT'] || (1025 + rand(0xffff-1025))
 		@lhost = datastore['LHOST'] || "0.0.0.0"
-		@path     = datastore['FILENAME']
+		@local_file = file
-		@filename = ::File.split(@path).last
+		@remote_file = remote_file
 
 		@tftp_client = Rex::Proto::TFTP::Client.new(
 			"LocalHost" => @lhost,
 			"LocalPort" => @lport,
 			"PeerHost"  => rhost,
 			"PeerPort"  => rport,
-			"LocalFile" => file,
+			"LocalFile" => @local_file,
-			"RemoteFile" => remote_file,
+			"RemoteFile" => @remote_file,
-			"Mode"      => mode
+			"Mode"      => mode,
+			"Action"    => action.name.to_s.downcase.intern
 		)
 	end
 
 	def run
-		print_status "Sending '#{file}' to #{@lhost}:#{@lport} as '#{remote_file}'"
+		return unless check_valid_filename()
-		@tftp_client.send_write_request do |msg|
+		run_upload()   if action.name == 'Upload'
-			print_tftp_status(msg)
+		run_download() if action.name == 'Download'
-		end
 		while true
 			if @tftp_client.complete
 				print_status [rtarget,"TFTP transfer operation complete."].join
+				if action.name == 'Download'
+					save_downloaded_file()
+				end
 				break
 			else
 				select(nil,nil,nil,1)
@@ -100,6 +134,38 @@ class Metasploit3 < Msf::Auxiliary
 		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 = File.open(fh,"rb") {|f| f.read f.stat.size} rescue nil
+		unless framework.db.active
+			print_status "No database connected, so not actually saving the data:"
+			print_line data
+		end
+		if data and not data.empty?
+			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.$/
@@ -107,7 +173,7 @@ class Metasploit3 < Msf::Auxiliary
 		when /^WRQ accepted/, /^Sending/, /complete!$/
 			print_good [rtarget,msg].join
 		else
-			vprint_status [rtarget,msg].join 
+			vprint_status [rtarget,msg].join
 		end
 	end