#!/usr/bin/env python3
import argparse
import collections
import itertools
import logging
import os
import random
import re
import shutil
import string
import subprocess
import sys
import tempfile


log = logging.getLogger('stage3img')


LVM_PV_PART = {
    'type': 0x8e00,
}

UNITS = {
    'k': 2 ** 10,
    'kb': 10 ** 3,
    'kib': 2 ** 10,
    'm': 2 ** 20,
    'mb': 10 ** 6,
    'mib': 2 ** 20,
    'g': 2 ** 30,
    'gb': 10 ** 9,
    'gib': 2 ** 30,
    't': 2 ** 40,
    'tb': 10 ** 12,
    'tib': 2 ** 40,
    'p': 2 ** 50,
    'pb': 10 ** 15,
    'pib': 2 ** 50,
    'e': 2 ** 60,
    'eb': 10 ** 18,
    'eib': 2 ** 60,
    'z': 2 ** 70,
    'zb': 10 ** 21,
    'zib': 2 ** 70,
    'y': 2 ** 80,
    'yb': 10 ** 24,
    'yib': 2 ** 80,
}

SIZE_RE = re.compile(
    r'^(?P<value>[0-9]+(?:\.[0-9]+)?)\s*'
    r'(?P<unit>\w+)?\s*$'
)

TRY_SSH_KEYS = [
    os.path.expanduser('~/.ssh/id_ed25519.pub'),
    os.path.expanduser('~/.ssh/id_rsa.pub'),
    os.path.expanduser('~/.ssh/id_ecdsa.pub'),
]


class CommandError(Exception):
    pass


class MaxLevelFilter(logging.Filter):
    '''Filter log records above a particular level

    :param max_level: The maximum level of log message to include
    '''

    def __init__(self, max_level=logging.WARNING, *args, **kwargs):
        logging.Filter.__init__(self, *args, **kwargs)
        self.max_level = max_level

    def filter(self, record):
        return record.levelno <= self.max_level


