Introduce virtual file system abstraction

COPYING

@@ -1,6 +1,7 @@
 The MIT License
 
 Copyright (c) 2009 Michael P. Soulier
+Copyright (c) 2012 Fabian Knittel <[email protected]>
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

bin/tftpy_server.py

@@ -21,7 +21,8 @@ def main():
                       '--root',
                       type='string',
                       help='path to serve from',
-                      default=None)
+                      default=None,
+                      action="append")
     parser.add_option('-d',
                       '--debug',
                       action='store_true',
@@ -38,7 +39,11 @@ def main():
         parser.print_help()
         sys.exit(1)
 
-    server = tftpy.TftpServer(options.root)
+    vfs = tftpy.TftpVfsStack()
+    for root in options.root:
+        vfs.mount(tftpy.TftpVfsNative(root), '/')
+
+    server = tftpy.TftpServerVfs(vfs)
     try:
         server.listen(options.ip, options.port)
     except tftpy.TftpException, err:

t/test.py

@@ -5,6 +5,8 @@
 import tftpy
 import os
 import time
+import tempfile
+import shutil
 
 log = tftpy.log
 
@@ -216,12 +218,12 @@ def testServerNoOptions(self):
         raddress = '127.0.0.2'
         rport = 10000
         timeout = 5
-        root = os.path.dirname(os.path.abspath(__file__))
+        vfs = tftpy.TftpVfsCompat(os.path.dirname(os.path.abspath(__file__)))
         # Testing without the dyn_func_file set.
         serverstate = tftpy.TftpContextServer(raddress,
                                               rport,
                                               timeout,
-                                              root)
+                                              vfs)
 
         self.assertTrue( isinstance(serverstate,
                                     tftpy.TftpContextServer) )
@@ -253,12 +255,12 @@ def testServerNoOptionsSubdir(self):
         raddress = '127.0.0.2'
         rport = 10000
         timeout = 5
-        root = os.path.dirname(os.path.abspath(__file__))
+        vfs = tftpy.TftpVfsCompat(os.path.dirname(os.path.abspath(__file__)))
         # Testing without the dyn_func_file set.
         serverstate = tftpy.TftpContextServer(raddress,
                                               rport,
                                               timeout,
-                                              root)
+                                              vfs)
 
         self.assertTrue( isinstance(serverstate,
                                     tftpy.TftpContextServer) )
@@ -286,5 +288,198 @@ def testServerNoOptionsSubdir(self):
         finalstate = serverstate.state.handle(ack, raddress, rport)
         self.assertTrue( finalstate is None )
 
