#!/usr/bin/python3
# --------------------------------------------------------------------
# Copyright © 2014 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/>.
# --------------------------------------------------------------------

'''
Apply an Ubuntu Core system image update to the running system.

WARNING : Should be run within a new mount namespace to avoid
          disturbing the existing (read-only) root filesystem.

NOTE    : The signature files associated with each image file is *NOT*
          re-verified, since this must already have been handled by
          system-image-cli(1).
'''

import sys
import os
import shutil
import subprocess
import tempfile
import argparse
import tarfile

DIR_MODE = 0o750

# FIXME: Should query system-image-cli (see bug LP:#1380574).
BASE_DIR = '/userdata/cache'

# Magic file that records time last system image was successfully applied.
TIMESTAMP_FILE = '/userdata/.last_update'

# Directory to mount writable root filesystem.
MOUNT_TARGET = BASE_DIR + '/system'

# Extension added to the command-file to show the new image is in the
# process of being applied.
IN_PROGRESS_SUFFIX = 'applying'

script_name = os.path.basename(__file__)


class Parser():
    def __init__(self, options, commands, remove_list):
        self.dispatcher = {
            'format': self._format,
            'load_keyring': self._load_keyring,
            'mount': self._mount,
            'unmount': self._unmount,
            'update': self._update,
        }

        self.options = options
        self.commands = commands
        self.remove_list = remove_list
        self.full_image = False

        # Identify the directory the files the command file refers to
        # live in.
        self.file_dir = os.path.dirname(self.options.cmdfile)

    def _format(self, args):
        try:
            target = args[0]
        except:
            sys.stderr.write('WARNING: expected target\n')
            return

        if target == 'system':
            self.full_image = True
        else:
            sys.stderr.write('WARNING: ingoring format target {}'
                             ' (unsupported operation)\n'
                             .format(target))
            return

    def _load_keyring(self, args):
        try:
            keyring = args[0]
            signature = args[1]
        except:
            sys.stderr.write('WARNING: expected keyring and signature\n')
            return

        keyring = os.path.join(self.file_dir, keyring)
        signature = os.path.join(self.file_dir, signature)

        self.remove_list.append(keyring)
        self.remove_list.append(signature)

        # already handled by system-image-cli(1) on image download
        sys.stderr.write('INFO: ignoring keyring {} and signature {}'
                         ' (already handled)\n'
                         .format(keyring, signature))

    def _mount(self, args):
        try:
            target_type = args[0]
        except:
            sys.stderr.write('WARNING: expected target type\n')
            return

        # the mount argument is imperative - this script decides what to
        # mount and where :)
        if target_type != 'system':
            sys.stderr.write('WARNING: unknown mount target type\n')
            return

        target = MOUNT_TARGET

        try:
            os.makedirs(target, mode=DIR_MODE, exist_ok=True)
        except Exception as e:
            st = os.stat(target)
            if st.mode == DIR_MODE:
                sys.exit('ERROR: failed to create directory {}: {}'
                         .format(target, e))

        args = ['mount', '-o', 'bind', '/', target]
        proc = subprocess.Popen(args,
                                stderr=subprocess.PIPE,
                                universal_newlines=True)
        if proc.wait() != 0:
            stderr = proc.communicate()[1]
            sys.exit('ERROR: failed to mount ({}): {}'
                     .format(args, stderr))

    def _unmount(self, args):
        try:
            target_type = args[0]
        except:
            sys.stderr.write('WARNING: expected target type\n')
            return

        # the unmount argument is imperative - this script decides what to
        # mount and where :)
        if target_type != 'system':
            sys.stderr.write('WARNING: unknown mount target type\n')
            return

        target = MOUNT_TARGET

        args = ['umount', target]
        proc = subprocess.Popen(args,
                                stderr=subprocess.PIPE,
                                universal_newlines=True)

        if proc.wait() != 0:
            stderr = proc.communicate()[1]
            sys.exit('ERROR: failed to unmount ({}): {}'
                     .format(args, stderr))

        try:
            os.rmdir(target)
        except Exception as e:
            sys.exit('ERROR: failed to remove directory {}: {}'
                     .format(target, e))

    def finish(self):
        '''
        Final tidy-up.
        '''

        for file in self.remove_list:
            if self.options.debug:
                print("DEBUG: Removing file {}".format(file))
            os.remove(file)

        # Update timestamp file
        open(TIMESTAMP_FILE, 'w').close()

    def _update(self, args):
        try:
            file = args[0]
            signature = args[1]
        except:
            sys.stderr.write('WARNING: expected file and signature\n')
            return

        file = os.path.join(self.file_dir, file)
        signature = os.path.join(self.file_dir, signature)

        for f in (file, signature):
            if not os.path.exists(f):
                sys.stderr.write('WARNING: ignoring missing file {}\n'
                                 .format(f))
                return

            self.remove_list.append(f)

        print('Applying update: {}'.format(file))
        tar = tarfile.open(file)

        # The tar file contains a system/ directory with the new files
        # to apply to the real system. It may also contain a top-level
        # file called 'removed' that lists files relative to the system
        # mount for files on the system that should be removed before
        # unpacking the rest of the archive.
        #
        # For example, if the file /etc/passwd should be removed, it
        # would be specified in the removed_file as 'system/etc/passwd'.

        found_removed_file = False
        removed_file = 'removed'

        if not self.full_image:
            if removed_file in tar.getnames():
                found_removed_file = True
                tmpdir = tempfile.mkdtemp(prefix=script_name)
                tar.extract(path=tmpdir, member=tar.getmember(removed_file))

                removed_file_path = os.path.join(tmpdir, removed_file)
                with open(removed_file_path, 'r') as fh:
                    to_remove = fh.readlines()

                # process backwards to work around bug LP:#1381134.
                for remove in sorted(to_remove, reverse=True):
                    remove = remove.strip()
                    final = os.path.join(BASE_DIR, remove)
                    if not os.path.exists(final):
                        # This scenario can only mean there is a bug
                        # with system-image generation (or someone made
                        # the image writable and removed some files
                        # manually).
                        print('WARNING: ignoring non-existent file {}'
                              .format(final))
                        continue

                    if self.options.debug:
                        print('DEBUG: Removing file {}'.format(final))
                    try:
                        if os.path.isdir(final) and not os.path.islink(final):
                            shutil.rmtree(final)
                        else:
                            os.remove(final)
                    except Exception as e:
                        sys.stderr.write('WARNING: failed to remove {}: {}\n'
                                         .format(final, e))
                shutil.rmtree(tmpdir)

        # Unpack everything else on top of the system partition
        tar.extractall(path=BASE_DIR)
        tar.close()

        if found_removed_file:
            os.remove(os.path.join(BASE_DIR, removed_file))

    def run(self):
        '''
        Execute the commands in the command file
        '''
        for cmdline in self.commands:
            cmdline = cmdline.strip()

            args = cmdline.split()
            cmd = args[0]
            args = args[1:]

            if self.dispatcher[cmd]:
                if self.options.debug:
                    print('DEBUG: Running dispatcher {} ({})'
                          .format(cmd, args))
                self.dispatcher[cmd](args)
            else:
                sys.stderr.write('WARNING: ignoring bogus input line: {}\n'
                                 .format(cmdline))


def main():
    parser = argparse.ArgumentParser(description='System image Upgrader')

    parser.add_argument(
        '--debug',
        nargs='?', const=1, default=0, type=int,
        help='Dump debug info (specify numeric value to increase verbosity)')

    parser.add_argument(
        '-t', '--tmpdir',
        help='Specify name for pre-existing temporary directory to use')

    parser.add_argument(
        'cmdfile', action="store",
        help='Name of file containing commands to execute')

    options = parser.parse_args()

    try:
        with open(options.cmdfile, 'r') as fh:
            commands = fh.readlines()
    except:
        sys.exit('Failed to read command file: {}'.format(options.cmdfile))

    # Rename the file to show the upgrade is in progress
    in_progress = '{}.{}'.format(options.cmdfile, IN_PROGRESS_SUFFIX)
    shutil.move(options.cmdfile, in_progress)

    # Remember to remove it
    remove_list = [in_progress]

    parser = Parser(options, commands, remove_list)

    parser.run()
    parser.finish()


if __name__ == '__main__':
    main()
