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)
