Updating TFTP client. Now with grown-up thread handling.

No longer blocks on successful connections.

lib/rex/proto/tftp/client.rb

@@ -14,8 +14,9 @@ module TFTP
 class Client
 
 	attr_accessor :local_host, :local_port, :peer_host, :peer_port
-	attr_accessor :thread, :context, :sock, :write_sock
+	attr_accessor :threads, :context, :server_sock, :client_sock
 	attr_accessor :local_file, :remote_file, :mode
+	attr_accessor :complete
 
 	# Returns an array of [code, type, msg]
 	def parse_tftp_msg(str)
@@ -25,6 +26,7 @@ class Client
 	end
 
 	def initialize(params)
+		self.threads = []
 		self.local_host = params["LocalHost"] || "0.0.0.0"
 		self.local_port = params["LocalPort"] || (1025 + rand(0xffff-1025))
 		self.peer_host = params["PeerHost"] || (raise ArgumentError, "Need a peer host.")
@@ -33,7 +35,6 @@ class Client
 		self.local_file = params["LocalFile"] || (raise ArgumentError, "Need a file to send.")
 		self.remote_file = params["RemoteFile"] || ::File.split(self.local_file).last
 		self.mode = params["Mode"] || "octet"
-		self.sock = nil
 	end
 
 	def blockify_file
@@ -53,10 +54,10 @@ class Client
 		end
 		data_blocks.each_with_index do |data_block,idx|
 			req = [3, (idx + 1), data_block].pack("nnA*")
-			if self.sock.sendto(req, host, port) > 0
+			if self.server_sock.sendto(req, host, port) > 0
 				sent_data += data_block.size
 			end
-			res = self.sock.recvfrom(65535)
+			res = self.server_sock.recvfrom(65535)
 			if res
 				code, type, msg = parse_tftp_msg(res[0])
 				if code == 4
@@ -72,35 +73,36 @@ class Client
 		end
 		if block_given?
 			if(sent_data == expected_size)
-				yield "Transfer complete!"
+				yield "Upload complete!"
 			else
-				yield "Transfer complete, but with errors."
+				yield "Upload complete, but with errors."
 			end
 		end
 	end
 
-	def monitor_socket
-			yield "Listening for incoming ACKs" if block_given?
-			res = self.sock.recvfrom(65535)
-			if res[0] and res[0] =~ /^\x00\x04/
-				send_data(res[1], res[2]) {|msg| yield msg}
-			end
-		stop
-	end
-
-	def start_client_port
-		self.sock = Rex::Socket::Udp.create(
+	def start_server_socket
+		self.server_sock = Rex::Socket::Udp.create(
 			'LocalHost' => local_host,
 			'LocalPort' => local_port,
 			'Context'   => context
 		)
-		if self.sock and block_given?
+		if self.server_sock and block_given?
 			yield "Started TFTP client listener on #{local_host}:#{local_port}"
 		end
 
-		self.thread = Rex::ThreadFactory.spawn("TFTPClientMonitor", false) {
-			monitor_socket {|msg| yield msg}
+		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
@@ -113,38 +115,38 @@ class Client
 
 	def send_write_request(&block)
 		if block_given?
-			start_client_port {|msg| yield msg}
+			start_server_socket {|msg| yield msg}
 		else
-			start_client_port 
+			start_server_socket 
 		end
-		self.write_sock = Rex::Socket::Udp.create(
+		self.client_sock = Rex::Socket::Udp.create(
 			'PeerHost'  => peer_host,
 			'PeerPort'  => peer_port,
 			'LocalHost' => local_host,
 			'LocalPort' => local_port,
 			'Context'   => context
 		)
-		self.write_sock.sendto(wrq_packet, peer_host, peer_port)
-		res = self.write_sock.recvfrom(65535)
-		if res 
+		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])
-			case code
-			when 1
-				yield "WRQ accepted, sending file" if block_given?
-				self.write_sock.close rescue nil
-			when nil
-				stop
-			else
-				yield("Aborting, got code:%d, type:%d, message:'%s'" % [code, type, msg]) if block_given?
-				stop
-			end
+			yield("Aborting, got code:%d, type:%d, message:'%s'" % [code, type, msg]) if block_given?
+			stop
 		end
 	end
 
 	def stop
-		self.thread.kill
-		self.sock.close rescue nil # might be closed already
-		self.write_sock.close rescue nil # might be closed already
+		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

modules/auxiliary/admin/tftp/tftp_upload_file.rb

@@ -67,7 +67,6 @@ class Metasploit3 < Msf::Auxiliary
 
 	#
 	# TFTP is a funny service and needs to kind of be a server on our side, too.
-	# Setup is called only once
 	def setup
 		@rport = datastore['RPORT'] || 69
 		@lport = datastore['LPORT'] || (1025 + rand(0xffff-1025))
@@ -89,21 +88,27 @@ class Metasploit3 < Msf::Auxiliary
 	def run
 		print_status "Sending '#{file}' to #{@lhost}:#{@lport} as '#{remote_file}'"
 		@tftp_client.send_write_request do |msg|
-			case msg
-			when /Aborting/, /errors.$/
-				print_error [rtarget,msg].join
-				@tftp_client.thread.kill
-			when /^WRQ accepted/, /^Sending/, /complete!$/
-				print_good [rtarget,msg].join
-			when /complete!$/
+			print_tftp_status(msg)
+		end
+		while true
+			if @tftp_client.complete
+				print_status [rtarget,"TFTP transfer operation complete."].join
+				break
 			else
-				vprint_status [rtarget,msg].join 
+				select(nil,nil,nil,1)
 			end
 		end
 	end
 
-	def cleanup
-		# Need to kill the server in case of interruption
+	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