146 lines
6.5 KiB
Python
146 lines
6.5 KiB
Python
import socket, os, re, time, random
|
|
import select
|
|
from TftpShared import *
|
|
from TftpPacketTypes import *
|
|
from TftpPacketFactory import *
|
|
from TftpStates import *
|
|
|
|
class TftpServer(TftpSession):
|
|
"""This class implements a tftp server object."""
|
|
|
|
def __init__(self, tftproot='/tftpboot', dyn_file_func=None):
|
|
"""Class constructor. It takes two optional arguments. tftproot is
|
|
the path to the tftproot directory to serve files from and/or write
|
|
them to. dyn_file_func is a callable that must return a file-like
|
|
object to read from during downloads. This permits the serving of
|
|
dynamic content."""
|
|
self.listenip = None
|
|
self.listenport = None
|
|
self.sock = None
|
|
# FIXME: What about multiple roots?
|
|
self.root = os.path.abspath(tftproot)
|
|
self.dyn_file_func = dyn_file_func
|
|
# A dict of sessions, where each session is keyed by a string like
|
|
# ip:tid for the remote end.
|
|
self.sessions = {}
|
|
|
|
if os.path.exists(self.root):
|
|
log.debug("tftproot %s does exist" % self.root)
|
|
if not os.path.isdir(self.root):
|
|
raise TftpException, "The tftproot must be a directory."
|
|
else:
|
|
log.debug("tftproot %s is a directory" % self.root)
|
|
if os.access(self.root, os.R_OK):
|
|
log.debug("tftproot %s is readable" % self.root)
|
|
else:
|
|
raise TftpException, "The tftproot must be readable"
|
|
if os.access(self.root, os.W_OK):
|
|
log.debug("tftproot %s is writable" % self.root)
|
|
else:
|
|
log.warning("The tftproot %s is not writable" % self.root)
|
|
else:
|
|
raise TftpException, "The tftproot does not exist."
|
|
|
|
def listen(self,
|
|
listenip="",
|
|
listenport=DEF_TFTP_PORT,
|
|
timeout=SOCK_TIMEOUT):
|
|
"""Start a server listening on the supplied interface and port. This
|
|
defaults to INADDR_ANY (all interfaces) and UDP port 69. You can also
|
|
supply a different socket timeout value, if desired."""
|
|
tftp_factory = TftpPacketFactory()
|
|
|
|
# Don't use new 2.5 ternary operator yet
|
|
# listenip = listenip if listenip else '0.0.0.0'
|
|
if not listenip: listenip = '0.0.0.0'
|
|
log.info("Server requested on ip %s, port %s"
|
|
% (listenip, listenport))
|
|
try:
|
|
# FIXME - sockets should be non-blocking?
|
|
self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
self.sock.bind((listenip, listenport))
|
|
except socket.error, err:
|
|
# Reraise it for now.
|
|
raise
|
|
|
|
log.info("Starting receive loop...")
|
|
while True:
|
|
# Build the inputlist array of sockets to select() on.
|
|
inputlist = []
|
|
inputlist.append(self.sock)
|
|
for key in self.sessions:
|
|
inputlist.append(self.sessions[key].sock)
|
|
|
|
# Block until some socket has input on it.
|
|
log.debug("Performing select on this inputlist: %s" % inputlist)
|
|
readyinput, readyoutput, readyspecial = select.select(inputlist,
|
|
[],
|
|
[],
|
|
SOCK_TIMEOUT)
|
|
|
|
deletion_list = []
|
|
|
|
# Handle the available data, if any. Maybe we timed-out.
|
|
for readysock in readyinput:
|
|
# Is the traffic on the main server socket? ie. new session?
|
|
if readysock == self.sock:
|
|
log.debug("Data ready on our main socket")
|
|
buffer, (raddress, rport) = self.sock.recvfrom(MAX_BLKSIZE)
|
|
|
|
log.debug("Read %d bytes" % len(buffer))
|
|
|
|
recvpkt = tftp_factory.parse(buffer)
|
|
# Forge a session key based on the client's IP and port,
|
|
# which should safely work through NAT.
|
|
key = "%s:%s" % (raddress, rport)
|
|
|
|
if not self.sessions.has_key(key):
|
|
log.debug("Creating new server context for "
|
|
"session key = %s" % key)
|
|
self.sessions[key] = TftpContextServer(raddress,
|
|
rport,
|
|
timeout,
|
|
self.root,
|
|
self.dyn_file_func)
|
|
self.sessions[key].start(buffer)
|
|
else:
|
|
log.warn("received traffic on main socket for "
|
|
"existing session??")
|
|
|
|
else:
|
|
# Must find the owner of this traffic.
|
|
for key in self.sessions:
|
|
if readysock == self.sessions[key].sock:
|
|
log.info("Matched input to session key %s"
|
|
% key)
|
|
try:
|
|
self.sessions[key].cycle()
|
|
if self.sessions[key].state == None:
|
|
log.info("Successful transfer.")
|
|
deletion_list.append(key)
|
|
except TftpException, err:
|
|
deletion_list.append(key)
|
|
log.error("Fatal exception thrown from "
|
|
"session %s: %s"
|
|
% (key, str(err)))
|
|
break
|
|
|
|
else:
|
|
log.error("Can't find the owner for this packet. "
|
|
"Discarding.")
|
|
|
|
log.debug("Looping on all sessions to check for timeouts")
|
|
now = time.time()
|
|
for key in self.sessions:
|
|
try:
|
|
self.sessions[key].checkTimeout(now)
|
|
except TftpException, err:
|
|
log.error(str(err))
|
|
deletion_list.append(key)
|
|
|
|
log.debug("Iterating deletion list.")
|
|
for key in deletion_list:
|
|
if self.sessions.has_key(key):
|
|
log.debug("Deleting session %s" % key)
|
|
del self.sessions[key]
|