Spotter-VM/build/usr/lib/python3.6/lxcbuild/imagebuilder.py

# -*- coding: utf-8 -*-

import os
import shutil
import subprocess
import sys

from lxcmgr import lxcmgr
from lxcmgr.paths import LXC_STORAGE_DIR

class ImageExistsError(Exception):
    pass

class ImageNotFoundError(Exception):
    pass

class ImageBuilder:
    def __init__(self, image):
        self.image = image
        self.script = []
        self.script_eof = None

    def build(self):
        with open(self.image.lxcfile, 'r') as f:
            for line in f:
                line = line.strip()
                if self.script_eof:
                    if line == self.script_eof:
                        self.script_eof = None
                        self.run_script(self.script)
                    else:
                        self.script.append(line)
                elif line:
                    self.process_line(*line.split(None, 1))

    def process_line(self, directive, args):
        if 'RUN' == directive:
            self.script = []
            self.script_eof = args
        elif 'IMAGE' == directive:
            self.set_name(args)
        elif 'LAYER' == directive:
            self.add_layer(args)
        elif 'MERGE' == directive:
            self.merge_layers(args.split())
        elif 'COPY' == directive:
            srcdst = args.split()
            self.copy_files(srcdst[0], srcdst[1] if len(srcdst) == 2 else '')
        elif 'ENV' == directive:
            self.add_env(*args.split(None, 1))
        elif 'USER' == directive:
            self.set_user(*args.split())
        elif 'CMD' == directive:
            self.set_cmd(args)
        elif 'WORKDIR' == directive:
            self.set_cwd(args)
        elif 'HALT' == directive:
            self.set_halt(args)
        elif 'READY' == directive:
            self.set_ready(args)

    def get_layer_path(self, layer):
        return os.path.join(LXC_STORAGE_DIR, layer)

    def run_script(self, script):
        lxcmgr.create_container(self.image.name, self.image.conf)
        sh = os.path.join(LXC_STORAGE_DIR, self.image.name, 'run.sh')
        with open(sh, 'w') as f:
            f.write('#!/bin/sh\nset -ev\n\n{}\n'.format('\n'.join(script)))
        os.chmod(sh, 0o700)
        os.chown(sh, 100000, 100000)
        subprocess.run(['lxc-execute', self.image.name, '--', '/bin/sh', '-lc', '/run.sh'], check=True)
        os.unlink(sh)
        if not self.image.scratch_build:
            lxcmgr.destroy_container(self.image.name)

    def set_name(self, name):
        self.image.name = name
        self.image.conf['layers'] = [name]
        image_path = self.get_layer_path(name)
        if os.path.exists(image_path):
            if self.image.force_build:
                self.clean()
            else:
                raise ImageExistsError(image_path)
        os.makedirs(image_path, 0o755, True)
        os.chown(image_path, 100000, 100000)

    def add_layer(self, name):
        layer_path = self.get_layer_path(name)
        if not os.path.exists(layer_path):
            raise ImageNotFoundError(layer_path)
        self.image.conf['layers'].insert(1, name)

    def merge_layers(self, cmd):
        layers = [self.get_layer_path(layer) for layer in self.image.conf['layers']]
        subprocess.run(cmd + layers[::-1], check=True)

    def copy_files(self, src, dst):
        dst = os.path.join(LXC_STORAGE_DIR, self.image.name, dst)
        if src.startswith('http://') or src.startswith('https://'):
            unpack_http_archive(src, dst)
        else:
            copy_tree(os.path.join(self.image.build_dir, src), dst)
        shift_uid(dst)

    def add_env(self, key, value):
        if 'env' not in self.image.conf:
            self.image.conf['env'] = []
        self.image.conf['env'].append([key, value])

    def set_user(self, uid, gid):
        self.image.conf['uid'] = uid
        self.image.conf['gid'] = gid

    def set_cmd(self, cmd):
        self.image.conf['cmd'] = cmd

    def set_cwd(self, cwd):
        self.image.conf['cwd'] = cwd

    def set_halt(self, halt):
        self.image.conf['halt'] = halt

    def set_ready(self, cmd):
        self.image.conf['ready'] = cmd

    def clean(self):
        lxcmgr.destroy_container(self.image.name)
        shutil.rmtree(self.get_layer_path(self.image.name))

def unpack_http_archive(src, dst):
        xf = 'xzf'
        if src.endswith('.bz2'):
            xf = 'xjf'
        elif src.endswith('.xz'):
            xf = 'xJf'
        with subprocess.Popen(['wget', src, '-O', '-'], stdout=subprocess.PIPE) as wget:
            with subprocess.Popen(['tar', xf, '-', '-C', dst], stdin=wget.stdout) as tar:
                wget.stdout.close()
                tar.wait()

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)

def shift_uid(dir):
    shift_uid_entry(dir, os.stat(dir, follow_symlinks=True))
    shift_uid_recursively(dir)

def shift_uid_recursively(dir):
    for entry in os.scandir(dir):
        shift_uid_entry(entry.path, entry.stat(follow_symlinks=False))
        if entry.is_dir():
            shift_uid_recursively(entry.path)

def shift_uid_entry(path, stat):
    uid = stat.st_uid
    gid = stat.st_gid
    do_chown = False
    if uid < 100000:
        uid = uid + 100000
        do_chown = True
    if gid < 100000:
        gid = gid + 100000
        do_chown = True
    if do_chown:
        os.lchown(path, uid, gid)