#!/usr/bin/python3

import os
import shutil
import subprocess
import sys

LXC_ROOT = '/var/lib/lxc'
CONFIG_TEMPLATE = '''# Image name
lxc.uts.name = {name}

# Network
lxc.net.0.type = veth
lxc.net.0.link = lxcbr0
lxc.net.0.flags = up

# Volumes
lxc.rootfs.path = {rootfs}

# Mounts
lxc.mount.entry = shm dev/shm tmpfs rw,nodev,noexec,nosuid,relatime,mode=1777,create=dir 0 0
lxc.mount.entry = /etc/hosts etc/hosts none bind 0 0
lxc.mount.entry = /etc/resolv.conf etc/resolv.conf none bind 0 0
{mounts}

# Init
lxc.init.cmd = {cmd}
lxc.init.uid = {uid}
lxc.init.gid = {gid}
lxc.init.cwd = {cwd}

# Environment
{env}

# Halt
lxc.signal.halt = SIGTERM

# Log
lxc.console.size = 1MB
lxc.console.logfile = /var/log/lxc/{name}.log

# Other
lxc.arch = x86_64
lxc.cap.drop = sys_admin
lxc.hook.start-host = /usr/bin/vmmgr register-container
lxc.hook.post-stop = /usr/bin/vmmgr unregister-container
lxc.include = /usr/share/lxc/config/common.conf
'''

class LXCImage:
    def __init__(self, build_path):
        self.name = None
        self.layers = []
        self.mounts = []
        self.env = []
        self.uid = 0
        self.gid = 0
        self.cmd = '/bin/true'
        self.cwd = '/'

        if os.path.isfile(build_path):
            self.lxcfile = os.path.realpath(build_path)
            self.build_dir = os.path.dirname(self.lxcfile)
        else:
            self.build_dir = os.path.realpath(build_path)
            self.lxcfile = os.path.join(self.build_dir, 'lxcfile')

    def build(self):
        with open(self.lxcfile, 'r') as fd:
            lxcfile = [l.strip() for l in fd.readlines()]

        script = []
        script_eof = None

        for line in lxcfile:
            if script_eof:
                if line == script_eof:
                    script_eof = None
                    self.run_script(script)
                elif in_script:
                    script.append(line)
            elif line.startswith('RUN'):
                script = []
                script_eof = line.split()[1]
            elif line.startswith('IMAGE'):
                self.set_name(line.split()[1])
            elif line.startswith('LAYER'):
                self.add_layer(line.split()[1])
            elif line.startswith('LAYERFIX'):
                self.fix_layers(line.split()[1])
            elif line.startswith('COPY'):
                srcdst = line.split()
                self.copy_files(srcdst[1], srcdst[2] if len(srcdst) == 3 else '')
            elif line.startswith('MOUNT'):
                srcdst = line.split()
                self.add_mount(srcdst[1], srcdst[2])
            elif line.startswith('ENV'):
                env = line.split()
                self.add_env(env[1], env[2])
            elif line.startswith('USER'):
                uidgid = line.split()
                self.set_user(uidgid[1], uidgid[2])
            elif line.startswith('CMD'):
                self.set_cmd(' '.join(line.split()[1:]))
            elif line.startswith('WORKDIR'):
                self.set_cwd(line.split()[1])
        # Add the final layer which can be treated as ephemeral
        self.add_layer('{}/delta0'.format(self.name))

    def rebuild_config(self):
        if len(self.layers) == 1:
            rootfs_path = self.layers[0]
        else:
            # Multiple lower overlayfs layers are ordered from right to left (lower2:lower1:rootfs:upper)
            rootfs_path = 'overlay:{}:{}'.format(':'.join(self.layers[:-1][::-1]), self.layers[-1])
        mount_entries = '\n'.join(['lxc.mount.entry = {} none bind 0 0'.format(m) for m in self.mounts])
        env_entries = '\n'.join(['lxc.environment = {}'.format(e) for e in self.env])
        with open(os.path.join(LXC_ROOT, self.name, 'config'), 'w') as fd:
            fd.write(CONFIG_TEMPLATE.format(name=self.name, rootfs=rootfs_path, mounts=mount_entries, env=env_entries, uid=self.uid, gid=self.gid, cmd=self.cmd, cwd=self.cwd))

    def run_script(self, script):
        sh = os.path.join(self.layers[-1], 'run.sh')
        with open(sh, 'w') as fd:
            fd.write('#!/bin/sh\nset -ev\n\n{}\n'.format('\n'.join(script)))
        os.chmod(sh, 0o700)
        subprocess.run(['lxc-execute', '-n', self.name, '--', '/bin/sh', '-lc', '/run.sh'], check=True)
        os.unlink(sh)

    def set_name(self, name):
        self.name = name
        os.makedirs(os.path.join(LXC_ROOT, self.name), 0o755, True)

    def add_layer(self, layer):
        layer = os.path.join(LXC_ROOT, layer)
        self.layers.append(layer)
        os.makedirs(layer, 0o755, True)
        self.rebuild_config()

    def fix_layers(self, cmd):
        subprocess.run([cmd]+self.layers, check=True)

    def copy_files(self, src, dst):
        src = os.path.join(self.build_dir, src)
        dst = os.path.join(self.layers[-1], dst)
        copy_tree(src, dst)

    def add_mount(self, src, dst):
        self.mounts.append('{} {}'.format(src, dst))
        self.rebuild_config()

    def add_env(self, key, value):
        self.env.append('{}={}'.format(key, value))
        self.rebuild_config()

    def set_user(self, uid, gid):
        self.uid = uid
        self.gid = gid
        self.rebuild_config()

    def set_cmd(self, cmd):
        self.cmd = cmd
        self.rebuild_config()

    def set_cwd(self, cwd):
        self.cwd = cwd
        self.rebuild_config()

def copy_tree(src, dst):
    if not os.path.isdir(src):
        shutil.copy2(src, dst)
    else:
        os.makedirs(dst, exist_ok=True)
        for name in os.listdir(src):
            copy_tree(os.path.join(src, name), os.path.join(dst, name))
    shutil.copystat(src, dst)

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print('Usage: lxc-build <buildpath>\n  where the buildpath can be either specific lxcfile or a directory containing one')
    else:
        i = LXCImage(sys.argv[1])
        i.build()