lib/rex/proto/tftp/server.rb
# -*- coding: binary -*-
require 'rex/socket'
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 >= Constants::ERRCODES.length)
# ignore..
return
end
pkt = [Constants::OpError, num].pack('nn')
pkt << Constants::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)
# 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] == Constants::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 = [Constants::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 = [Constants::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 Constants::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, Constants::ErrFileNotFound)
else
transfer = {
:type => Constants::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, Constants::ErrFileNotFound)
end
when Constants::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 => Constants::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, Constants::ErrIllegalOperation)
end
when Constants::OpAck
# Process ACK packets
block = buf.unpack('n')[0]
#puts "%s %d" % [start, block]
tr = find_transfer(Constants::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, Constants::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 Constants::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(Constants::OpWrite, from, (block-1))
if not tr
# If we didn't find it, send an error.
send_error(from, Constants::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, Constants::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] == Constants::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 = [Constants::OpOptAck].pack('n')
to_ack.each { |el|
data << el[0] << "\x00" << el[1] << "\x00"
}
send_packet(from, data)
end
end
end
end
end