src/nitsi/disk.py

   1 #!/usr/bin/python3
   2
   3 import collections
   4 import guestfs
   5 import logging
   6 import os
   7 import tarfile
   8 import tempfile
   9
  10 logger = logging.getLogger("nitsi.disk")
  11
  12 class DiskExeption(Exception):
  13     def __init__(self, message):
  14         self.message = message
  15
  16 class Disk():
  17     def __init__(self, disk):
  18         self.log = logger.getChild(os.path.basename(disk))
  19         self.log.debug("Initiated a disk class for {}".format(disk))
  20         self.con = guestfs.GuestFS(python_return_dict=True)
  21         self.con.add_drive_opts(disk, format="qcow2")
  22
  23     def mount(self):
  24         self.log.info("Trying to mount all partions of the found operating system")
  25         for mountpoint, partition in self.mountpoints.items():
  26             self.log.info("Trying to mount the partion {} under {}".format(partition, mountpoint))
  27             self.con.mount(partition, mountpoint)
  28
  29     def inspect(self):
  30         self.con.launch()
  31
  32         # This inspects the given Disk and
  33         # returns all root filesystems of the operating systems on this disk
  34         root_filesystem = self.con.inspect_os()
  35
  36         # If we found more the one rootfile system we throw an error
  37         # because we do not support multi boot systems
  38         if len(root_filesystem) > 1:
  39             raise DiskExeption("Found more then one operating system on the given disk.")
  40
  41         # We cannot go on if we found no operating system
  42         if len(root_filesystem) == 0:
  43             raise DiskExeption("Found no operating system on the given disk.")
  44
  45         root_filesystem = root_filesystem.pop()
  46
  47         # Get mountpoints of all filesystems accociated with this operationg system
  48         tmp = self.con.inspect_get_mountpoints(root_filesystem)
  49
  50         # Sort mountpoints by length to mount / before /usr
  51         # We are using an ordered dict here to keep the order of the keys
  52         self.mountpoints = collections.OrderedDict(sorted(tmp.items(), key=lambda t: len(t[0])))
  53
  54         self.log.debug(self.mountpoints)
  55
  56         # Print some nice information about the found operationg system
  57         self.log.info("Arch of the installed operating system: {}".format(
  58             self.con.inspect_get_arch(root_filesystem)))
  59         self.log.info("Distribution of the installed operating system: {}".format(
  60             self.con.inspect_get_distro(root_filesystem)))
  61
  62
  63     def copy_in(self, fr, to):
  64         self.log.info("Going to copy some files into the image.")
  65         tmp = tempfile.mkstemp()
  66         tmp = tmp[1] + ".tar"
  67         self.log.debug("Path of tarfile is: {}".format(tmp))
  68         with tarfile.open(tmp, "w") as tar:
  69             for file in fr:
  70                 self.log.debug("Adding {} to be copied into the image".format(file))
  71                 tar.add(file, arcname=os.path.basename(file))
  72
  73         self.log.info("Going to copy the files into the image")
  74         self.con.tar_in_opts(tmp, to)
  75         self.log.debug(self.con.ls(to))
  76
  77     def umount(self):
  78         self.log.info("Trying to unmount all partions of the found operating system")
  79         for mountpoint, partition in reversed(self.mountpoints.items()):
  80             self.log.info("Trying to unmount the partion {} under {}".format(partition, mountpoint))
  81             self.con.umount_opts(mountpoint)
  82
  83     def close(self):
  84         self.log.info("Flush the image and closing the connection")
  85         self.con.shutdown()
  86         self.con.close()
  87
  88 # test  = Disk("/var/lib/libvirt/images/ipfire-bob.qcow2")
  89 # test.mount("1efb5389-0949-46bb-b688-5246acba9f6d", "/")
  90 # test.copy_in("/home/jonatan/nitsi/libguestfs-test", "/root/")
  91 # test.umount("/")
  92 # test.close()