# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#!/usr/bin/env python3
# -*- encoding=utf8 -*-
#
# Author 2011 Hsin-Yi Chen
#
# This is a 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 2 of the License, or (at your
# option) any later version.
#
# This software 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 software; if not, write to the Free Software Foundation, Inc., 59 Temple
# Place, Suite 330, Boston, MA 02111-1307 USA
import dbus
import os
import distro
import shutil
import subprocess

from ubunturecovery import recovery_common as magic

def create_g2ldr(chroot, rp_mount, install_mount):
    '''Create a g2ldr compatible image using the install
       chroot: chroot to launch commands in
       rp_mount: mountpoint to find partition# and install g2ldr
       install_mount: mountpoint to find target OS UUID'''
    bus = dbus.SystemBus()
    udisk_bus_name = 'org.freedesktop.UDisks'
    dev_bus_name   = 'org.freedesktop.UDisks.Device'
    uuid = ''
    partition = '-1'

    #Find the UUID we are installing to
    obj = bus.get_object(udisk_bus_name, '/org/freedesktop/UDisks')
    iface = dbus.Interface(obj, udisk_bus_name)
    devices = iface.EnumerateDevices()
    for device in devices:
        obj = bus.get_object(udisk_bus_name, device)
        dev = dbus.Interface(obj, 'org.freedesktop.DBus.Properties')
        mount = dev.Get(dev_bus_name, 'DeviceMountPaths')
        if mount and install_mount and install_mount in mount:
            uuid = dev.Get(dev_bus_name, 'IdUuid')
        elif mount and rp_mount and rp_mount in mount:
            partition = str(dev.Get(dev_bus_name, 'DeviceMinor'))

    #The file that the Windows BCD will chainload
    shutil.copy('/usr/lib/grub/i386-pc/g2ldr.mbr', rp_mount)

    # This BCD configuration will load from the recovery partition if it exists
    #   a /grub/grub.cfg.  As specified by the active partition flag
    # If it doesn't, then it will look for /boot/grub/grub.cfg on the OS partition
    #   as specified by the UUID.
    cfg = GrubCfg.create(os.path.join(chroot, 'tmp', 'bcd.cfg'),
                          '/usr/share/ubuntu/grub/bcd.cfg')
    cfg.set(uuid=uuid, rp_number=partition)
    cfg.save()

    #Build the Grub2 Core Image
    build_command = ['grub-mkimage', '-O', 'i386-pc',
                                     '-c', '/tmp/bcd.cfg',
                                     '-o', '/tmp/core.img',
                                     'biosdisk', 'part_msdos', 'part_gpt', 'ls',
                                     'fat', 'ntfs', 'ntfscomp', 'search',
                                     'linux', 'vbe', 'boot', 'minicmd',
                                     'cat', 'cpuid', 'chain', 'halt', 'linux16',
                                     'echo', 'test', 'configfile', 'ext2',
                                     'keystatus', 'help', 'boot', 'loadenv' ]
    if chroot == '/target':
        from ubiquity import install_misc
        install_misc.chrex(chroot, *build_command)
    else:
        from ubiquity import misc
        misc.execute_root(*build_command)

    with open(os.path.join(rp_mount, 'g2ldr'), 'wb') as wfd:
        files =  ['/usr/lib/grub/i386-pc/g2hdr.bin',
                  os.path.join(chroot, 'tmp', 'core.img')]
        for fname in files:
            with open(fname, 'rb') as rfd:
                wfd.write(rfd.read())