+    def testServerInsecurePath(self):
+        raddress = '127.0.0.2'
+        rport = 10000
+        timeout = 5
+        vfs = tftpy.TftpVfsCompat(os.path.dirname(os.path.abspath(__file__)))
+        serverstate = tftpy.TftpContextServer(raddress,
+                                              rport,
+                                              timeout,
+                                              vfs)
+        rrq = tftpy.TftpPacketRRQ()
+        rrq.filename = '../setup.py'
+        rrq.mode = 'octet'
+        rrq.options = {}
+
+        # Start the download.
+        self.assertRaisesRegexp(tftpy.TftpException, "bad file path",
+                serverstate.start, rrq.encode().buffer)
+
+    def testServerSecurePath(self):
+        raddress = '127.0.0.2'
+        rport = 10000
+        timeout = 5
+        vfs = tftpy.TftpVfsCompat(os.path.dirname(os.path.abspath(__file__)))
+        serverstate = tftpy.TftpContextServer(raddress,
+                                              rport,
+                                              timeout,
+                                              vfs)
+        rrq = tftpy.TftpPacketRRQ()
+        rrq.filename = '100KBFILE'
+        rrq.mode = 'octet'
+        rrq.options = {}
+
+        # Start the download.
+        serverstate.start(rrq.encode().buffer)
+        # Should be in expectack state.
+        self.assertTrue(isinstance(serverstate.state,
+                                    tftpy.TftpStateExpectACK))
+
+class TestTftpyVfsReadOnlyDynFileFunc(unittest.TestCase):
+    def testRead(self):
+        state = {'called':False}
+        the_path = 'a path'
+        def dyn_func(path):
+            state['called'] = True
+            return path
+        vfs = tftpy.TftpVfsReadOnlyDynFileFunc(dyn_func)
+        ret = vfs.open_read(the_path)
+        self.assertEqual(the_path, ret)
+        self.assert_(state['called'])
+
+    def testWrite(self):
+        state = {'called':False}
+        the_path = 'a-path'
+        def dyn_func(path):
+            state['called'] = True
+            return path
+        vfs = tftpy.TftpVfsReadOnlyDynFileFunc(dyn_func)
+        ret = vfs.open_write(the_path)
+        self.assertEqual(None, ret)
+        self.assert_(not state['called'])
+
+class TestTftpyVfsNative(unittest.TestCase):
+    def setUp(self):
+        self.write_root = tempfile.mkdtemp()
+    def tearDown(self):
+        shutil.rmtree(self.write_root, ignore_errors=True)
+
+    def testReadExisting(self):
+        # Copy file to the temporary tftp root
+        root = os.path.dirname(os.path.abspath(__file__))
+        the_path = '100KBFILE'
+        shutil.copy(os.path.join(root, the_path), self.write_root)
+
+        vfs = tftpy.TftpVfsNative(self.write_root)
+        fp = vfs.open_read(the_path)
+        self.assert_(fp is not None)
+        try:
+            orig_fp = open(os.path.join(self.write_root, the_path), 'rb')
+            try:
+                self.assertEqual(orig_fp.read(), fp.read())
+            finally:
+                orig_fp.close()
+        finally:
+            fp.close()
+
+    def testReadNonExistent(self):
+        # The temporary tftp root is empty
+        the_path = '100KBFILE'
+
+        vfs = tftpy.TftpVfsNative(self.write_root)
+        fp = vfs.open_read(the_path)
+        self.assert_(fp is None)
+
+    def testNonExistentRoot(self):
+        non_existent_root = os.path.join(self.write_root, 'non-existent')
+        self.assertRaisesRegexp(tftpy.TftpException, 'tftproot does not exist',
+                tftpy.TftpVfsNative, non_existent_root)
+
+    def testWriteSubdir(self):
+        """Write a test string and read it back."""
+        the_dir = 'a-path'
+        the_fn = os.path.join(the_dir, 'a-file')
+        data = 'test string'
+        vfs = tftpy.TftpVfsNative(self.write_root)
+        fp = vfs.open_write(the_fn)
+        self.assert_(fp is not None)
+        fp.write(data)
+        fp.close()
+        self.assert_(os.path.exists(os.path.join(self.write_root, the_dir)))
+        self.assert_(os.path.isdir(os.path.join(self.write_root, the_dir)))
+        self.assert_(os.path.exists(os.path.join(self.write_root, the_fn)))
+
+    def testUnsafeRead(self):
+        the_path = os.path.join(os.path.pardir, '100KBFILE')
+        vfs = tftpy.TftpVfsNative(self.write_root)
+        self.assertRaisesRegexp(tftpy.TftpException, "bad file path",
+                vfs.open_read, the_path)
+
+class TestTftpyVfsStack(unittest.TestCase):
+    def setUp(self):
+        self.vfs = tftpy.TftpVfsStack()
+
+    def testReadEmpty(self):
+        self.assert_(self.vfs.open_read('path') is None)
+
+    def testWriteEmpty(self):
+        self.assert_(self.vfs.open_write('path') is None)
+
+    class MockVfsAccept(object):
+        def __init__(self):
+            self.read_fp = object()
+            self.write_fp = object()
+            self.path = None
+        def open_read(self, path):
+            self.path = path
+            return self.read_fp
+        def open_write(self, path):
+            self.path = path
+            return self.write_fp
+
+    class MockVfsReject(object):
+        def __init__(self):
+            self.path = None
+        def open_read(self, path):
+            self.path = path
+            return None
+        def open_write(self, path):
+            self.path = path
+            return None
+
+    def testReadRoot(self):
+        fs1 = self.MockVfsAccept()
+        self.vfs.mount(fs1, '/')
+        ret = self.vfs.open_read('path')
+        self.assert_(ret is fs1.read_fp)
+        self.assertEqual('/path', fs1.path)
+
+    def testWriteRoot(self):
+        fs1 = self.MockVfsAccept()
+        self.vfs.mount(fs1, '/')
+        ret = self.vfs.open_write('path')
+        self.assert_(ret is fs1.write_fp)
+        self.assertEqual('/path', fs1.path)
+
+    def testFirstRoot(self):
+        """Return first valid match"""
+        fs1 = self.MockVfsReject()
+        fs2 = self.MockVfsAccept()
+        fs3 = self.MockVfsAccept()
+        self.vfs.mount(fs1, '/')
+        self.vfs.mount(fs2, '/')
+        self.vfs.mount(fs3, '/')
+        ret = self.vfs.open_read('path')
+        self.assert_(ret is fs2.read_fp)
+        self.assert_(ret is not fs3.read_fp)
+        self.assertEqual('/path', fs1.path)
+        self.assertEqual('/path', fs2.path)
+        self.assertEqual(None, fs3.path)
+
+    def testIterateSubPaths(self):
+        """Visit all providers that have a matching base path."""
+        fs1 = self.MockVfsReject()
+        fs2 = self.MockVfsReject()
+        fs3 = self.MockVfsReject()
+        self.vfs.mount(fs1, '/base')
+        self.vfs.mount(fs2, '/base/somewhere')
+        self.vfs.mount(fs3, '/not-relevant')
+        ret = self.vfs.open_read('/base/somewhere/path')
+        self.assert_(ret is None)
+        self.assertEqual('/somewhere/path', fs1.path)
+        self.assertEqual('/path', fs2.path)
+        self.assertEqual(None, fs3.path)
+
 if __name__ == '__main__':
     unittest.main()

