#!/usr/bin/ruby

# a full ruby TFTP client or server
# (c) Yoann Guillot 2008
# distributes under the terms of the WTFPLv2

require 'socket'

# one tftp packet
class Tftp
	Opcode = {
		:RRQ => 1,
		:WRQ => 2,
		:DATA => 3,
		:ACK => 4,
		:ERROR => 5,
		:OACK => 6,
	}
	attr_accessor :opcode, :file, :nr, :data
	def initialize(mode, a, b=nil)
		case @opcode = mode
		when :RRQ, :WRQ; @file, @data = a, b
		when :DATA; @nr, @data = a, b
		when :ACK; @nr = a
		when :ERROR; @nr, @data = a, b
		when :OACK; @data = a
		else raise
		end
	end

	def encode
		ret = [Opcode[@opcode]].pack('n')
		case @opcode
		when :RRQ, :WRQ; mode = @data.delete(:mode) || 'octet' ; ret << @file << 0 << mode << 0 << @data.to_a.map { |e| e.to_s + 0.chr }.join
		when :DATA; ret << [@nr].pack('n') << @data
		when :ACK; ret << [@nr].pack('n')
		when :ERROR; ret << [@nr].pack('n') << @data << 0
		when :OACK; ret << @data.to_a.flatten.join(0.chr) << 0
		end
	end

	def self.decode(str)
		case op = Opcode.index(str[0, 2].unpack('n')[0])
		when :RRQ, :WRQ; opts = str[2..-1].split(0.chr) ; new(op, opts.shift, Hash[:mode, *opts])
		when :DATA; new(op, str[2, 2].unpack('n')[0], str[4..-1])
		when :ACK; new(op, str[2, 2].unpack('n')[0])
		when :ERROR; new(op, str[2, 2].unpack('n')[0], str[4...-1])
		when :OACK; new(op, Hash[str[2..-1].split(0.chr)])
		else puts "unknown packet #{str.inspect}"
		end
	end
end

class Tftpd
	def initialize(addr=0, port=69)
		@lfd = UDPSocket.open
		@lfd.bind(addr, port)
	end

	def repl(*a)
		@ofd.send(Tftp.new(*a).encode, 0, @c_ip, @c_p)
	end

	def mainloop
		loop do
			puts 'wait cx' if $VERBOSE
			nil while not IO.select([@lfd], nil, nil, 1)
			handle_newclient
		end
	end

	def handle_newclient
		packet, from = @lfd.recvfrom(516)
		p = Tftp.decode(packet)
		@c_ip, @c_p = from[3], from[1]
		puts "#{Time.now.strftime('%H:%M:%S')} rq from #@c_ip:#@c_p - #{p.inspect}" if $VERBOSE
		#@ofd = UDPSocket.open
		@ofd = @lfd
		p.file.tr!('/', '_') if not $insecure
		if not p or p.opcode != :RRQ or not File.exist? p.file or p.data[:mode] != 'octet'
			puts "not sending #{f}"
			repl :ERROR, 0, 'fail'
		elsif neg_opts(p.data, p.file)
		else
			puts "sending #{p.file}"
			File.open(p.file, 'rb') { |fd| send_file(fd) }
		end
		#@ofd.close
		@c_ip = @c_p = nil
	end

	def neg_opts(opts, file)
		@blksize = 512
		if opts.length > 1	# options ?
			resp = {}
			opts.each { |k, v|
				case k
				when 'tsize'; resp[k] = File.size(file)
				when 'blksize'; @blksize = v.to_i ; resp[k] = @blksize
				when :mode
				else puts "unsupported option #{k.inspect} => #{v.inspect}"
				end
			}
			repl :OACK, resp

			wait_ack(0)	# pass return value
		end
	end

	def send_file(fd)
		fd.seek(0, File::SEEK_END)
		sz = fd.tell
		fd.pos = 0
		i = 0
		loop do
			i += 1
			i = 0 if i > 0xffff
			buf = fd.read(@blksize).to_s
			repl(:DATA, i, buf)
			$stderr.print "\r#{fd.pos/@blksize}/#{sz/@blksize}     " if $stderr.tty? and $VERBOSE and fd.pos/@blksize % 137 == 9
			puts if buf.length != @blksize and $VERBOSE
			break if wait_ack(i) or buf.length != @blksize
		end
	end

	def wait_ack(nr)
		if not IO.select([@lfd], nil, nil, 4)
			puts "close ack timeout" if $VERBOSE
			repl :ERROR, 0, 'ack timeout'
			return true
		end
		ack = Tftp.decode(@ofd.recv(@blksize+4))
		if not ack or ack.opcode != :ACK or ack.nr != nr
			puts "close bad ack #{ack.inspect} (want #{nr})" if $VERBOSE
			repl :ERROR, 0, 'bad ack'
			return true
		end
	end
end

class Tftpc
	def initialize(addr=0, port=69)
		@fd = UDPSocket.open
		@c_ip, @c_p = addr, port
	end

	def repl(*a)
		@fd.send(Tftp.new(*a).encode, 0, @c_ip, @c_p)
	end

	def get(file)
		repl(:RRQ, file, :mode => 'octet')
		@blksize = 512

		File.open(file, 'wb') { |fd|
		loop do
			if not IO.select([@fd], nil, nil, 4)
				repl :ERROR, 0, 'data timeout'
				raise 'timeout'
			end
			p = @fd.recv(@blksize+4)
			p = Tftp.decode(p)
			raise p.inspect if not p or p.opcode != :DATA
			repl(:ACK, p.nr)
			fd.write p.data
			$stderr.print "\r#{fd.pos}    " if $stderr.tty? and $VERBOSE
			break if p.data.length != @blksize
		end
		}
		puts "\r#{File.size(file)} bytes ok" if $VERBOSE
	end
end

if __FILE__ == $0
	$VERBOSE = true if ARGV.delete '-v'
	$insecure = true if ARGV.delete '--insecure'	# don't strip / in paths
	case $0
	when /tftpd/i
		abort "usage: tftpd [[<rootdir>] <ip> [<port>]]" if ARGV.delete '--help' or ARGV.delete '-h'
		Dir.chdir ARGV.shift if ARGV[0] and File.directory? ARGV[0]
		Tftpd.new(*ARGV).mainloop
	else
		abort "usage: tftp <file> <ip> [<port>]" if ARGV.empty? or ARGV.delete '--help' or ARGV.delete '-h'
		Tftpc.new(*ARGV[1, 2]).get(ARGV[0])
	end
end