# 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 gi
gi.require_version('UDisks', '2.0')
from gi.repository import UDisks

from ubunturecovery import metaclass

class udisks1_dev():
    def __init__(self):
        self.label = ''
        self.filesystem = ''
        self.devicefile = ''
        self.slave = ''
        self.parent = ''
        self.number = 0
        self.uuid = ''
        self.DeviceIsPartition = 0

def find_partition(label):
    return Udisks().find_partition(label)

class _DisksHnd(object, metaclass=metaclass.Singleton):
    def __init__(self):
        self.bus = dbus.SystemBus()
        try:
            self.obj = self.bus.get_object(self.busname, self.objpath)
            self.iface = dbus.Interface(self.obj, self.ifacename)
        except dbus.DBusException as msg:
            self.error(msg)
        self._disks = []

    def error(self, msg):
            print("%s, %s Failed" % (str(msg), self.__class__.__name__))

    def gettype(self):
        return self.__class__.__name__.lower()

    def find_partition(self, label):
        # TODO: refine functionality
        self.udisks = UDisks.Client.new_sync(None)
        manager = self.udisks.get_object_manager()

        for obj in manager.get_objects():
            block = obj.get_block()
            if not block:
                continue

            dev = udisks1_dev()
            check_label = block.get_cached_property("IdLabel")
            if check_label.get_string().lower() == label.lower():
                partition = obj.get_partition()
                if not partition:
                    continue
                dev.label = check_label.get_string()
                dev.filesystem = block.get_cached_property('IdType').get_string()
                dev.devicefile = block.get_cached_property('Device').get_bytestring().decode()
                dev.uuid = block.get_cached_property('IdUUID').get_string()
                dev.Size = block.get_cached_property('Size').get_uint64()
                dev.number = partition.get_cached_property('Number').get_uint32()
                dev.parent = partition.get_cached_property('Table').get_string()
                dev.slave = self.udisks.get_object(dev.parent).get_block().get_cached_property('Device').get_bytestring().decode()
                dev.DeviceIsPartition = 1
                break

        if dev.label == label:
            return dev
        else:
            return None

    def devices(self):
        try:
            return self._devices()
        except dbus.DBusException as msg:
            self.error(msg)

    def _devices(self):
        if self.gettype() != 'hal':
            self.udisks = UDisks.Client.new_sync(None)
            manager = self.udisks.get_object_manager()
            default_raid_vendor = 'Linux Software RAID'
            raids = {}

            # Get all mdraid devices
            for obj in manager.get_objects():
                mdraid = obj.get_mdraid()
                if mdraid:
                    level = mdraid.get_cached_property("Level").get_string()
                    if level in ('raid0', 'raid1', 'raid4', 'raid5', 'raid6', 'raid10'):
                        uuid = mdraid.get_cached_property("UUID").get_string()
                        raids[uuid] = level.upper()

            for obj in manager.get_objects():
                loop = obj.get_loop()
                block = obj.get_block()
                partition = obj.get_partition()

                if loop or \
                        partition or \
                        not block or \
                        block.get_cached_property("ReadOnly").get_boolean():
                    continue

                id_type = block.get_cached_property("IdType").get_string()
                if id_type == 'isw_raid_member':
                    continue

                device = block.get_cached_property('Device').get_bytestring().decode()

                # Check if the disk is the type of mdraid
                if device.startswith('/dev/md'):
                    if block.get_cached_property('Size').get_uint64() == int(0):
                        continue

                dev = udisks1_dev()

                if partition:
                    dev.DeviceIsPartition = 1
                    dev.number = partition.get_cached_property('Number').get_uint32()
                    dev.parent = partition.get_cached_property('Table').get_string()
                    if device.startswith('/dev/md'):
                        dev.slave = self.udisks.get_object(dev.parent).get_block().get_cached_property('Device').get_bytestring().decode()
                    else:
                        dev.slave = self.udisks.get_object(dev.parent).get_block().get_cached_property('PreferredDevice').get_bytestring().decode()

                if block:
                    dev.label = block.get_cached_property('IdLabel').get_string()
                    if device.startswith('/dev/md'):
                        dev.devicefile = block.get_cached_property('Device').get_bytestring().decode()
                    else:
                        dev.devicefile = block.get_cached_property('PreferredDevice').get_bytestring().decode()
                    dev.device = dev.devicefile
                    dev.filesystem = block.get_cached_property('IdType').get_string()
                    dev.uuid = block.get_cached_property('IdUUID').get_string()
                    dev.IsSystem = block.get_cached_property('HintSystem').get_string()

                    drive_name = block.get_cached_property('Drive').get_string()
                    if device.startswith('/dev/md'):
                        name = block.get_cached_property("PreferredDevice").get_bytestring().decode('utf-8').split("/")[-1]
                        uuid = block.get_cached_property("Id").get_string().split("-")[-1]
                        model = "%s %s %s" % (name, raids[uuid], uuid)
                        dev.ConnectionBus = ''
                        dev.Removable = 'false'
                        dev.Vendor = default_raid_vendor
                        dev.Model = model
                        dev.Size = block.get_cached_property('Size').get_uint64()
                        dev.ReadOnly = block.get_cached_property('ReadOnly').get_string()
                    elif drive_name != '/':
                        drive = self.udisks.get_object(drive_name).get_drive()
                        if drive:
                            dev.ConnectionBus = drive.get_cached_property('ConnectionBus').get_string()
                            dev.Removable = drive.get_cached_property('Removable').get_string()
                            dev.Vendor = drive.get_cached_property('Vendor').get_string()
                            dev.Model = drive.get_cached_property('Model').get_string()
                            dev.Size = drive.get_cached_property('Size').get_uint64()
                            dev.ReadOnly = drive.get_cached_property('ReadOnly')

                yield dev

