lxc-network/lxc_net/container.py

""" Container objects """

from . import settings
from . import util
from .xml_template import XMLTemplate

import libvirt
import tempfile
import subprocess
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,
        ]

        subprocess.run(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,
        ]
        subprocess.run(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(),
        ]
        subprocess.run(command_chown)
        command_chmod = [
            "sudo",
            "chmod",
            "-R",
            "700",
            self.temp_dir.resolve(),
        ]
        subprocess.run(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)):
        super().__init__(conn)

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

        self.networks = networks
        self.mem = mem

        self.overlay_root = None

        self.lxc_container = 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(self):
        if self.lxc_container:
            raise self.AlreadyExists()

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

        xml = XMLTemplate("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 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

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

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

    def __del__(self):
        self.cleanup()