class GrubCfg(object):

    def __init__(self, filepath):
        self.path = filepath
        self.__str__ = self.__repr__
        dist, ver, codename = distro.linux_distribution(full_distribution_name=False)
        self._kwargs = {
            'os': "{} {}".format(dist, ver)
        }
        self._tpl = ''
        self._header = []

    @classmethod
    def create(cls, filepath, tplpath):
        obj = cls(filepath)
        obj.load_tpl(tplpath)
        return obj

    def load_tpl(self, path):
        with open(path, 'r') as f:
            self._tpl = f.read().split('\n')

    def set(self, **kwargs):
        self._kwargs.update(kwargs)

    def addheaders(self, lines):
        self._header += lines

    def __repr__(self):
        return self.genstr()

    def genstr(self):
        return self._header and '\n'.join(self._header) + '\n' + self.genbody() or self.genbody()

    def genbody(self):
        dist = distro.linux_distribution(full_distribution_name=False)

        #starting with 10.10, we replace the whole drive string (/dev/sdX,msdosY)
        #earlier releases are hardcoded to (hd0,Y)
        rp_number = self._kwargs.get('rp_number', '2')
        if float(dist[1]) >= 10.10:
            platinfo = magic.PlatInfo()
            rp_number = platinfo.disk_layout + rp_number

        ret = []
        for line in self._tpl:
            if "#RECOVERY_TEXT#" in line:
                line = line.replace("#RECOVERY_TEXT#", self._kwargs.get('recovery_text'))
            if "#UUID#" in line:
                line = line.replace("#UUID#", self._kwargs.get('uuid'))
            if "#PARTITION#" in line:
                line = line.replace("#PARTITION#", self._kwargs.get('rp_number'))
            if "#OS#" in line:
                ostext = self._kwargs.get('os', "%s %s" % (dist[0], dist[1]))
                line = line.replace("#OS#", ostext)
            if "#EXTRA#" in line:
                line = line.replace("#EXTRA#", "%s" % self.extra_cmdline())
            ret.append(line)
        return '\n'.join(ret)

    def extra_cmdline(self):
        extra_cmdline = self._kwargs.get('ako')
        if extra_cmdline:
            #remove any duplicate entries
            ka_list = magic.find_extra_kernel_options().split(' ')
            ako_list = extra_cmdline.split(' ')
            for var in ka_list:
                found = False
                for item in ako_list:
                    left = item.split('=')[0].strip()
                    if left and left in var:
                        found = True
                #propagate anything but BOOT_IMAGE (it gets added from isolinux)
                if not found and not 'BOOT_IMAGE' in var:
                    extra_cmdline += ' ' + var
        else:
            extra_cmdline = magic.find_extra_kernel_options()
        return extra_cmdline

    def save(self):
        with open(self.path, 'w') as fwd:
            fwd.write(self.genstr())

