# -*- coding: utf-8 -*-
# --------------------------------------------------------------------
# Copyright © 2014-2015 Canonical Ltd.
#
# This program is 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, version 3 of the License, or
# (at your option) any later version.
#
# 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 os
import tempfile
import tarfile
import shutil
import unittest

# file mode to use for creating test directories.
TEST_DIR_MODE = 0o750

script_name = os.path.basename(__file__)


def make_tmp_dir(tag=None):
    '''
    Create a temporary directory.
    '''

    if tag:
        prefix = '{}-{}-'.format(script_name, tag)
    else:
        prefix = script_name

    return tempfile.mkdtemp(prefix=prefix)


def append_file(path, contents):
    '''
    Append to a regular file (create it doesn't exist).
    '''

    dirname = os.path.dirname(path)
    os.makedirs(dirname, mode=TEST_DIR_MODE, exist_ok=True)

    with open(path, 'a') as fh:
        if contents:
            fh.writelines(contents)

            if not contents.endswith('\n'):
                fh.write('\n')


def create_file(path, contents):
    '''
    Create a regular file.
    '''
    append_file(path, contents)


def mock_get_root_partitions_by_label():
    '''
    Fake disk partition details for testing.
    '''
    matches = []

    matches.append(('system-a', '/dev/sda3', '/'))
    matches.append(('system-b', '/dev/sda4', '/writable/cache/system'))

    return matches


class UpdateTree():
    '''
    Representation of a directory tree that will be converted into an
    update archive.
    '''
    TEST_REMOVED_FILE = 'removed'
    TEST_SYSTEM_DIR = 'system/'

    def __init__(self):

        # Directory tree used to construct the tar file from.
        # Also used to hold the TEST_REMOVED_FILE file.
        self.dir = make_tmp_dir(tag='UpdateTree-tar-source')

        self.removed_file = os.path.join(self.dir, self.TEST_REMOVED_FILE)

        # Directory to place create/modify files into.
        self.system_dir = os.path.join(self.dir, self.TEST_SYSTEM_DIR)

        # Directory used to write the generated tarfile to.
        # This directory should also be used to write the command file
        # to.
        self.tmp_dir = make_tmp_dir(tag='UpdateTree-cache')

    def destroy(self):
        if os.path.exists(self.dir):
            shutil.rmtree(self.dir)

        if os.path.exists(self.tmp_dir):
            shutil.rmtree(self.tmp_dir)

    def add_to_removed_file(self, removed_files):
        '''
        Add the specified list of files to the removed file.

        The 'removed' file is simply a file with a well-known name that
        contains a list of files (one per line) to be removed from a
        system before the rest of the update archive is unpacked.

        :param removed_files: list of file names to add to the removed file.

        '''
        # all files listed in the removed list must be system files
        final = list(map(lambda a:
                     os.path.normpath('{}{}'.format(self.TEST_SYSTEM_DIR,
                                                    a)), removed_files))

        contents = "".join(final)
        append_file(self.removed_file, contents)

    def tar_filter(self, member):
        '''
        Function to filter the tarinfo members before creating the
        archive.
        '''
        # members are created with relative paths (no leading slash)
        path = os.sep + member.name

        if member.name == '/.':
            return None

        i = path.find(self.dir)
        assert(i == 0)

        # remove the temporary directory elements
        # (+1 for the os.sep we added above)
        member.name = path[len(self.dir)+1:]

        return member

    def create_archive(self, name):
        '''
        Create an archive with the specified name from the UpdateTree
        object. Also creates a fake signature file alongside the archive
        file since this is currently required by the upgrader (although
        it is not validated).

        :param name: name of tarfile.
        :param name: full path to xz archive to create.
        :return full path to tar file with name @name.
        '''

        self.tar_path = os.path.join(self.tmp_dir, name)
        tar = tarfile.open(self.tar_path, 'w:xz')

        # We can't just add recursively since that would attempt to add
        # the parent directory. However, the real update tars don't
        # include that, and attempting to ignore the parent directory
        # results in an empty archive. So, walk the tree and add
        # file-by-file.
        for path, names, files in os.walk(self.dir):
            for file in files:
                full = os.path.join(path, file)
                tar.add(full, recursive=False, filter=self.tar_filter)
            if not files and not names:
                # add (empty) directories
                tar.add(path, recursive=False, filter=self.tar_filter)

        tar.close()

        signature = '{}.asc'.format(self.tar_path)

        with open(signature, 'w') as fh:
            fh.write('fake signature file')

        return self.tar_path


class UbuntuCoreUpgraderTestCase(unittest.TestCase):
    '''
    Base class for Upgrader tests.

    Most of the tests follow a standard pattern:

    1) Create an UpdateTree object:

         update = UpdateTree()

       This creates 2 temporary directories:

       - self.system_dir: Used to generate an update archive from.

       - self.tmp_dir: Used to write the generated archive file to. The
         intention is that this directory should also be used to hold
         the command file.

    2) Removal tests call update.add_to_removed_file(file) to add a
       particular file to the removals file in the update archive.

    3) Create/Modify tests create files below update.system_dir.

    4) Create an update archive (which includes the removals file
       and all files below update.system_dir):

         archive = update.create_archive(self.TARFILE)

    5) Create a command file (which tells the upgrader what to do
       and which archive files to apply):

         make_command_file(...)

    6) Create a victim directory. This is a temporary directory where
       the upgrade will happen.

    7) Start the upgrade:

         call_upgrader(...)

    8) Perform checks on the victim directory to ensure that upgrade
       did what was expected.

    '''

    TARFILE = 'update.tar.xz'

    # result of last test run. Hack to deal with fact that even if a
    # test fails, unittest still calls .tearDown() (whomever thought
    # that was a good idea...?)
    currentResult = None

    def setUp(self):
        '''
        Test setup.
        '''
        # Create an object to hold the tree that will be converted into
        # an upgrade archive.
        self.update = UpdateTree()

        # The directory which will have the update archive applied to
        # it.
        self.victim_dir_base = make_tmp_dir(tag='victim')
        self.victim_dir = os.path.join(self.victim_dir_base, 'system')
        os.makedirs(self.victim_dir, mode=0o750)

    def tearDown(self):
        '''
        Test cleanup.
        '''

        if not self.currentResult.wasSuccessful():
            # Do not clean up - the only sane option if a test fails.
            return

        self.update.destroy()
        self.update = None

        shutil.rmtree(self.victim_dir)
        self.victim_dir = None

    def run(self, result=None):
        self.currentResult = result
        unittest.TestCase.run(self, result)
