#!/usr/bin/env python3
# Copyright 2015 Canonical Ltd.
# Written by:
#   Shawn Wang <shawn.wang@canonical.com>
#
# 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/>.

"""Show the recovery partition information for the preinstalled OS."""

import os
import re
import subprocess
import sys
import tempfile
import unittest
import xml.dom.minidom as minidom

from guacamole import Command

try:
    from unittest import mock
except ImportError:
    from plainbox.vendor import mock

RECOVERY_PACKAGES = ["dell-recovery", "ubuntu-recovery"]


def get_recovery_package():
    """
    Test with RECOVERY_PACKAGES.

    to check recovery application is installed or not

    :return:
        string of package_version or None
    """
    for pkg in RECOVERY_PACKAGES:
        output = subprocess.check_output(["apt-cache", "policy", pkg],
                                         universal_newlines=True)
        for line in output.split("\n"):
            if line.startswith("  Installed:"):
                ver = line.split(": ")[1]
                return "{}_{}".format(pkg, ver.strip())
    return None


RECOVERY_LABELS = {"HP_TOOLS": "HP",
                   "PQSERVICE": "UBUNTU",
                   "BACKUP": "TEST",
                   "INSTALL": "DELL",
                   "OS": "DELL",
                   "RECOVERY": "DELL"}


_escape_pattern = re.compile(r'\\x([0-9a-fA-F][0-9a-fA-F])')


def lsblk_unescape(label):
    """Un-escape text escaping done by lsblk(8)."""
    return _escape_pattern.sub(
        lambda match: chr(int(match.group(1), 16)), label)


def get_recovery_partition():
    """
    Get the type and location of the recovery partition.

    :return:
        (recovery_type, recovery_partition) or None

    Use lsblk(8) to inspect available block devices looking
    for a partition with FAT or NTFS and a well-known label.
    """
    cmd = ['lsblk', '-o', 'TYPE,FSTYPE,NAME,LABEL', '--raw']
    for line in subprocess.check_output(cmd).splitlines()[1:]:
        type, fstype, name, label = line.split(b' ', 3)
        # Skip everything but partitions
        if type != b'part':
            continue
        # Skip everything but FAT and NTFS
        if fstype != b'vfat' and fstype != b'ntfs':
            continue
        label = lsblk_unescape(label.decode('utf-8'))
        recovery_type = RECOVERY_LABELS.get(label)
        # Skip unknown labels
        if recovery_type is None:
            continue
        recovery_partition = '/dev/{}'.format(name.decode('utf-8'))
        return (recovery_type, recovery_partition)


class FunctionTests(unittest.TestCase):

    """Tests for several functions."""

    @mock.patch('subprocess.check_output')
    def test_get_recovery_package(self, mock_subprocess_check_output):
        """Smoke test for get_recovery_package()."""
        mock_subprocess_check_output.return_value = """\
dell-recovery:
  Installed: 1.11
  Candidate: 1.11
  Version table:
     1.11
        500 https://archive/cesg-mirror/ test/public amd64 Packages
"""
        self.assertEqual(get_recovery_package(),
                         "dell-recovery_1.11")

    @mock.patch('subprocess.check_output')
    def test_get_recovery_partition(self, mock_subprocess_check_output):
        """Smoke test for get_recovery_partition()."""
        mock_subprocess_check_output.return_value = (
            b'TYPE FSTYPE NAME LABEL\n'
            b'disk linux_raid_member sda fx:2x250GB\n'
            b'raid1 bcache md127 \n'
            b'disk ext4 bcache0 Ultra\n'
            b'disk linux_raid_member sdb fx:2x250GB\n'
            b'raid1 bcache md127 \n'
            b'disk ext4 bcache0 Ultra\n'
            b'disk  sdc \n'
            b'part btrfs sdc1 vol1\n'
            b'disk  sdd \n'
            b'part ntfs sdd1 Windows\x208.1\n'
            b'part  sdd2 \n'
            b'part ext4 sdd5 Utopic\n'
            b'part swap sdd6 \n'
            b'disk bcache sde \n'
            b'disk ext4 bcache0 Ultra\n'
            b'disk  sdf \n'
            b'part ntfs sda3 RECOVERY\n')
        self.assertEqual(get_recovery_partition(), ("DELL", "/dev/sda3"))

    def test_lsblk_unescape(self):
        """Smoke tests for lsblk_unescape()."""
        self.assertEqual(lsblk_unescape('Windows\\x208.1'), 'Windows 8.1')
        self.assertEqual(lsblk_unescape('Windows XP'), 'Windows XP')


class MountedPartition(object):

    """
    Mount Manager to mount partition on tempdir.

    e.g.
    with MountedPartition("/dev/sda1") as tmp:
        print("This is the mount point: {}".format(tmp))
        do_stuff()
    """

    def __init__(self, part):
        """
        Prepare the mntdir point.

        :param part: string of the partition device file, like /dev/sda2
        """
        self.part = part
        self.mntdir = tempfile.mkdtemp()

    def __enter__(self):
        """
        __enter__ method for python's with statement.

        Mount the partition device to the mntdir.
        """
        cmd = ["mount", self.part, self.mntdir]
        subprocess.check_output(cmd, universal_newlines=True)
        return self.mntdir

    def __exit__(self, type, value, traceback):
        """
        __exit__ method for python's with statement.

        Unmount and remove the mntdir.
        """
        subprocess.check_output(["umount", self.mntdir],
                                universal_newlines=True)
        os.rmdir(self.mntdir)