class GrubInstaller(object):
    """Object for handling grub configures and grub installing

    Keyword Arguments:
        - targetdir --  A directory that gub configures will be copied to
        - tpldir -- A directory that grub configure templates located
        - efi -- running as EFI mode? (default: False)
    """

    def __init__(self, targetdir, tpldir, efi=False):
        self.targetdir = targetdir
        self.targetgrubdir = os.path.join(self.targetdir, 'boot', 'grub')
        self.tpldir = tpldir
        self.efi = efi
        self.conf = None

        from ubiquity import misc
        with misc.raised_privileges():
            subprocess.call(['mkdir', '-p', self.targetgrubdir])

    def cfgpath(self, cfgname):
        return os.path.join(self.targetgrubdir, cfgname)

    def tplpath(self, tplname):
        return os.path.join(self.tpldir, tplname)

    def create_cfg(self, filename, tplname, **kwargs):
        """create grub configure file

        @param filename grub configure file name
        @param tplname template name
        @param **kwargs
        """
        cfg = GrubCfg.create(self.cfgpath(filename),
                             self.tplpath(tplname))
        cfg.set(**kwargs)
        return cfg

    def install_rp_grubcfg(self, files, data, dual=False, force=False):
        """installing grub configures for a recovery partition

        @param files grub configure name
        @param force overwrite if this is True
        @param **kwargs datas for filling template
        """
        for tplname, cfgname in files.items():
            cfgpath = self.cfgpath(cfgname)
            if os.path.exists(self.cfgpath(cfgname)):
                if force:
                    shutil.move(cfgpath, cfgpath + '.old')
                else:
                    # file exists, skip it
                    continue
            # create new grub configure
            cfg = self.create_cfg(cfgname, tplname, **data)
            # setting for support recovery hotkey if needed
            if self.conf and \
               self.conf.rec_hotkey_trigger == 'mbr' and \
               cfgname == 'common.cfg':
                _headers = [
                    'parttool (hd0,{}) boot-'.format(data.get('rp_number')),
                    'parttool (hd0,{}) type=0x{}'.format(data.get('rp_number'), self.conf.rec_hotkey_filesystemtype),
                    'parttool (hd0,{}) boot+'.format(self.conf.os_part)]
                cfg.addheaders(_headers)
            cfg.save()
            #Allow these to be invoked from a recovery solution launched
            #by the BCD.
            if dual:
                shutil.copy(cfgpath, os.path.join('/tmp', cfgname))

        # Install unicode.pf2 font
        shutil.copy('/usr/share/grub/unicode.pf2', self.targetgrubdir)

        # Add Cedar Trail graphics to the blacklist.
        with open(os.path.join(self.targetgrubdir, 'cedartraillist.txt'), 'w') as wfd:
            wfd.write("v8086d0be.*\n")

        # Put a text efi file
        if self.efi:
            with open(os.path.join(self.targetgrubdir, 'efi'), 'w') as wfd:
                wfd.write("EFI\n")

    def install_rp_grub(self, rp_devicefile, grub_devicefile):
        """installing grub core.img in recovery partiton to
           MBR or specified partition

        @param rp_devicefile recovery partition devicefile
        @param grub_devicefile  partition device that grub core img will be installed
        """
        from ubiquity import misc
        if self.efi:
            efidir = os.path.join('/mnt', 'boot', 'efi')
            with misc.raised_privileges():
                if not os.path.isdir(efidir):
                    os.makedirs(efidir)
            if not misc.execute_root('mount', grub_devicefile, efidir):
                raise RuntimeError("Error mounting %s" % grub_devicefile)
            self.install_grub(rp_devicefile, True)
            misc.execute_root('umount', efidir)
        else:
            self.install_grub(rp_devicefile, True)

    def install_grub(self, install_dev, recovery=False):
        """installing grub to target device

        @param install_dev the device grub will be install
        @param recovery Install grub in recovery partition.
        """
        from ubiquity import misc
        cmds = ['/usr/sbin/grub-install']
        if self.efi:
            cmds.append('--target=x86_64-efi')
        else:
            cmds.append('--target=i386-pc')
        cmds.append('--force')
        if self.targetdir:
            cmds.append('--root-directory=' + self.targetdir)
        if recovery and self.efi:
            cmds.append('--removable')
        if install_dev:
            cmds.append(install_dev)
        if not misc.execute_root(*cmds):
            msg = "Error installing grub"
            if install_dev:
                msg + ' to {}'.format(install_dev)
            raise RuntimeError(msg)

        # Remove EFI files to avoid some issue when installing in UEFI legacy mode. (LP: #1090265)
        if recovery and self.targetdir and not self.efi:
            efidir = os.path.join(self.targetdir, 'efi')
            if os.path.isdir(efidir):
                cmds = ['rm', '-fr', efidir]
                if not misc.execute_root(*cmds):
                    msg = "Error removing EFI files"
                    if install_dev:
                        msg + ' from {}'.format(install_dev)
                    raise RuntimeError(msg)

    def build_g2ldr(self, files):
        """building grub in logic partitoin

        @param files grub configure template and files
        """
        create_g2ldr('/', '/mnt', '')
        if not os.path.isdir(os.path.join('/mnt', 'boot', 'grub')):
            os.makedirs(os.path.join('/mnt', 'boot', 'grub'))
        for item in files:
            shutil.copy(os.path.join('/tmp', files[item]), \
                        os.path.join('/mnt', 'boot', 'grub', files[item]))
