# $Id$
require 'rex/socket'
require 'rex/proto/tftp'

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 >= ERRCODES.length)
			# ignore..
			return
		end
		pkt = [OpError, num].pack('nn')
		pkt << 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, true)

		# 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] == 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 = [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 = [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 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, ErrFileNotFound)
				else
					transfer = {
						:type => 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, ErrFileNotFound)
			end

		when 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 => 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, ErrIllegalOperation)
			end

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

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

			tr = find_transfer(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, 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 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(OpWrite, from, (block-1))
			if not tr
				# If we didn't find it, send an error.
				send_error(from, 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, 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] == 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 = [OpOptAck].pack('n')
		to_ack.each { |el|
			data << el[0] << "\x00" << el[1] << "\x00"
		}

		send_packet(from, data)
	end

end

end
end
end