class Udisks(_DisksHnd):
    busname = 'org.freedesktop.UDisks2'
    objpath = '/org/freedesktop/UDisks2'
    ifacename = 'org.freedesktop.UDisks2'
    devifacename = 'org.freedesktop.DBus.Properties'

    @property
    def disks(self):
        if not self._disks:
            self._disks = self.find_disks()
        return self._disks

    def find_disks(self):
        """get disks in this machine"""
        disks = []
        for dev in self.devices():
            #Skip USB, Removable Disks, Partitions, External, Loopback, Readonly
            if dev.ConnectionBus == 'usb' or \
               dev.Removable == 1 or \
               dev.DeviceIsPartition == 1 or \
               dev.IsSystem == 0 or \
               dev.ReadOnly == 1 :
                continue

            #if we made it this far, add it
            devicefile = dev.devicefile
            devicemodel = dev.Model
            devicevendor = dev.Vendor
            devicesize = dev.Size
            devicesize_gb = "%i" % (devicesize / 1000000000)
            disks.append([devicefile, devicesize, "%s GB %s %s (%s)" % (devicesize_gb, devicevendor, devicemodel, devicefile)])

        #If multiple candidates were found, record in the logs
        if len(disks) == 0:
            raise RuntimeError("Unable to find and candidate hard disks to install to.")
        if len(disks) > 1:
            disks.sort()
        return disks

    def disable_swap(self):
        """Disables any swap partitions in use"""
        from ubiquity import misc
        if misc.execute_root('swapoff', '-a') is False:
            raise RuntimeError("Error when disabling swap")

    def find_factory_rp_stats(self, label):
        """Uses udisks to find the RP of a system and return stats on it
           Only use this method during bootstrap."""
        recovery = {}
        dev = self.find_partition(label)
        if not dev:
            return None
        recovery["label" ] = label
        recovery["device"] = dev.devicefile
        recovery["fs"    ] = dev.filesystem
        recovery["slave" ] = dev.slave
        recovery["number"] = dev.number
        recovery["parent"] = dev.parent
        recovery["uuid"]   = dev.uuid
        recovery['size_gb'] = dev.Size / 1000000000
        return recovery
