Refactor disk class
[nitsi.git] / 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()