[PATCH 5/7] Refactoring of Disk class

[Michael Tremer <michael.tremer@ipfire.org>]

[-- Attachment #1: Type: text/plain, Size: 5643 bytes --]

This is an example of how the code could take advantage of
Python contexts. The with statement will open the context
and when it is left (either normally or because of an exception),
the __exit__() method is called and can be used to cleanup.

Cleaning up is in this case umounting the image file, but could
also be used to shutdown and destroy a virtual machine or
network.

Signed-off-by: Michael Tremer <michael.tremer(a)ipfire.org>
---
 src/nitsi/disk.py    | 92 +++++++++++++++++++++++++++++---------------
 src/nitsi/machine.py | 13 ++-----
 2 files changed, 65 insertions(+), 40 deletions(-)

diff --git a/src/nitsi/disk.py b/src/nitsi/disk.py
index 29782574282e..5f465956c849 100755
--- a/src/nitsi/disk.py
+++ b/src/nitsi/disk.py
@@ -8,43 +8,75 @@ import tempfile
 
 logger = logging.getLogger("nitsi.disk")
 
+class Disk(object):
+    def __init__(self, path, part_uuid):
+        self.path = path
+        self.part_uuid = part_uuid
 
-class disk():
-    def __init__(self, disk):
-        self.log = logger.getChild(os.path.basename(disk))
-        self.log.debug("Initiated a disk class for {}".format(disk))
+        self.log = logger.getChild(os.path.basename(self.path))
+        self.log.debug("Initiated a disk class for {}".format(self.path))
+
+        # Create GuestFS object for this disk
         self.con = guestfs.GuestFS(python_return_dict=True)
-        self.con.add_drive_opts(disk, format="qcow2")
+        self.con.add_drive_opts(self.path, format="qcow2")
 
-    def mount(self, uuid, path):
-        self.log.debug("Trying to mount the partion with uuid: {} under {}".format(uuid, path))
+    def __enter__(self):
+        # Launch the GuestFS backend
         self.con.launch()
-        part = self.con.findfs_uuid(uuid)
-        self.con.mount(part, path)
 
-    def copy_in(self, fr, to):
-        self.log.debug("Going to copy some files into the image.")
-        tmp = tempfile.mkstemp()
-        tmp = tmp[1] + ".tar"
-        with tarfile.open(tmp, "w") as tar:
-            for file in fr:
-                self.log.debug("Adding {} to be copied into the image".format(file))
-                tar.add(file, arcname=os.path.basename(file))
+        # Find partition to mount
+        part = self.con.findfs_uuid(self.part_uuid)
+        if not part:
+            raise RuntimeError("Could not find partition %s" % self.part_uuid)
 
-        self.log.debug("Going to copy the files into the image")
-        self.con.tar_in_opts(tmp, to)
+        # Mounting partition
+        self._mount(part)
 
-    def umount(self, path):
-        self.log.debug("Unmounting the image")
-        self.con.umount_opts(path)
+        return MountedDisk(self, part)
 
-    def close(self):
-        self.log.debug("Flush the image and closing the connection")
+    def __exit__(self, type, value, traceback):
+        # Umount the disk image
+        self._umount()
+
+        # Shuts down the GuestFS backend
         self.con.shutdown()
-        self.con.close()
 
-# test  = disk("/var/lib/libvirt/images/alice.qcow2")
-# test.mount("45598e92-3487-4a1b-961d-79aa3dd42a7d", "/")
-# test.copy_in("/home/jonatan/nitsi/libguestfs-test", "/root/")
-# test.umount("/")
-# test.close()
+    def _mount(self, part, path="/"):
+        """
+            Mounts the selected partition to path
+        """
+        self.log.debug("Mounting partition %s to %s" % (part, path))
+
+        self.con.mount(part, path)
+
+    def _umount(self, path="/"):
+        self.log.debug("Umounting %s" % path)
+
+        self.con.umount_opts(path)
+
+
+class MountedDisk(object):
+    """
+        Provides commands that can only be executed when the disk is mounted.
+
+        Leaving the context will automatically umount the disk.
+    """
+    def __init__(self, disk, part):
+        self.disk = disk
+        self.part = part
+
+    def copy_in(self, files, destination):
+        self.disk.log.debug("Going to copy some files into the image")
+
+        with tempfile.NamedTemporaryFile("wb", suffix=".tar") as f:
+            with tarfile.open(fileobj=f, mode="w") as t:
+                for file in files:
+                    self.disk.log.debug("Adding %s" % file)
+                    t.add(file, arcname=os.path.basename(file))
+
+            self.disk.log.debug("Going to copy the files into the image")
+            self.disk.con.tar_in_opts(f.name, destination)
+
+
+# with disk("/var/lib/libvirt/images/alice.qcow2", "45598e92-3487-4a1b-961d-79aa3dd42a7d") as test:
+#     test.copy_in("/home/jonatan/nitsi/libguestfs-test", "/root/")
diff --git a/src/nitsi/machine.py b/src/nitsi/machine.py
index 0b4617aed8dc..9efa4dde5688 100644
--- a/src/nitsi/machine.py
+++ b/src/nitsi/machine.py
@@ -40,8 +40,7 @@ class machine():
         if not os.path.isfile(self.image):
             self.log.error("No such file: {}".format(self.image))
 
-        self.root_uid = root_uid
-        self.disk = disk.disk(image)
+        self.disk = disk.Disk(image, root_uid)
 
         self.username = username
         self.password = password
@@ -132,11 +131,5 @@ class machine():
         return self.serial_con.command(cmd)
 
     def copy_in(self, fr, to):
-        try:
-            self.disk.mount(self.root_uid, "/")
-            self.disk.copy_in(fr, to)
-        except BaseException as e:
-            self.log.error(e)
-        finally:
-            self.disk.umount("/")
-            self.disk.close()
+        with self.disk as d:
+            d.copy_in(fr, to)
-- 
2.17.0