#! /usr/bin/python
# Copyright 2011-2013  Lars Wirzenius
# Copyright 2012  Codethink Limited
# Copyright 2014-2015 Neil Williams <codehelp@debian.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import sys
import time
import shutil
import cliapp
import logging
import datetime
import tempfile
import subprocess
from vmdebootstrap.base import (
    Base,
    runcmd,
    cleanup_apt_cache,
)
from vmdebootstrap.grub import GrubHandler
from vmdebootstrap.extlinux import ExtLinux
from vmdebootstrap.codenames import Codenames
from vmdebootstrap.filesystem import Filesystem
from vmdebootstrap.uefi import Uefi
from vmdebootstrap.constants import arch_table

__version__ = '1.0'

# pylint: disable=invalid-name,line-too-long,missing-docstring


class VmDebootstrap(cliapp.Application):  # pylint: disable=too-many-public-methods

    def __init__(self, progname=None, version=__version__, description=None, epilog=None):
        super(VmDebootstrap, self).__init__(progname, version, description, epilog)
        self.remove_dirs = []
        self.mount_points = []
        self.bootdir = None
        self.handlers = {
            Base.name: Base(),
            Uefi.name: Uefi(),
            Codenames.name: Codenames(),
            GrubHandler.name: GrubHandler(),
            ExtLinux.name: ExtLinux(),
            Filesystem.name: Filesystem(),
        }

    def add_settings(self):
        default_arch = subprocess.check_output(
            ["dpkg", "--print-architecture"]).strip()

        self.settings.boolean(
            ['verbose'], 'report what is going on')
        self.settings.string(
            ['image'], 'put created disk image in FILE',
            metavar='FILE')
        self.settings.bytesize(
            ['size'], 'create a disk image of size SIZE (%default)',
            metavar='SIZE', default='1G')
        self.settings.bytesize(
            ['bootsize'], 'create boot partition of size '
            'SIZE (%default)',
            metavar='BOOTSIZE', default='0%')
        self.settings.string(
            ['boottype'], 'specify file system type for /boot/',
            default='ext2')
        self.settings.bytesize(
            ['bootoffset'], 'Space to leave at start of the '
            'image for bootloader',
            default='0')
        self.settings.boolean(
            ['use-uefi'], 'Setup image for UEFI boot',
            default=False)
        self.settings.bytesize(
            ['esp-size'], 'Size of EFI System Partition - '
            'requires use-uefi',
            default='5mib')
        self.settings.string(
            ['part-type'], 'Partition type to use for this image',
            default='msdos')
        self.settings.string(
            ['roottype'], 'specify file system type for /',
            default='ext4')
        self.settings.bytesize(
            ['swap'], 'create swap space of size SIZE (min 256Mb)')
        self.settings.string(
            ['foreign'], 'set up foreign debootstrap environment '
            'using provided program (ie binfmt handler)')
        self.settings.string_list(
            ['debootstrapopts'],
            'pass additional options to debootstrap'),
        self.settings.boolean(
            ['extlinux'], 'install extlinux?',
            default=True)
        self.settings.string(
            ['tarball'], "tar up the disk's contents in FILE",
            metavar='FILE')
        self.settings.string(
            ['apt-mirror'], 'configure apt to use MIRROR',
            metavar='URL')
        self.settings.string(
            ['mirror'], 'use MIRROR as package source (%default)',
            metavar='URL',
            default='http://http.debian.net/debian/')
        self.settings.string(
            ['arch'], 'architecture to use (%default)',
            metavar='ARCH',
            default=default_arch)
        self.settings.string(
            ['distribution'], 'release to use (%default)',
            metavar='NAME',
            default='stable')
        self.settings.string_list(
            ['package'], 'install PACKAGE onto system')
        self.settings.string_list(
            ['custom-package'], 'install package in DEB file '
            'onto system (not from mirror)',
            metavar='DEB')
        self.settings.boolean(
            ['no-kernel'], 'do not install a linux package')
        self.settings.string(
            ['kernel-package'], 'install PACKAGE instead of '
            'the default kernel package',
            metavar='PACKAGE')
        self.settings.boolean(
            ['enable-dhcp'], 'enable DHCP on eth0')
        self.settings.string(
            ['root-password'], 'set root password',
            metavar='PASSWORD')
        self.settings.boolean(
            ['lock-root-password'], 'lock root account so they '
            'cannot login?')
        self.settings.string(
            ['customize'], 'run SCRIPT after setting up system',
            metavar='SCRIPT')
        self.settings.string(
            ['hostname'], 'set name to HOSTNAME (%default)',
            metavar='HOSTNAME',
            default='debian')
        self.settings.string_list(
            ['user'], 'create USER with PASSWORD',
            metavar='USER/PASSWORD')
        self.settings.boolean(
            ['serial-console'], 'configure image to use a serial console')
        self.settings.string(
            ['serial-console-command'], 'command to manage the '
            'serial console, appended to /etc/inittab (%default)',
            metavar='COMMAND',
            default='/sbin/getty -L ttyS0 115200 vt100')
        self.settings.boolean(
            ['sudo'], 'install sudo, and if user is created, '
            'add them to sudo group')
        self.settings.string(
            ['owner'], 'the user who will own the image when '
            'the build is complete.')
        self.settings.boolean(
            ['squash'], 'use squashfs on the final image.')
        self.settings.string(
            ['squash-file'], 'filename for the squashfs '
            '- cannot be used with --image',
            metavar='FILE',
            default='rootfs.squash')
        self.settings.boolean(
            ['configure-apt'], 'Create an apt source based on '
            'the distribution and mirror selected.')
        self.settings.boolean(
            ['mbr'], 'Run install-mbr (default if extlinux used)')
        self.settings.boolean(
            ['grub'], 'Install and configure grub2 - '
            'disables extlinux.')
        self.settings.boolean(
            ['sparse'], 'Do not fill the image with zeros to '
            'keep a sparse disk image',
            default=False)
        self.settings.boolean(
            ['pkglist'], 'Create a list of package names '
            'included in the image.')
        self.settings.boolean(
            ['no-acpid'], 'do not install the acpid package',
            default=False)

    def process_args(self, args):  # pylint: disable=too-many-branches,too-many-statements

        for _, handler in self.handlers.items():
            handler.define_settings(self.settings)

        distro = self.handlers[Codenames.name]
        if not self.settings['image'] and not (
                self.settings['tarball'] or self.settings['squash']):
            raise cliapp.AppException(
                'You must give disk image filename or use either a '
                'tarball filename or use squash')
        if self.settings['image'] and not self.settings['size']:
            raise cliapp.AppException(
                'If disk image is specified, you must give image size.')
        if not distro.debian_info.valid(self.settings['distribution']):
            if not distro.ubuntu_info.valid(self.settings['distribution']):
                raise cliapp.AppException(
                    '%s is not a valid Debian or Ubuntu suite or codename.'
                    % self.settings['distribution'])
        uefi = self.handlers[Uefi.name]
        uefi.check_settings()
        if os.geteuid() != 0:
            sys.exit("You need to have root privileges to run this script.")
        self.start_ops()

    def _image_preparations(self):
        uefi = self.handlers[Uefi.name]
        base = self.handlers[Base.name]
        filesystem = self.handlers[Filesystem.name]
        extlinux = self.handlers[ExtLinux.name]
        base.create_empty_image()
        self.partition_image()
        extlinux.install_mbr()
        filesystem.setup_kpartx()
        rootdev = filesystem.devices['rootdev']
        roottype = filesystem.devices['roottype']
        bootdev = filesystem.devices['bootdev']
        if self.settings['swap'] > 0:
            base.message("Creating swap space")
            runcmd(['mkswap', filesystem.devices['swapdev']])
        filesystem.mkfs(rootdev, fstype=roottype)
        rootdir = self.mount(rootdev)
        filesystem.devices['rootdir'] = rootdir
        if self.settings['use-uefi']:
            self.bootdir = uefi.prepare_esp(rootdir, bootdev)
            logging.debug("mounting %s", self.bootdir)
            self.mount(bootdev, self.bootdir)
            logging.debug(runcmd(['mount']))
        elif bootdev:
            boottype = self.settings['boottype']
            filesystem.mkfs(bootdev, fstype=boottype)
            self.bootdir = '%s/%s' % (rootdir, 'boot/')
            filesystem.devices['bootdir'] = self.bootdir
            os.mkdir(self.bootdir)
            self.mount(bootdev, self.bootdir)

    def _image_operations(self, rootdir, rootdev):
        if not self.settings['image']:
            return
        logging.debug("rootdir=%s rootdev=%s", rootdir, rootdev)
        grub = self.handlers[GrubHandler.name]
        extlinux = self.handlers[ExtLinux.name]
        base = self.handlers[Base.name]
        uefi = self.handlers[Uefi.name]
        filesystem = self.handlers[Filesystem.name]
        if self.settings['use-uefi']:
            bootdir = self.bootdir
            logging.debug(
                "rootdir=%s rootdev=%s bootdir=%s",
                rootdir, rootdev, bootdir)
            logging.debug(runcmd(['mount']))
            if not os.path.ismount(bootdir):
                logging.warning("%s had to be remounted", bootdir)
                self.mount(bootdir)
            grub.install_grub_uefi(rootdir)
            uefi.configure_efi(rootdir)
            grub.install_extra_grub_uefi(rootdir)
            uefi.configure_extra_efi(rootdir)
        elif self.settings['grub']:
            if not grub.install_grub2(rootdev, rootdir):
                extlinux.install_extlinux(rootdev, rootdir)
        elif self.settings['extlinux']:
            extlinux.install_extlinux(rootdev, rootdir)
        base.append_serial_console(rootdir)
        self.optimize_image(rootdir)
        filesystem.squash_image()

    def start_ops(self):
        base = self.handlers[Base.name]
        filesystem = self.handlers[Filesystem.name]
        try:
            if self.settings['image']:
                self._image_preparations()
                rootdir = filesystem.devices['rootdir']
                rootdev = filesystem.devices['rootdev']
            else:
                rootdir = self.mkdtemp()
                rootdev = filesystem.devices['rootdev']
            logging.debug("rootdir=%s rootdev=%s", rootdir, rootdev)
            self.debootstrap(rootdir)
            filesystem.set_hostname()
            filesystem.create_fstab()
            self.install_debs(rootdir)
            base.set_root_password(rootdir)
            base.create_users(rootdir)
            filesystem.remove_udev_persistent_rules()
            self.setup_networking(rootdir)
            filesystem.configure_apt()
            base.customize(rootdir)
            cleanup_apt_cache(rootdir)
            filesystem.update_initramfs()
            self._image_operations(rootdir, rootdev)
            filesystem.list_installed_pkgs()

            if self.settings['foreign']:
                os.unlink(
                    '%s/usr/bin/%s' %
                    (rootdir, os.path.basename(self.settings['foreign'])))

            if self.settings['tarball']:
                base.create_tarball(rootdir)
            else:
                filesystem.squash_filesystem()
            filesystem.chown()

        except BaseException as e:
            base.message('EEEK! Something bad happened...')
            rootdir = filesystem.devices['rootdir']
            if rootdir:
                db_log = os.path.join(rootdir, 'debootstrap', 'debootstrap.log')
                if os.path.exists(db_log):
                    shutil.copy(db_log, os.getcwd())
            base.message(e)
            self.cleanup_system()
            raise
        else:
            self.cleanup_system()

    def mkdtemp(self):
        dirname = tempfile.mkdtemp()
        self.remove_dirs.append(dirname)
        logging.debug('mkdir %s', dirname)
        return dirname

    def mount(self, device, path=None):
        base = self.handlers[Base.name]
        if not path:
            mount_point = self.mkdtemp()
        else:
            mount_point = path
        base.message('Mounting %s on %s' % (device, mount_point))
        runcmd(['mount', device, mount_point])
        self.mount_points.append(mount_point)
        logging.debug('mounted %s on %s', device, mount_point)
        return mount_point

    def partition_image(self):
        """
        Uses fat16 (msdos) partitioning by default, use part-type to change.
        If bootoffset is specified, the first actual partition
        starts at that offset to allow customisation scripts to
        put bootloader images into the space, e.g. u-boot.
        """
        base = self.handlers[Base.name]
        base.message('Creating partitions')
        uefi = self.handlers[Uefi.name]
        runcmd(['parted', '-s', self.settings['image'],
                'mklabel', self.settings['part-type']])
        partoffset = 0
        extent = base.check_swap_size()

        # uefi
        uefi.partition_esp()

        # /boot partitioning offset calculation
        # returns partoffset
        if self.settings['bootoffset'] and self.settings['bootoffset'] is not '0':
            # turn v.small offsets into something at least possible to create.
            if self.settings['bootoffset'] < 1048576:
                partoffset = 1
                logging.info(
                    "Setting bootoffset %smib to allow for %s bytes",
                    partoffset, self.settings['bootoffset'])
            else:
                partoffset = self.settings['bootoffset'] / (1024 * 1024)
                base.message(
                    "Using bootoffset: %smib %s bytes" %
                    (partoffset, self.settings['bootoffset']))

        # /boot creation - move into base but keep the check
        # needs extent, partoffset, bootsize: no return
        if self.settings['bootsize'] and self.settings['bootsize'] is not '0%':
            if self.settings['grub'] and not partoffset:
                partoffset = 1
            bootsize = self.settings['bootsize'] / (1024 * 1024)
            bootsize += partoffset
            base.message("Using bootsize %smib: %s bytes" % (bootsize, self.settings['bootsize']))
            logging.debug("Starting boot partition at %sMb", bootsize)
            runcmd(['parted', '-s', self.settings['image'],
                    'mkpart', 'primary', 'fat16', str(partoffset), str(bootsize)])
            logging.debug("Starting root partition at %sMb", partoffset)
            runcmd(['parted', '-s', self.settings['image'],
                    'mkpart', 'primary', str(bootsize), extent])

        # uefi - make rootfs partition after end of ESP
        # needs extent
        elif self.settings['use-uefi']:
            uefi.make_root(extent)

        # no boot partition
        else:
            runcmd(['parted', '-s', self.settings['image'],
                    'mkpart', 'primary', '0%', extent])

        # whatever we create, something needs the boot flag
        runcmd(['parted', '-s', self.settings['image'],
                'set', '1', 'boot', 'on'])

        # return to doing swap setup
        base.make_swap(extent)

    def _bootstrap_packages(self):
        base = self.handlers[Base.name]
        uefi = self.handlers[Uefi.name]
        grub = self.handlers[GrubHandler.name]
        distro = self.handlers[Codenames.name]

        include = self.settings['package']
        include.extend(base.base_packages())
        include.extend(uefi.efi_packages())
        include.extend(grub.grub_packages())
        include.extend(distro.kernel_package())
        return list(set(include))

    def _debootstrap_second_stage(self, rootdir):
        base = self.handlers[Base.name]
        # set a noninteractive debconf environment for secondstage
        env = {
            "DEBIAN_FRONTEND": "noninteractive",
            "DEBCONF_NONINTERACTIVE_SEEN": "true",
            "LC_ALL": "C"
        }
        # add the mapping to the complete environment.
        env.update(os.environ)
        # First copy the binfmt handler over
        base.message('Setting up binfmt handler')
        shutil.copy(self.settings['foreign'], '%s/usr/bin/' % rootdir)
        # Next, run the package install scripts etc.
        base.message('Running debootstrap second stage')
        runcmd(['chroot', rootdir,
                '/debootstrap/debootstrap', '--second-stage'],
               env=env)

    def debootstrap(self, rootdir):
        base = self.handlers[Base.name]
        include = self._bootstrap_packages()
        base.message(
            'Debootstrapping %s [%s]' % (
                self.settings['distribution'], self.settings['arch']))

        args = ['debootstrap', '--arch=%s' % self.settings['arch']]

        if self.settings['package']:
            args.append(
                '--include=%s' % ','.join(include))
        if self.settings['foreign']:
            args.append('--foreign')
        if self.settings['debootstrapopts']:
            for opt in self.settings['debootstrapopts']:
                for part in opt.split(' '):
                    args.append('--%s' % part)
        self.logging.debug("debootstrap arguments: %s", args)
        args += [self.settings['distribution'],
                 rootdir, self.settings['mirror']]
        logging.debug(" ".join(args))
        runcmd(args)
        if self.settings['foreign']:
            self._debootstrap_second_stage(rootdir)

    def install_debs(self, rootdir):
        base = self.handlers[Base.name]
        if not self.settings['custom-package']:
            return
        base.message('Installing custom packages')
        tmp = os.path.join(rootdir, 'tmp', 'install_debs')
        os.mkdir(tmp)
        for deb in self.settings['custom-package']:
            shutil.copy(deb, tmp)
        filenames = [os.path.join('/tmp/install_debs', os.path.basename(deb))
                     for deb in self.settings['custom-package']]
        out, err, _ = \
            self.runcmd_unchecked(['chroot', rootdir, 'dpkg', '-i'] + filenames)
        logging.debug('stdout:\n%s', out)
        logging.debug('stderr:\n%s', err)
        out = runcmd(['chroot', rootdir,
                      'apt-get', '-f', '--no-remove', 'install'])
        logging.debug('stdout:\n%s', out)
        shutil.rmtree(tmp)

    def optimize_image(self, rootdir):
        """
        Filing up the image with zeros will increase its compression rate
        """
        if not self.settings['sparse']:
            zeros = os.path.join(rootdir, 'ZEROS')
            self.runcmd_unchecked(['dd', 'if=/dev/zero', 'of=' + zeros, 'bs=1M'])
            runcmd(['rm', '-f', zeros])

    def setup_networking(self, rootdir):
        base = self.handlers[Base.name]
        base.message('Setting up networking')
        distro = self.handlers[Codenames.name]
        ifc_file = os.path.join(rootdir, 'etc', 'network', 'interfaces')
        ifc_d = os.path.join(rootdir, 'etc', 'network', 'interfaces.d')

        # unconditionally write for wheezy (which became oldstable on 2015.04.25)
        if distro.was_oldstable(datetime.date(2015, 4, 26)):
            with open(ifc_file, 'w') as netfile:
                netfile.write('source /etc/network/interfaces.d/*\n')
        elif not os.path.exists(ifc_file):
            with open(ifc_file, 'a') as netfile:
                netfile.write('source-directory /etc/network/interfaces.d\n')

        if not os.path.exists(ifc_d):
            os.mkdir(ifc_d)
        ethpath = os.path.join(ifc_d, 'setup')
        with open(ethpath, 'w') as eth:
            eth.write('auto lo\n')
            eth.write('iface lo inet loopback\n')

            if self.settings['enable-dhcp']:
                eth.write('\n')
                eth.write('auto eth0\n')
                eth.write('iface eth0 inet dhcp\n')

    def cleanup_system(self):
        base = self.handlers[Base.name]
        # Clean up after any errors.

        base.message('Cleaning up')

        # Umount in the reverse mount order
        if self.settings['image']:
            for i in range(len(self.mount_points) - 1, -1, -1):
                mount_point = self.mount_points[i]
                try:
                    runcmd(['umount', mount_point], ignore_fail=False)
                except cliapp.AppException:
                    logging.debug("umount failed, sleeping and trying again")
                    time.sleep(5)
                    runcmd(['umount', mount_point], ignore_fail=False)

            runcmd(['kpartx', '-d', self.settings['image']], ignore_fail=True)

        for dirname in self.remove_dirs:
            shutil.rmtree(dirname)

if __name__ == '__main__':
    VmDebootstrap(version=__version__).run()