tftpy/TftpContexts.py

@@ -67,7 +67,7 @@ def add_dup(self, pkt):
 class TftpContext(object):
     """The base class of the contexts."""
 
-    def __init__(self, host, port, timeout, dyn_file_func=None):
+    def __init__(self, host, port, timeout):
         """Constructor for the base context, setting shared instance
         variables."""
         self.file_to_transfer = None
@@ -94,7 +94,6 @@ def __init__(self, host, port, timeout, dyn_file_func=None):
         self.last_update = 0
         # The last packet we sent, if applicable, to make resending easy.
         self.last_pkt = None
-        self.dyn_file_func = dyn_file_func
         # Count the number of retry attempts.
         self.retry_count = 0
 
@@ -194,18 +193,17 @@ def cycle(self):
 
 class TftpContextServer(TftpContext):
     """The context for the server."""
-    def __init__(self, host, port, timeout, root, dyn_file_func=None):
+    def __init__(self, host, port, timeout, vfs):
         TftpContext.__init__(self,
                              host,
                              port,
                              timeout,
-                             dyn_file_func
                              )
+        self._vfs = vfs
+
         # At this point we have no idea if this is a download or an upload. We
         # need to let the start state determine that.
         self.state = TftpStateServerStart(self)
-        self.root = root
-        self.dyn_file_func = dyn_file_func
 
     def __str__(self):
         return "%s:%s %s" % (self.host, self.port, self.state)
@@ -237,6 +235,18 @@ def end(self):
         log.debug("Set metrics.end_time to %s" % self.metrics.end_time)
         self.metrics.compute()
 
+    def open_read(self, path):
+        """Return a file-like object for the virtual path `path` or ``None`` if
+        the path does not exist.  Throws :class:`TftpException` in case the path
+        is invalid."""
+        return self._vfs.open_read(path)
+
+    def open_write(self, path):
+        """Return a file-like object for the virtual path `path` or ``None`` if
+        the path is not writable.  Throws :class:`TftpException` in case the
+        path is invalid."""
+        return self._vfs.open_write(path)
+
 class TftpContextClientUpload(TftpContext):
     """The upload context for the client during an upload.
     Note: If input is a hyphen, then we will use stdin."""