class MountedPartitionTests(unittest.TestCase):

    """Unittest of MountedPartition."""

    @mock.patch('subprocess.check_output')
    def test_with_of_MountedPartition(self, mock_subprocess_check_output):
        """Test mount point."""
        test_dir = ""
        with MountedPartition("/dev/test") as tmp:
            test_dir = tmp
            self.assertTrue(os.path.exists(test_dir))
            mock_subprocess_check_output.assert_has_calls(
                [mock.call(['mount', '/dev/test', test_dir],
                           universal_newlines=True)])
        self.assertFalse(os.path.exists(test_dir))
        mock_subprocess_check_output.assert_has_calls(
            [mock.call(['umount', test_dir],
                       universal_newlines=True)])


class RecoveryVersion(Command):

    """
    print the version of recovery image.

    @EPILOG@

    This commands prints information such as:

    image_version: xxx
    bto_version: REV_xxx.iso (dell only)
    """

    def invoked(self, ctx):
        """
        Guacamole method called when the command is invoked.

        /etc/buildstamp is a image information file,
        it created by the oem image builder.

        oilpalm Fri, 20 Jun 2014 04:02:07 +0000
        somerville-trusty-amd64-20140620-0

        If /etc/buildstamp exist, print out the second line (image iso name).

        For Dell-recovery partition, /etc/buildstamp shows base image info.
        If recovery_partition/bto.xml,
        print out the bto_version (read from xml file).
        """
        if os.path.isfile("/etc/buildstamp"):
            with open('/etc/buildstamp', 'rt', encoding='UTF-8') as stream:
                data = stream.readlines()
                print("image_version: {}".format(data[1].strip()))

        with MountedPartition(ctx.recovery_partition) as mntdir:
            fname = "{}/bto.xml".format(mntdir)
            if os.path.isfile(fname):
                o = minidom.parse("{}/bto.xml".format(mntdir))
                bto_version = o.getElementsByTagName("iso")[0].firstChild.data
                print("bto_version: {}".format(bto_version))


class RecoveryFile(Command):

    """
    display a single file from the recovery partition

    This command can be used to ``cat`` any file from the recovery partition
    """

    def register_arguments(self, parser):
        """
        Guacamole method used by the argparse ingredient.

        :param parser:
            Argument parser (from :mod:`argparse`) specific to this command.
        """
        parser.add_argument('file', help='name of the file to display')

    def invoked(self, ctx):
        """
        Guacamole method used by the command ingredient.

        :param ctx:
            The guacamole context object. Context provides access to all
            features of guacamole. The argparse ingredient adds the ``args``
            attribute to it. That attribute contains the result of parsing
            command line arguments.
        :returns:
            The return code of the command. Guacamole translates ``None`` to a
            successful exit status (return code zero).
        """
        with MountedPartition(ctx.recovery_partition) as mnt:
            return subprocess.call([
                'cat', '--', os.path.join(mnt, ctx.args.file)])


class RecoveryCheckType(Command):

    """
    test if the recovery partition is of the given type.

    This command can be used for scripted tests, to see if the recovery
    partition on the current system is of a concrete type or not (e.g.
    DELL-specific)

    @EPILOG@

    The exit code is 0 if the recovery partition type matches and 1 otherwise.
    """

    def register_arguments(self, parser):
        """
        Guacamole method used by the argparse ingredient.

        :param parser:
            Argument parser (from :mod:`argparse`) specific to this command.
        """
        parser.add_argument(
            'type', help="expected type of the recovery partition")

    def invoked(self, ctx):
        """
        Guacamole method used by the command ingredient.

        :param ctx:
            The guacamole context object. Context provides access to all
            features of guacamole. The argparse ingredient adds the ``args``
            attribute to it. That attribute contains the result of parsing
            command line arguments.
        :returns:
            The return code of the command. Guacamole translates ``None`` to a
            successful exit status (return code zero).
        """
        if ctx.recovery_type != ctx.args.type:
            return 1


class RecoveryInfo(Command):

    """
    Inspect the recovery partition.

    This command can be used to inspect the recovery partition. It has several
    sub-commands that do various tasks.  If the system has no recovery
    partition, the command exits with the error code 1.
    """

    sub_commands = (
        ('version', RecoveryVersion),
        ('file', RecoveryFile),
        ('checktype', RecoveryCheckType),
    )

    def invoked(self, ctx):
        """
        Guacamole method used by the command ingredient.

        :param ctx:
            The guacamole context object. Context provides access to all
            features of guacamole. The argparse ingredient adds the ``args``
            attribute to it. That attribute contains the result of parsing
            command line arguments.
        :returns:
            The return code of the command. Guacamole translates ``None`` to a
            successful exit status (return code zero).
        """
        partition = get_recovery_partition()

        if partition is None:
            print("Recovery partition not found", file=sys.stderr)
            return 1
        (recovery_type, recovery_partition) = partition
        ctx.recovery_partition = recovery_partition
        ctx.recovery_type = recovery_type


class RecoveryInfoTests(unittest.TestCase):

    """Tests for RecoveryInfo."""

    @mock.patch('__main__.get_recovery_package')
    @mock.patch('__main__.get_recovery_partition')
    def test_smoke(self, mock_get_recovery_partition,
                   mock_get_recovery_package):
        """Smoke tests for running recovery_info."""
        mock_get_recovery_partition.return_value = ("DELL", "/dev/sda3")
        mock_get_recovery_package.return_value = "dell-recovery_1.11"
        self.assertEqual(RecoveryInfo().main(argv=[], exit=False), 0)
        self.assertEqual(
            RecoveryInfo().main(argv=["checktype", "HP"], exit=False), 1)
        self.assertEqual(
            RecoveryInfo().main(argv=["checktype", "DELL"], exit=False), 0)


if __name__ == '__main__':
    if '--test' in sys.argv:
        sys.argv.remove('--test')
        unittest.main()
    else:
        RecoveryInfo().main()
