#!/usr/bin/python3

# Copyright (C) 2009 Canonical Ltd.

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3,
# as published by the Free Software Foundation.
#
# 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 dbus
from gi.repository import GObject, GLib, UDisks
import dbus.service
import logging
import os
logging.basicConfig(level=logging.DEBUG)

from dbus.mainloop.glib import DBusGMainLoop
from usbcreator.misc import (
    USBCreatorProcessException,
    find_on_path,
    popen,
    sane_path,
    )

USBCREATOR_IFACE = 'com.ubuntu.USBCreator'
PROPS_IFACE = 'org.freedesktop.DBus.Properties'

no_options = GLib.Variant('a{sv}', {})

sane_path()

def _get_object_path_from_device(device_name):
    if device_name.startswith('/dev/'):
        return '/org/freedesktop/UDisks2/block_devices/' + device_name[5:]
    return device_name

def _get_parent_object(udisks, device_name):
    obj = udisks.get_object(_get_object_path_from_device(device_name))
    if obj.get_partition_table():
        return obj
    partition = obj.get_partition()
    parent = partition.get_cached_property('Table').get_string()    
    return udisks.get_object(parent)
    
def unmount_all(udisks, parent):
    '''Unmounts the device or any partitions of the device.'''
    parent_path = parent.get_object_path()
    manager = udisks.get_object_manager()
    for obj in manager.get_objects():
        block = obj.get_block()
        partition = obj.get_partition()
        fs = obj.get_filesystem()
        if not (block and partition and fs):
            continue
        block_name = block.get_cached_property('Device').get_bytestring().decode('utf-8')
        table = partition.get_cached_property('Table').get_string()
        mounts = fs.get_cached_property('MountPoints').get_bytestring_array()
        if table == parent_path and len(mounts):
            logging.debug('Unmounting %s' % block_name)
            # We explictly avoid catching errors here so that failure to
            # unmount a partition causes the format method to fail with the
            # error floating up to the frontend.
            fs.call_unmount_sync(no_options, None)
            
    fs = parent.get_filesystem()
    if not fs:
        return
    dev_name = parent.get_block().get_cached_property('Device').get_bytestring().decode('utf-8')
    mounts = fs.get_cached_property('MountPoints').get_bytestring_array()
    if len(mounts):
        logging.debug('Unmounting %s' % dev_name)
        fs.call_unmount_sync(no_options, None)

def check_system_internal(device):
    block = device.get_block()
    if block.get_cached_property('HintSystem').get_boolean():
        raise dbus.DBusException('com.ubuntu.USBCreator.Error.SystemInternal')

def mem_free():
    # Largely copied from partman-base.
    free = 0
    with open('/proc/meminfo') as meminfo:
        for line in meminfo:
            if line.startswith('MemFree:'):
                free += int(line.split()[1]) / 1024.0
            if line.startswith('Buffers:'):
                free += int(line.split()[1]) / 1024.0
    return free

