lxc-network/lxc_net/container.py

""" Container objects """

from . import settings
from . import util
from .jinja_template import JinjaTemplate

import libvirt
import tempfile
from pathlib import Path
import getpass
import weakref


class OverlayDirectory:
    """ Allocates an overlayfs with a given lower dir and a temporary upper dir. The
    overlayfs is destroyed when this object is deallocated """

    class DoesNotExist(Exception):
        """ Raised if `lower` does not exist """

        def __init__(self, path):
            self.path = path

        def __str__(self):
            return "Path {} does not exist".format(self.path)

    def __init__(self, lower, name=None):
        if not name:
            name = ""

        self.lower = Path(lower)

        if not self.lower.exists():
            raise self.DoesNotExist(self.lower)

        root_temp_dir = Path(settings.OVERLAYFS_BASE_DIR)
        if not root_temp_dir.exists():
            root_temp_dir.mkdir()

        self.overlay_temp_dir = tempfile.TemporaryDirectory(
            prefix=name + ("_" if name else ""), dir=root_temp_dir.resolve(),
        )
        self.temp_dir = Path(self.overlay_temp_dir.name)
        self.temp_dir_cleaner = None

        self.upper = self.temp_dir / "upper"
        self.work = self.temp_dir / "work"
        self.mount_point = self.temp_dir / "mount"

        self.upper.mkdir()
        self.work.mkdir()
        self.mount_point.mkdir()

        self.mounted = False

        self._mount()

    def _mount(self):
        if self.mounted:
            return

        command = [
            "sudo",
            "mount",
            "-t",
            "overlay",
            "overlay",
            "-o",
            "lowerdir={lower},upperdir={upper},workdir={work}".format(
                lower=self.lower.resolve(),
                upper=self.upper.resolve(),
                work=self.work.resolve(),
            ),
            self.mount_point,
        ]

        util.run_cmd_retry(command)

        self.mounted = True

        self.temp_dir_cleaner = weakref.finalize(
            self.overlay_temp_dir, self.cleanup_mount
        )

    def _umount(self):
        if not self.mounted:
            return

        command = [
            "sudo",
            "umount",
            self.mount_point,
        ]
        util.run_cmd_retry(command)
        self.mounted = False

    def _set_perms_for_cleanup(self):
        """ Sets the permissions to allow cleanup """

        if not self.temp_dir.exists():
            return  # Already deleted

        # The overlayfs runs as root, so we need to allow our current user to delete
        # the temp directory files
        command_chown = [
            "sudo",
            "chown",
            "-R",
            getpass.getuser(),
            self.temp_dir.resolve(),
        ]
        util.run_cmd_retry(command_chown)

        command_chmod = [
            "sudo",
            "chmod",
            "-R",
            "700",
            self.temp_dir.resolve(),
        ]
        util.run_cmd_retry(command_chmod)

    def cleanup_mount(self):
        self._umount()
        self._set_perms_for_cleanup()


class Container(util.LibvirtObject):
    class AlreadyExists(Exception):
        def __str__(self):
            return "This container is already instanciated"

    def __init__(self, conn, networks, name=None, mem=int(1e6), enable_v4=True):
        """ Parameters:
        * conn: connection to libvirt,
        * networks: iterable of Network instances this container is connected to,
        * name: name of the container. Defaults to something id-based.
        * mem: KiB of memory available to this container,
        * enable_v4: is IPv4 enabled for this container? Defaults to True. If False,
          this container won't be given an IPv4 address.
        """

        super().__init__(conn)

        if not name:
            name = str(self.id)
        self.name = settings.PREFIX + "_dom_" + name

        self.networks = networks
        self.mem = mem
        self.enable_v4 = enable_v4

        self.overlay_root = None

        self.lxc_container = None

    @property
    def up(self):
        return self.lxc_container is not None

    def get_jinja_networks(self):
        """ Get a jinja2 environment-compatible network list """
        return [
            {"mac": util.MACAddress(net.id, self.id), "name": net.name, "net": net}
            for net in self.networks
        ]

    def _create_network_files(self):
        """ Creates systemd-networkd .network files to set IP addresses """
        if not self.overlay_root:
            raise Exception("No root directory specified yet")
        net_config_templ = JinjaTemplate("nic.network")
        net_conf_dir = self.overlay_root.mount_point / "etc/systemd/network/"
        for net in self.networks:
            net_config = net_config_templ.inst(
                mac=util.MACAddress(net.id, self.id),
                ipv4=util.Addrv4(net.id, self.id)
                if self.enable_v4 and net.enable_v4
                else None,
                ipv6=util.Addrv6(net.id, self.id),
            )
            net_config_path = net_conf_dir / "11-{link:02d}-{name}.network".format(
                link=net.id, name=net.name
            )
            with open(net_config_path, "w") as handle:
                handle.write(net_config)

    def create(self):
        if self.lxc_container:
            raise self.AlreadyExists()

        self.overlay_root = OverlayDirectory(settings.BASE_SYSTEM_ROOT, name=self.name)

        self._create_network_files()

        xml = JinjaTemplate("container.xml").inst(
            name=self.name,
            uuid=self.uuid,
            mem=self.mem,
            filesystem_root=self.overlay_root.mount_point,
            networks=self.get_jinja_networks(),
        )

        self.lxc_container = self.conn.createXML(xml)

    def notify_cleanup(self):
        """ This method can be called before `cleanup` to notify the host that a
        cleanup will be performed soon, speeding the process up. """
        if self.lxc_container:
            try:
                self.lxc_container.shutdown()
            except libvirt.libvirtError as exn:
                if not str(exn).startswith("Domain not found:"):
                    raise exn
                # Else, the machine was already stopped: everything is fine

    def cleanup(self):
        if self.lxc_container:
            try:
                self.lxc_container.destroy()
            except libvirt.libvirtError as exn:
                if not str(exn).startswith("Domain not found:"):
                    raise exn
                # Else, the machine was already stopped: everything is fine
            self.lxc_container = None
        if self.overlay_root:
            self.overlay_root.cleanup_mount()
            self.overlay_root.overlay_temp_dir.cleanup()
            self.overlay_root = None

    def __enter__(self):
        self.create()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()

    def __del__(self):
        self.cleanup()