# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utility module for working with USB mass storage device.

This module uses a combination of ``pyudev`` and ``psutil`` to list and filter
USB mass storage devices.

.. note:

    According to the udiskctl man page
    http://manpages.ubuntu.com/manpages/trusty/man8/udisks.8.html#contenttoc5

    Therefore, applications can not rely on tools like e.g.  udisksctl(1) to be
    available. Additionally, there is no guarantee that the options,
    command-line switches etc. of command-line tools or similar will remain
    stable.

    Instead, applications should only use the D-Bus API, the libudisks2
    library or tools such as dbus-send(1) or gdbus(1) to interact with
    udisksd(8).
"""
from __future__ import absolute_import
import os
import re
import shutil
import stat
import subprocess
import sys
import time
import psutil
from gazoo_device import errors
from gazoo_device import gdm_logger
from gazoo_device.utility import usb_utils


logger = gdm_logger.get_gdm_logger()


def _validate_suid_bit(cmd_str):
    if sys.platform == "darwin":
        return
    binary = cmd_str.split()[0]
    if not os.stat(binary).st_mode and stat.S_ISUID:
        raise errors.GazooDeviceError("Missing SUID bit on {binary} "
                                      "Fix using 'sudo chmod u+s {binary}'".format(binary=binary))


class UsbStorage(object):
    """USB mass storage device class.

    Used in 'with block' and will mount the USB mass storage device upon entering the with
    block and eject the USB mass storage device upon exiting the 'with block'. If
    the UsbStorage class is not used within a 'with block' then the methods of
    mount, unmount, and eject can be called separately.
    """
    commands_dict = {
        "darwin": {
            "EJECT": 'diskutil eject {}',
            "MOUNT": 'diskutil mount {}',
            "UNMOUNT": 'diskutil unmount {}'
        },
        "linux": {
            "EJECT": '/usr/bin/eject {}',
            "MOUNT": '/usr/bin/udisksctl mount --no-user-interaction -b {}',
            "UNMOUNT": '/bin/umount {}'
        }}

    def __init__(self, serial_number, add_log_note_func=None, eject_on_exit=True):
        """Initialize the USB mass storage device instance with the information specified.

        Args:
            serial_number (str): serial number to identify device.
            add_log_note_func (func): function to log usb serial methods to device logs.
            eject_on_exit (bool): whether to eject or unmount on exit from context manager.
        """
        self._serial_number = serial_number
        self.add_log_note_func = add_log_note_func
        self.eject_on_exit = eject_on_exit
        if sys.platform == "darwin":
            self.commands = self.commands_dict["darwin"]
        else:
            self.commands = self.commands_dict["linux"]

    def __enter__(self):
        self.mount()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.eject_on_exit:
            self.eject()
        else:
            self.unmount()

    @property
    def disk(self):
        """Disk that device mounted to. Empty if not mounted."""
        devices_info = self._get_devices_info()
        return devices_info[0].disk if devices_info else ""

    @property
    def mount_point(self):
        """Returns the mass storage mount point to the caller.

        Returns:
            str: mount path location for caller to use in opening files on the mass storage
                    device
        """
        usb_endpoint_infos = self._get_devices_info()
        if usb_endpoint_infos:
            partitions = psutil.disk_partitions()
            for usb_info_inst in usb_endpoint_infos:
                for partition in partitions:
                    if partition.device == usb_info_inst.disk:
                        return partition.mountpoint
        return None

    def add_log_note(self, msg):
        """Adds note in device logs if method available."""
        if self.add_log_note_func:
            self.add_log_note_func(msg)

    def copy_file_to(self, source_file_path, relative_destination_dir=None):
        """Copies source file to destination directory specified.

        Args:
            source_file_path (str): path to source file to copy to mass storage device
            relative_destination_dir (str): optional path relative to the mount path of the mass
                                            storage device

        Raises:
            GazooDeviceError: if storage is not mounted, the mount path doesn't exist,
                             or if the source file doesn't exist after copying the source
                             file to the mount path + relative destination directory destination

        Note:
            The storage device must be mounted before calling this method. If the relative
            destination directory specified doesn't exist it will be created before the
            source file is copied
        """
        if not self.is_mounted():
            raise errors.GazooDeviceError("Device copy file failed. "
                                          "Device with Serial {} is not mounted.".
                                          format(self._serial_number))

        mount_point = self.mount_point
        if relative_destination_dir is not None:
            dst_file_path = os.path.join(mount_point, relative_destination_dir)
        else:
            dst_file_path = mount_point

        if not os.path.exists(dst_file_path):
            logger.debug("Creating destination path %s", dst_file_path)
            os.makedirs(dst_file_path)

        file_name = os.path.basename(source_file_path)
        dst_file_path = os.path.join(dst_file_path, file_name)
        msg = "Copying {} to {}".format(source_file_path, dst_file_path)
        logger.debug(msg)
        self.add_log_note(msg)
        shutil.copyfile(source_file_path, dst_file_path)
        if not os.path.exists(dst_file_path):
            raise errors.GazooDeviceError("Device copy file failed. "
                                          "File does not exist at destination {}"
                                          "Check destination path permissions.".
                                          format(dst_file_path))
        msg = "Copying {} to {} completed".format(source_file_path, dst_file_path)
        logger.debug(msg)
        self.add_log_note(msg)

    def eject(self, timeout=100.0):
        """Unmounts and ejects the mass storage device within the timeout specified.

        Args:
            timeout (float): time to wait to confirm the mass storage has been unmounted and
                           ejected

        Raises:
            GazooDeviceError: if timeout is exceeded before storage is confirmed as being
                             unmounted or ejected.
        """
        start_time = time.time()
        self.add_log_note("{} eject requested.".format(self._serial_number))
        if self.is_device_ejected(timeout=.1):
            self.add_log_note("{} already ejected.".format(self._serial_number))
            return
        self.unmount(timeout=timeout)
        _validate_suid_bit(self.commands["EJECT"])
        try:
            cmd_list = self.commands["EJECT"].format(self.disk).split()
            subprocess.check_output(cmd_list)
        except subprocess.CalledProcessError as err:
            raise errors.GazooDeviceError("Device eject failed. "
                                          "{} produced this error: {!r}".format(
                                              cmd_list,
                                              err.output))
        if self.is_device_ejected(timeout=timeout):
            self.add_log_note("{} ejected in {}s".format(
                self._serial_number, int(time.time() - start_time)))
            return
        raise errors.GazooDeviceError("Device eject failed. "
                                      "Unable to eject USB mass storage {} in {}s"
                                      "Check if device is currently in use. ".
                                      format(self.mount_point, timeout))

    def is_device_ejected(self, timeout=20):
        """Check whether device is ejected from the system.

        Args:
            timeout (int):  Number of seconds to wait for successful eject.
        Returns:
            bool: True if device is ejected.
        """
        exit_time = time.time() + timeout
        while time.time() < exit_time:
            devices_info = self._get_devices_info()
            if not devices_info:
                return True
            if sys.platform != "darwin":
                device_id_path = devices_info[0].disk
                with open('/etc/fstab', 'r') as open_file:
                    text = open_file.read()
                if not re.search(device_id_path, text):
                    return True
            time.sleep(.5)
        return False

    def is_found(self):
        """Returns True if the mass storage device has been detected and is available to mount.

        Returns:
            bool: True if mass storage device is detected
        """
        return bool(self._get_devices_info())

    def is_mounted(self, timeout=20):
        """Returns True if the mass storage device is currently mounted.

        Args:
            timeout (float): time to wait for it to mount before timeout.
        Returns:
            bool: True if mass storage device is mounted

        Note:
           Uses a timeout as not all devices are using mount to mount the device and it
           can take a few seconds for the computers to register the external mount.
        """
        start_time = time.time()
        exit_time = start_time + timeout
        while time.time() < exit_time:
            if self.mount_point:
                msg = "{} found mounted at {} in {}s".format(
                    self._serial_number, self.mount_point, int(time.time() - start_time))
                self.add_log_note(msg)
                logger.debug(msg)
                return True
        return False

    def mount(self, timeout=180):
        """Mounts the mass storage device within the timeout specified.

        Args:
            timeout (float): time to wait to confirm storage has been mounted and/or after
                             manually attempting to mount the storage

        Raises:
            GazooDeviceError: if timeout is exceeded before storage is confirmed as
            being mounted.

        Note:
            Since there is a udev rule that is responsible for performing the actual mounting of
            the device this method will wait up to the timeout specified before attempting to
            manually mount the device and then wait again for the timeout specified to confirm
            that the device has been successfully mounted
        """
        self.add_log_note("{} mount requested.".format(self._serial_number))
        if self.is_mounted():
            return
        if not self.is_found():
            self.add_log_note("Error: {} not found on machine.".format(self._serial_number))
            raise errors.GazooDeviceError("Device mount failed. "
                                          "USB mass storage matching Serial#={} was not found. "
                                          "\nCheck cables and device or try rebooting.".
                                          format(self._serial_number))
        msg = "Mounting {}".format(self.disk)
        logger.debug(msg)
        self.add_log_note(msg)
        _validate_suid_bit(self.commands["MOUNT"])
        try:
            cmd_list = self.commands["MOUNT"].format(self.disk).split()
            subprocess.check_output(cmd_list)
        except subprocess.CalledProcessError as err:
            if isinstance(err.output, bytes):
                err.output = err.output.decode('utf-8')
            if 'Not authorized' in err.output:
                err = "Try running 'sudo chmod u+s {}' to fix permission error.".format(
                    ' '.join(cmd_list[0]))
            raise errors.GazooDeviceError("Device mount failed. Err: {!r}".format(err))

        if not self.is_mounted(timeout=timeout):
            raise errors.GazooDeviceError("Device mount failed. Not mounted after {}s".format(
                timeout))

    def move_file_from(self, relative_source_file_path, destination, delete_file=True):
        """Moves file from storage mount point to destination directory specified.

        Args:
            relative_source_file_path (str): relative path to file on device to copy over.
            destination (str): destination directory or file name path
            delete_file (bool): delete file from device.

        Raises:
            GazooDeviceError: if storage is not mounted, the mount path doesn't exist,
                             or if the source file doesn't exist

        Note:
            The storage device must be mounted before calling this method.

        """
        if not self.is_mounted():
            raise errors.GazooDeviceError("Move_file_from failed. "
                                          "Device with Serial {} is not mounted.".
                                          format(self._serial_number))

        source_file_path = os.path.join(self.mount_point,
                                        relative_source_file_path)

        if not os.path.exists(source_file_path):
            raise errors.GazooDeviceError("Device copy file failed. "
                                          "Source file {} doesn't exist. ".
                                          format(source_file_path))

        msg = "Copying {} to {}".format(source_file_path, destination)
        logger.debug(msg)
        self.add_log_note(msg)
        shutil.copy(source_file_path, destination)

        if delete_file:
            logger.debug("Removing %s", source_file_path)
            os.remove(source_file_path)

        if not os.path.exists(destination):
            raise errors.GazooDeviceError("Device copy file failed. "
                                          "Destination {} does not exist. "
                                          "Check destination path permissions.".
                                          format(destination))
        msg = "Copying {} to {} completed.".format(relative_source_file_path, destination)
        logger.debug(msg)
        self.add_log_note(msg)

    def unmount(self, timeout=180.0):
        """Unmounts the mass storage device within the timeout specified.

        Args:
            timeout (float): time to wait to confirm storage and successfully unmounted

        Raises:
            GazooDeviceError: if timeout is exceeded before storage is confirmed as being
                             unmounted
        """
        start_time = time.time()
        exit_time = start_time + timeout
        self.add_log_note("{} unmount requested".format(self._serial_number))
        mount_point = self.mount_point
        if not mount_point:
            self.add_log_note("{} already unmounted".format(self._serial_number))
            return
        self.add_log_note(
            "{} unmounting from {}".format(
                self._serial_number,
                mount_point))

        _validate_suid_bit(self.commands["UNMOUNT"])
        try:
            cmd_list = self.commands["UNMOUNT"].format(mount_point).split()
            subprocess.check_output(cmd_list)
        except subprocess.CalledProcessError as err:
            raise errors.GazooDeviceError("Device unmount failed. "
                                          "Error when calling /bin/umount: {}".format(
                                              err.output))

        while time.time() < exit_time:
            if not self.mount_point:
                msg = "{} successfully unmounted in {}s".format(
                    self._serial_number,
                    int(time.time() - start_time))
                logger.debug(msg)
                self.add_log_note(msg)
                return

        raise errors.GazooDeviceError("Device unmount failed. "
                                      "Unmounting USB mass storage at {} failed after {}s".
                                      format(self.mount_point, timeout))

    def _get_devices_info(self):
        """Returns info for the device.

        Returns:
          list: one or more matching usbinfo objects for each usb endpoint that matches.
        """
        return usb_utils.find_matching_connections(
            {"serial_number": {"include_regex": self._serial_number}})