class Image(object):

    EXTRA_MOUNTS = [
        ('/tmp', 'tmpfs', '-t', 'tmpfs'),
        ('/run', 'tmpfs', '-t', 'tmpfs'),
        ('/dev', '/dev', '--rbind'),
        ('/proc', 'proc', '-t', 'proc'),
        ('/sys', '/sys', '--rbind'),
    ]
    SCRIPT_ENV = {
        'PATH': os.pathsep.join((
            '/usr/local/sbin',
            '/usr/sbin',
            '/sbin',
            '/usr/local/bin',
            '/usr/bin',
            '/bin',
        )),
    }

    default_fstype = 'ext4'

    def __init__(self, filename, fd=None):
        self.fd = fd
        self.filename = filename
        self.tempname = self.filename + '.tmp'
        self.loopdev = None
        self.partitions = None
        self.vgname = None
        self.volumes = None
        self.mountpoint = None
        self._extra_mounted = False

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        self.unmount()
        self.deactivate_vg()
        self.detach_loopback()
        if self.fd is not None:
            os.close(self.fd)
        if exc_type:
            log.error('Error occurred, removing image')
            try:
                os.unlink(self.tempname)
            except:
                log.exception('Failed to remove temporary image file:')
            try:
                os.unlink(self.filename)
            except:
                log.exception('Failed to remove image file:')
        else:
            os.rename(self.tempname, self.filename)


    @property
    def filesystems(self):
        if self.partitions:
            partnum = 1
            for part in self.partitions:
                if 'mountpoint' not in part:
                    continue
                dev = '{}p{}'.format(self.loopdev, partnum)
                yield (
                    part['mountpoint'],
                    dev,
                    part.get('fstype'),
                    part.get('fsopts', 'defaults'),
                )
                partnum += 1
        if self.volumes:
            for vol in self.volumes:
                if 'mountpoint' not in vol:
                    continue
                dev = '/dev/{}/{}'.format(self.vgname, vol['name'])
                yield (
                    vol['mountpoint'],
                    dev,
                    vol.get('fstype'),
                    vol.get('fsopts', 'defaults'),
                )

    @classmethod
    def create(cls, filename):
        fd = os.open(filename, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
        return cls(filename, fd)

    def resize(self, size):
        with open(self.tempname, 'wb') as f:
            f.truncate(size)

    def attach_loopback(self):
        self.loopdev = run_cmd('losetup', '-f', '--show', self.tempname)
        log.info('Connected {} to {}'.format(self.tempname, self.loopdev))

    def detach_loopback(self):
        if not self.loopdev:
            return
        log.info('Detaching {}'.format(self.loopdev))
        run_cmd('losetup', '-d', self.loopdev)

    def partition(self, partitions):
        if not self.loopdev:
            self.attach_loopback()
        log.info('Partitioning {} with GPT'.format(self.loopdev))
        cmd = ['sgdisk', '-a', '4096', '-Z', '-g']
        for idx, part in enumerate(partitions):
            partnum = idx + 1
            cmd += ('-n', '{}::{}'.format(partnum, part.get('size', '')))
            if 'type' in part:
                cmd += ('-t', '{}:{:X}'.format(partnum, part['type']))
            if 'name' in part:
                cmd += ('-c', '{}:{}'.format(partnum, part['name']))
        cmd.append(self.loopdev)
        run_cmd_logged(*cmd)
        self.partitions = partititons

    def setup_lvm(self, vgname, volumes):
        if not self.loopdev:
            self.attach_loopback()
        pvscan(self.loopdev)
        if self.partitions:
            pv = '{}p{}'.format(self.loopdev, len(self.partitions))
        else:
            pv = self.loopdev
        pvcreate(pv)
        vgcreate(vgname, pv)
        self.vgname = vgname
        for vol in volumes:
            lvcreate(vgname, vol['name'], vol.get('size'))
        self.volumes = volumes

    def deactivate_vg(self):
        if not self.vgname:
            return
        run_cmd_logged('vgchange', '-an', self.vgname)

    def make_filesystems(self):
        for mountpoint, dev, fstype, fsopts in self.filesystems:
            if not fstype:
                fstype = self.default_fstype
            log.info('Creating {} filesystem on {}'.format(fstype, dev))
            run_cmd('mkfs.{}'.format(fstype), dev)

    def mount(self):
        self.mountpoint = tempfile.mkdtemp()
        for mountpoint, dev, fstype, fsopts in sorted(self.filesystems):
            path = os.path.join(self.mountpoint, mountpoint[1:])
            if not os.path.isdir(path):
                os.makedirs(path)
            log.info('Mounting {} on {}'.format(dev, path))
            run_cmd_logged('mount', dev, path)

    def mount_extra(self):
        if self._extra_mounted:
            return
        for item in self.EXTRA_MOUNTS:
            mountpoint, dev, args = item[0], item[1], item[2:]
            path = os.path.join(self.mountpoint, mountpoint[1:])
            if not os.path.isdir(path):
                os.makedirs(path)
            run_cmd_logged('mount', dev, path, *args)

    def unmount(self):
        if not self.mountpoint:
            return
        log.info('Unmounting {}'.format(self.mountpoint))
        run_cmd_logged('umount', '-R', self.mountpoint)
        os.rmdir(self.mountpoint)
        self.mountpoint = None
        self._mount_extra = False

    def extract(self, filename):
        if not self.mountpoint:
            self.mount()
        log.info('Extracting {} to {}'.format(filename, self.mountpoint))
        run_cmd_logged('tar', '-xha', '--numeric-owner', '-f', filename,
                       '-C', self.mountpoint)

    def run_script(self, script, chroot=True):
        if not self._extra_mounted:
            self.mount_extra()
        name = os.path.basename(script)
        dest = os.path.join(self.mountpoint, 'tmp', name)
        log.debug('Copying {} to {}'.format(script, dest))
        shutil.copy(script, dest)
        os.chmod(dest, 0o0755)
        kwargs = {
            'env': self.SCRIPT_ENV.copy(),
        }
        if chroot:
            cmd = ('chroot', self.mountpoint, os.path.join('/tmp', name))
        else:
            cmd = (dest,)
            kwargs['env']['IMAGE_ROOT'] = self.mountpoint
            kwargs['cwd'] = self.mountpoint
        log.info('Running script: {}'.format(name))
        try:
            run_cmd_logged(*cmd, **kwargs)
        finally:
            os.unlink(dest)

    def write_fstab(self, tmpfstmp=True):
        tmpl = '{dev}\t{mountpoint}\t{fstype}\t{fsopts}\t0 {passno}\n'
        log.info('Writing fstab')
        with open(os.path.join(self.mountpoint, 'etc/fstab'), 'w') as fstab:
            for mountpoint, dev, fstype, fsopts in sorted(self.filesystems):
                line = tmpl.format(
                    dev=dev,
                    mountpoint=mountpoint,
                    fstype=fstype or self.default_fstype,
                    fsopts=fsopts,
                    passno=1 if mountpoint == '/' else 2,
                )
                log.debug(line.rstrip('\n'))
                fstab.write(line)
            if tmpfstmp:
                line = tmpl.format(
                    dev='tmpfs',
                    mountpoint='/tmp',
                    fstype='tmpfs',
                    fsopts='defaults',
                    passno=0,
                )
                log.debug(line.rstrip('\n'))
                fstab.write(line)


def gen_vgname(length=8):
    name = random.choice(string.ascii_lowercase)
    while len(name) < length - 1:
        name += random.choice(string.ascii_letters + string.digits)
    name += random.choice(string.ascii_lowercase)
    return name


def inject_ssh_keys(root, ssh_keys):
    ssh_dir = os.path.join(root, 'root', '.ssh')
    if not os.path.isdir(ssh_dir):
        os.makedirs(ssh_dir)
    ssh_keys = set(ssh_keys)
    try:
        ssh_keys.remove(None)
    except KeyError:
        pass
    else:
        for path in TRY_SSH_KEYS:
            if os.path.exists(path):
                ssh_keys.append(path)
                break
    with open(os.path.join(ssh_dir, 'authorized_keys'), 'a') as dest:
        for keyfile in ssh_keys:
            with open(keyfile) as src:
                pubkey = src.read().strip()
            log.info('Injecting SSH public key: {}'.format(keyfile))
            dest.write(pubkey + '\n')


def lvcreate(vg_name, name, size=None):
    cmd = ['lvcreate', '-n', name]
    if size is not None:
        if '%' in size:
            cmd += ('-l', size)
        else:
            cmd += ('-L', size)
    else:
        cmd += ('-l', '100%FREE')
    cmd.append(vg_name)
    run_cmd_logged(*cmd)


def parse_volumes(volstrings):
    pat = re.compile(r'(?<!\\)=')
    keys = ('size', 'fstype', 'name', 'fsopts', 'type')
    volumes = []
    for vol in volstrings:
        parts = pat.split(vol, 1)
        mountpoint = parts[0]
        try:
            pairs = itertools.zip_longest(keys, parts[1].split(';'))
            params = dict((k, v) for k, v in pairs if v not in ('', None))
        except IndexError:
            params = {}
        params['mountpoint'] = mountpoint
        if 'name' not in params:
            if mountpoint.startswith('/'):
                name = mountpoint[1:].replace('/', '-')
            else:
                name = mountpoint
            if name == '':
                name = 'root'
            params['name'] = name
        if mountpoint == '/boot':
            volumes.insert(0, (mountpoint, params))
        elif mountpoint == '/':
            volumes.insert(1, (mountpoint, params))
        else:
            volumes.append((mountpoint, params))
    return collections.OrderedDict(volumes)


def parse_size(size):
    if isinstance(size, int):
        return size

    m = SIZE_RE.match(size.lower())
    if not m:
        raise ValueError('Invalid size: {}'.format(size))
    parts = m.groupdict()
    if parts['unit'] in (None, 'b', 'byte', 'bytes'):
        factor = 1
    else:
        try:
            factor = UNITS[parts['unit']]
        except KeyError:
            raise ValueError('Invalid size : {}'.format(size))
    return int(float(parts['value']) * factor)


def pvcreate(*pvs):
    run_cmd_logged('pvcreate', *pvs)


def pvscan(dev=None):
    cmd = ['pvscan']
    if dev:
        cmd += ('--cache', dev)
    run_cmd_logged(*cmd)


def run_cmd(*args, **kwargs):
    cmd = list(args)
    log.debug('EXEC {}'.format(' '.join(cmd)))
    with open(os.devnull) as nul:
        p = subprocess.Popen(cmd, stdin=nul, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE, **kwargs)
        out, err = p.communicate()
    if p.returncode != 0:
        msg = err.strip().decode()
        if not msg:
            msg = 'Command exited with status {}'.format(p.returncode)
        log.error(msg)
        raise CommandError(msg)
    else:
        return out.strip().decode()


def run_cmd_logged(*args, **kwargs):
    for line in run_cmd(*args, **kwargs).splitlines():
        log.info(line)


def vgcreate(name, *pvs):
    run_cmd_logged('vgcreate', name, *pvs)



def setup_logging(verbose=0):
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    stderr_handler = logging.StreamHandler()
    stderr_handler.setLevel(logging.ERROR)
    logger.addHandler(stderr_handler)
    stdout_handler = logging.StreamHandler(sys.stdout)
    stdout_handler.addFilter(MaxLevelFilter())
    logger.addHandler(stdout_handler)
    if verbose < 1:
        stdout_handler.setLevel(logging.WARNING)
    elif verbose < 2:
        stdout_handler.setLevel(logging.INFO)
    else:
        stdout_handler.setLevel(logging.DEBUG)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--verbose', '-v',
        action='count',
        default=0,
        help='Print additional status information',
    )
    parser.add_argument(
        '--size', '-s',
        type=parse_size,
        default='3G',
        help='Disk image size',
    )
    parser.add_argument(
        '--lvm', '-L',
        metavar='VGNAME',
        nargs='?',
        help='Create volumes using LVM, in the volume group VGNAME. If a '
             'volume group name is not specified, a unique name will be '
             'generated.'
    )
    parser.add_argument(
        '--no-lvm',
        dest='lvm',
        action='store_false',
        help='Do not use LVM',
    )
    parser.add_argument(
        '--volume', '-V',
        metavar='VOL',
        dest='volumes',
        action='append',
        help='Define a new volume, using the format '
             '<mountpoint>[=<size>[;<fstype>[;<name>[;<fsopts>'
             '[;<parttype>]]]]]',
    )
    parser.add_argument(
        '--default-fstype',
        metavar='FSTYPE',
        help='Default filesystem type to use for volumes that do not specify '
             'one of their own',
    )
    parser.add_argument(
        '--no-tmpfs-tmp',
        dest='tmpfstmp',
        action='store_false',
        default=True,
        help='Do not write an fstab entry for /tmp on tmpfs',
    )
    parser.add_argument(
        '--overlay', '-O',
        metavar='FILENAME',
        dest='overlays',
        action='append',
        default=[],
        help='Extract the contents of FILENAME onto the image',
    )
    parser.add_argument(
        '--script', '-S',
        metavar='FILENAME',
        dest='scripts',
        action='append',
        default=[],
        help='Run a script inside the image',
    )
    parser.add_argument(
        '--no-chroot',
        dest='chroot',
        action='store_false',
        default=True,
        help='Do not chroot into the image to run scripts - DANGEROUS!',
    )
    parser.add_argument(
        '--inject-ssh-key', '-i',
        metavar='FILENAME',
        nargs='?',
        dest='inject_ssh_keys',
        action='append',
        help='Pre-authorize SSH public keys for root',
    )
    parser.add_argument(
        'stagetbz',
        help='Path to stage tarball',
    )
    parser.add_argument(
        'image',
        nargs='?',
        help='Path to destination image file',
    )
    args = parser.parse_args()
    if args.image is None:
        basename = os.path.splitext(args.stagetbz)[0]
        if basename.endswith('.tar'):
            basename = basename[:-4]
        args.image = basename + '.img'
    if args.volumes is None:
        args.volumes = ['/']
    return args


def main():
    args = parse_args()

    setup_logging(args.verbose)

    if args.lvm is None:
        args.lvm = gen_vgname()

    volumes = parse_volumes(args.volumes)
    if args.lvm:
        try:
            partitions = [volumes.pop('/boot'), LVM_PV_PART]
        except KeyError:
            partitions = None
    else:
        partitions = volumes.values()

    try:
        image = Image.create(args.image)
    except OSError as e:
        sys.stderr.write('Failed to create image: {}\n'.format(e))
        raise SystemExit(os.EX_CANTCREAT)
    if args.default_fstype:
        image.default_fstype = args.default_fstype
    with image:
        image.resize(args.size)
        if partitions:
            image.partition(partitions)
        if args.lvm:
            image.setup_lvm(args.lvm, volumes.values())
        image.make_filesystems()
        image.extract(args.stagetbz)
        for overlay in args.overlays:
            image.extract(overlay)
        image.write_fstab(tmpfstmp=args.tmpfstmp)
        for script in args.scripts:
            image.run_script(script, chroot=args.chroot)
        if args.inject_ssh_keys:
            inject_ssh_keys(image.mountpoint, args.inject_ssh_keys)


if __name__ == '__main__':
    main()