Add the possibility to use a dynamic file loader function

ptftplib/tftpserver.py

@@ -110,6 +110,12 @@ def handle(self):
             self.wfile.write(response)
             self.wfile.flush()
 
+    def finish(self):
+        """ Workaround for following bug: http://bugs.python.org/issue1767511 """
+        wdata = self.wfile.getvalue()
+        if wdata:
+            self.socket.sendto(wdata, self.client_address)
+
     def finish_state(self, peer_state):
         self.server.clients[self.client_address] = peer_state
         return peer_state.next()
@@ -127,11 +133,23 @@ def serveRRQ(self, op, request):
         """
 
         try:
-            filename, mode, opts = proto.TFTPHelper.parseRRQ(request)
+            filenameOrig, mode, opts = proto.TFTPHelper.parseRRQ(request)
         except SyntaxError:
             # Ignore malformed RRQ requests
             return None
 
+        # we keep the filenameOrig for dynamic handler
+        # some clients request "\" and not "/" in their pathes
+        if os.name != 'nt':
+            filename = filenameOrig.replace('\\', '/')
+        else:
+            filename = filenameOrig.replace('/', '\\')
+
+        # absolute path requests are always relative to the tftp root directory
+        # if the client does something nasty we get it some lines down.
+        if filename and filename[0] == "/":
+            filename = filename[1:]
+
         peer_state = state.TFTPState(self.client_address, op,
                 self.server.root, filename, mode,
                 not self.server.strict_rfc1350)
@@ -145,8 +163,22 @@ def serveRRQ(self, op, request):
             return self.finish_state(peer_state)
 
         try:
-            peer_state.file = open(peer_state.filepath, 'rb')
-            peer_state.filesize = os.stat(peer_state.filepath)[stat.ST_SIZE]
+            # If the file exists, open it
+            # TODO: Windows clients request case insensitive, which may be a problem on *nix servers
+            if os.path.isfile(peer_state.filepath) and\
+                    os.access(peer_state.filepath, os.R_OK):
+                        peer_state.file = open(peer_state.filepath, 'rb')
+                        peer_state.filesize = os.stat(peer_state.filepath)[stat.ST_SIZE]
+            else:
+                # The file doen't exist, try the dynamic_file_handler
+                # if it is set
+                if hasattr(self, 'dynamic_file_handler') and\
+                    self.dynamic_file_handler is not None:
+                        # we send the original requested filename to the handler
+                        peer_state.file, peer_state.filesize = self.dynamic_file_handler(filenameOrig)
+                else:
+                    raise IOError('Cannot access file: %s' % filenameOrig)
+
             peer_state.packetnum = 0
             peer_state.state = state.STATE_SEND
 
@@ -468,7 +500,7 @@ def run(self):
 
 class TFTPServer(object):
     def __init__(self, iface, root, port=_PTFTPD_DEFAULT_PORT,
-                 strict_rfc1350=False, notification_callbacks={}):
+                 strict_rfc1350=False, notification_callbacks={}, dynamic_file_callback=None):
         self.iface, self.root, self.port, self.strict_rfc1350 = \
                 iface, root, port, strict_rfc1350
         self.client_registry = {}
@@ -477,8 +509,11 @@ def __init__(self, iface, root, port=_PTFTPD_DEFAULT_PORT,
             raise TFTPServerConfigurationError(
                 "The specified TFTP root does not exist")
 
+        class _TFTPServerHandler(TFTPServerHandler):
+            dynamic_file_handler = dynamic_file_callback
+
         self.ip, self.netmask, self.mac = get_ip_config_for_interface(self.iface)
-        self.server = SocketServer.UDPServer((self.ip, port), TFTPServerHandler)
+        self.server = SocketServer.UDPServer((self.ip, port), _TFTPServerHandler)
         self.server.root = self.root
         self.server.strict_rfc1350 = self.strict_rfc1350
         self.server.clients = self.client_registry