class USBCreator(dbus.service.Object):
    def __init__(self):
        bus_name = dbus.service.BusName(USBCREATOR_IFACE, bus=dbus.SystemBus())
        dbus.service.Object.__init__(self, bus_name, '/com/ubuntu/USBCreator')
        self.dbus_info = None
        self.polkit = None

    @dbus.service.method(USBCREATOR_IFACE, in_signature='', out_signature='b')
    def KVMOk(self):
        mem = mem_free()
        logging.debug('Asked to run KVM with %f M free' % mem)
        if mem >= 768 and find_on_path('kvm-ok') and find_on_path('kvm'):
            import subprocess
            if subprocess.call(['kvm-ok']) == 0:
                return True
        return False

    @dbus.service.method(USBCREATOR_IFACE, in_signature='sa{ss}', out_signature='')
    def KVMTest(self, device, env):
        '''Run KVM with the freshly created device as the first disk.'''
        for key in ('DISPLAY', 'XAUTHORITY'):
            if key not in env:
                logging.debug('Missing %s' % key)
                return
        udisks = UDisks.Client.new_sync(None)
        obj = _get_parent_object(udisks, device)
        # TODO unmount all the partitions.
        dev_file = obj.get_block().get_cached_property('Device').get_bytestring().decode('utf-8')
        if mem_free() >= 768:
            envp = []
            for k, v in env.items():
                envp.append('%s=%s' % (str(k), str(v)))
            cmd = ('kvm', '-m', '512', '-hda', str(dev_file))
            flags = (GObject.SPAWN_SEARCH_PATH)
            # Don't let SIGINT propagate to the child.
            GObject.spawn_async(cmd, envp=envp, flags=flags, child_setup=os.setsid)

    @dbus.service.method(USBCREATOR_IFACE, in_signature='ss', out_signature='',
                         sender_keyword='sender', connection_keyword='conn')
    def InstallEFI(self, target, efi_image, sender=None, conn=None):
        '''Unpacks bootx64.efi from an image in the proper location that uEFI
           firmware expects it'''
        self.check_polkit(sender, conn, 'com.ubuntu.usbcreator.bootloader')
        mount_point = self.MountISO(efi_image)
        if not os.path.isdir(os.path.join(target, 'efi', 'boot')):
            os.makedirs(os.path.join(target, 'efi', 'boot'))
        #really small file, don't bother with fancy copy methods
        with open(os.path.join(mount_point, 'efi', 'boot', 'bootx64.efi'), 'r') as rd:
            with open(os.path.join(target,  'efi', 'boot', 'bootx64.efi'), 'w') as wd:
                wd.write(rd.read())
        self.UnmountFile(mount_point)


    # TODO return boolean success
    @dbus.service.method(USBCREATOR_IFACE, in_signature='sbsb', out_signature='',
                         sender_keyword='sender', connection_keyword='conn')
    def InstallBootloader(self, device, allow_system_internal, grub_location,
                          syslinux_legacy, sender=None, conn=None):
        '''Install a bootloader to the boot code area, either grub or syslinux.

           The function takes a partition device file of the form /dev/sda1
           and an option grub_location argument for where grub is located

           For GRUB, it's expected that GRUB already exists, all that is
           installed is the bootsector code from grub-setup.

           For syslinux:
           Installs syslinux to the partition boot code area and writes the
           syslinux boot code to the disk code area.  The latter is done to
           handle cases where a bootloader is accidentally installed to the
           MBR, and to handle some buggy BIOSes.'''

        self.check_polkit(sender, conn, 'com.ubuntu.usbcreator.bootloader')
        udisks = UDisks.Client.new_sync(None)
        obj = udisks.get_object(_get_object_path_from_device(device))
        logging.debug('WTF?!: %s' % device)
        if not allow_system_internal:
            check_system_internal(obj)

        parent_obj = _get_parent_object(udisks, device)
        parent_file = parent_obj.get_block().get_cached_property('Device').get_bytestring().decode('utf-8')

        if grub_location:
            #Expect that boot.img and core.img both pre-built in grub_location
            popen(['dd', 'if=%s' % os.path.join(grub_location, 'boot.img'), 'of=%s' % parent_file,
                   'bs=446', 'count=1', 'conv=sync'])
            popen(['dd', 'if=%s' % os.path.join(grub_location, 'core.img'), 'of=%s' % parent_file,
                   'bs=512', 'count=62', 'seek=1', 'conv=sync'])
        else:
            if syslinux_legacy and find_on_path('syslinux-legacy'):
                syslinux = 'syslinux-legacy'
            else:
                syslinux = 'syslinux'
            popen([syslinux, '-f', device])
            # Write the syslinux MBR.
            popen(['dd', 'if=/usr/lib/%s/mbr.bin' % syslinux, 'of=%s' % parent_file,
                   'bs=446', 'count=1', 'conv=sync'])

        part = obj.get_partition()
        if part:
            number = part.get_cached_property('Number').get_uint32()
        else:
            return
        try:
            popen(['/sbin/parted', parent_file, 'set', str(number), 'boot', 'on'])
        except USBCreatorProcessException:
            # Don't worry about not being able to re-read the partition table.
            # TODO: As this will still be a problem for KVM users, this should
            # be fixed by unmounting all the partitions before we get to this
            # point, then remounting the target partition after.
            pass

    @dbus.service.method(USBCREATOR_IFACE, in_signature='sb', out_signature='',
                         sender_keyword='sender', connection_keyword='conn')
    def Format(self, device, allow_system_internal, sender=None, conn=None):
        self.check_polkit(sender, conn, 'com.ubuntu.usbcreator.format')

        # TODO evand 2009-08-25: Needs a confirmation dialog.
        # XXX test with a device that doesn't have a partition table.
        udisks = UDisks.Client.new_sync(None)
        dev = udisks.get_object(device)
        parent_dev = _get_parent_object(udisks, device)
        
        if not allow_system_internal:
            check_system_internal(dev)

        # TODO LOCK
        unmount_all(udisks, dev)
        # Do NOT use the disk if asked to format a partition.
        # We still need to obtain the disk device name to zero out the MBR
        part = dev.get_partition()
        block = dev.get_block()
        if part:
            # Create the partition
            #type, label, flags
            boot_dos_flag=64
            part.call_set_type_sync('0x0c', no_options, None)
            part.call_set_flags_sync(boot_dos_flag, no_options, None)
            block.call_format_sync('vfat', GLib.Variant('a{sv}', {'label': GLib.Variant('s', '')}), None)
        else:
            # Create a new partition table and a FAT partition.
            # Is this correct call to create partition table ?!
            block.call_format_sync('dos', GLib.Variant('a{sv}', {'erase': GLib.Variant('s', '')}), None)
            size = block.get_cached_property('Size').get_uint64()
            table = dev.get_partition_table()
            partition = table.call_create_partition_sync(0, size, '0x0c', '', no_options, None) 
            obj = udisks.get_object(partition)
            block = obj.get_block()
            block.call_format_sync('vfat', GLib.Variant('a{sv}', {'label': GLib.Variant('s', '')}), None)

        # Zero out the MBR.  Will require fancy privileges.
        parent_block = parent_dev.get_block()
        parent_file = parent_block.get_cached_property('Device').get_bytestring().decode('utf-8')
        popen(['dd', 'if=/dev/zero', 'of=%s' % parent_file, 'bs=446', 'count=1'])
        # TODO UNLOCK

    @dbus.service.method(USBCREATOR_IFACE, in_signature='ssb', out_signature='',
                         sender_keyword='sender', connection_keyword='conn')
    def Image(self, source, target, allow_system_internal,
              sender=None, conn=None):
        self.check_polkit(sender, conn, 'com.ubuntu.usbcreator.image')
        if not allow_system_internal:
            check_system_internal(target)
        cmd = ['dd', 'if=%s' % str(source), 'of=%s' % str(target), 'bs=1M']
        popen(cmd)

    @dbus.service.method(USBCREATOR_IFACE, in_signature='s', out_signature='s',
                         sender_keyword='sender', connection_keyword='conn')
    def MountISO(self, device, sender=None, conn=None):
        self.check_polkit(sender, conn, 'com.ubuntu.usbcreator.mount')
        import tempfile
        ret = tempfile.mkdtemp()
        device = device.encode('utf-8')
        popen(['mount', '-r', '-o', 'loop', device, ret])
        return ret

    @dbus.service.method(USBCREATOR_IFACE, in_signature='s', out_signature='',
                         sender_keyword='sender', connection_keyword='conn')
    def UnmountFile(self, device, sender=None, conn=None):
        self.check_polkit(sender, conn, 'com.ubuntu.usbcreator.mount')
        popen(['umount', device])

    @dbus.service.method(USBCREATOR_IFACE, in_signature='s', out_signature='',
                         sender_keyword='sender', connection_keyword='conn')
    def Unmount(self, device, sender=None, conn=None):
        self.check_polkit(sender, conn, 'com.ubuntu.usbcreator.mount')
        udisks = UDisks.Client.new_sync(None)
        parent = udisks.get_object(device)
        unmount_all(udisks, parent)

    @dbus.service.method(USBCREATOR_IFACE, in_signature='s', out_signature='',
                         sender_keyword='sender', connection_keyword='conn')
    def RemountRW(self, device, sender=None, conn=None):
        # Until udisks supports remounting devices.
        self.check_polkit(sender, conn, 'com.ubuntu.usbcreator.mount')
        popen(['mount', '-o', 'remount,rw', device])

    @dbus.service.method(USBCREATOR_IFACE, in_signature='', out_signature='',
                         sender_keyword='sender', connection_keyword='conn')
    def Shutdown(self, sender=None, conn=None):
        logging.debug('Shutting down.')
        loop.quit()

    # Taken from Jockey 0.5.3.
    def check_polkit(self, sender, conn, priv):
        if sender is None and conn is None:
            return
        if self.dbus_info is None:
            self.dbus_info = dbus.Interface(conn.get_object(
                                            'org.freedesktop.DBus',
                                            '/org/freedesktop/DBus/Bus',
                                            False), 'org.freedesktop.DBus')
        pid = self.dbus_info.GetConnectionUnixProcessID(sender)
        if self.polkit is None:
            self.polkit = dbus.Interface(dbus.SystemBus().get_object(
                                'org.freedesktop.PolicyKit1',
                                '/org/freedesktop/PolicyKit1/Authority',
                                False), 'org.freedesktop.PolicyKit1.Authority')
        try:
            # we don't need is_challenge return here, since we call with
            # AllowUserInteraction
            (is_auth, _, details) = self.polkit.CheckAuthorization(
                                    ('system-bus-name', {'name': dbus.String(sender,
                                        variant_level = 1)}), priv, {'': ''},
                                    dbus.UInt32(1), '', timeout=600)
        except dbus.DBusException as e:
            if e._dbus_error_name == 'org.freedesktop.DBus.Error.ServiceUnknown':
                # polkitd timed out, connect again
                self.polkit = None
                return self.check_polkit(sender, conn, priv)
            else:
                raise

        if not is_auth:
            logging.debug('_check_polkit_privilege: sender %s on connection %s '
                          'pid %i is not authorized for %s: %s' %
                          (sender, conn, pid, priv, str(details)))
            raise dbus.DBusException('com.ubuntu.USBCreator.Error.NotAuthorized')

DBusGMainLoop(set_as_default=True)
helper = USBCreator()
loop = GLib.MainLoop()
loop.run()
