#!/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 logging
import logging.handlers
import atexit
import shutil
import subprocess
import argparse

from ubuntucoreupgrader.unshare import unshare, UnshareFlags
from ubuntucoreupgrader.upgrader import Upgrader

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

# seconds to wait before a reboot (should one be required)
DEFAULT_REBOOT_DELAY = 2

DEFAULT_ROOT = '/'

log = logging.getLogger()


def shutdown_logging():
    logging.shutdown()


def setup_logger(options):

    atexit.register(shutdown_logging)

    if options.debug:
        log.setLevel(logging.DEBUG)

    if options.dry_run or options.check_reboot:
        # use default logger which displays output to stderr only.
        return

    # send data to syslog...
    handler = logging.handlers.SysLogHandler(address='/dev/log')

    # add an identifier
    script = os.path.splitext(os.path.basename(__file__))[0]
    handler.ident = '{}: '.format(script)

    log.addHandler(handler)

    # ... and stderr
    handler = logging.StreamHandler()
    log.addHandler(handler)


def parse_args():
    '''
    Handle command-line options.
    '''

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

    parser.add_argument(
        '--check-reboot',
        action='store_true',
        help='''
        Attempt to determine if a reboot may be required.

        This option is similar to --dry-run: no system changes are made.

        Note that the result of this command cannot be definitive since a
        reboot would be triggered if a service failed to start (but this
        option does not attempt to restart any services).
        ''')

    parser.add_argument(
        '--clean-only',
        action='store_true',
        help='Clean up from a previous upgrader run, but do not upgrade)')

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

    parser.add_argument(
        '-n', '--dry-run',
        action='store_true',
        help='''
        Simulate an update including showing processes that have locks
        on files
        ''')

    parser.add_argument(
        '--force-inplace-upgrade',
        action='store_true',
        help='Apply an upgrade to current rootfs even if running " \
        "on dual rootfs system')

    parser.add_argument(
        '--leave-files',
        action='store_true',
        help='Do not remove the downloaded system image files after upgrade')

    parser.add_argument(
        '--no-reboot',
        action='store_true',
        help='Do not reboot even if one would normally be required')

    parser.add_argument(
        '--reboot-delay',
        type=int, default=DEFAULT_REBOOT_DELAY,
        help='''
        Wait for specified number of seconds before rebooting
        (default={})
        '''.format(DEFAULT_REBOOT_DELAY))

    parser.add_argument(
        '--root-dir',
        default=DEFAULT_ROOT,
        help='Specify an alternative root directory (for testing ONLY)')

    parser.add_argument(
        '--show-other-details',
        action='store_true',
        help='Dump the details of the system-image vesion on the " \
        "other root partition')

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

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

    return parser.parse_args()


def upgrade(options, commands, remove_list):
    upgrader = Upgrader(options, commands, remove_list)
    upgrader.run()


def prepare_upgrade(options):

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

    remove_list = []

    setup_logger(options)

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

        # Remember to remove it
        remove_list.append(in_progress)

    upgrade(options, commands, remove_list)


def main():
    options = parse_args()

    if options.reboot_delay < 0:
        sys.exit('ERROR: must specify a positive number')

    if not os.path.exists(options.root_dir):
        sys.exit('ERROR: root directory does not exist: {}'
                 .format(options.root))

    if options.check_reboot:
        # check options must be inert
        options.dry_run = True

    if options.dry_run or options.root_dir != '/':
        prepare_upgrade(options)
        sys.exit(0)

    if options.show_other_details:
        upgrader = Upgrader(options, [], None)
        upgrader.show_other_partition_details()
        sys.exit(0)

    if not options.cmdfile:
        sys.exit('ERROR: need command file')

    if not os.path.exists(options.cmdfile):
        sys.exit('ERROR: command file does not exist: {}'
                 .format(options.cmdfile))

    if os.getuid() != 0:
        sys.exit("ERROR: need to be root")

    pid = os.fork()
    if pid == 0:
        # Writable changes are possible on the read-only root filesystem since
        # the initramfs creates a read-only bind mount *on top* of the real
        # read-only root filesystem mount.
        #
        # This bind mount acts as an insulating layer/cloak, hiding any
        # superblock changes (as made by this script) from the main system.
        # equivalent to "unshare --mount".
        unshare(UnshareFlags.CLONE_NEWNS | UnshareFlags.CLONE_NEWPID)

        # we need to fork again after CLONE_NEWPID, otherwise we get a
        # out of memory error
        pid = os.fork()
        if pid == 0:
            # Peel back the read-only bind mount layer added by the
            # initramfs to allow changes to be made to the underlying
            # filesystem.
            print("I: preparing root filesystem")
            subprocess.check_call(["umount", "/"])

            # Make root filesystem writable.
            #
            # This change is invisible to the rest of the system, since it is
            # operating in the _original_ mount namespace (where the read-only
            # bind mount on top of the real read-only mount below it masks this
            # change).
            subprocess.check_call(["mount", "-o", "remount,rw", "/"])

            # Check the newly-downloaded image, then unpack it on top of
            # the root filesystem.
            prepare_upgrade(options)

            # available in py3.3
            os.sync()

            # Reinstate read-only rootfs in this mount namespace.
            #
            # Note that this operation appears to be redundant (since we
            # are about to destroy the new mount namespace), but it
            # seems safer to restore the original "mount stack" since
            # the live system is relying on it. It certainly shouldn't
            # cause any problems being tidy :)
            print("I: finalising")
            subprocess.check_call(["mount", "-o", "remount,ro", "/"])

            sys.exit(0)

        # second fork
        (pid, status) = os.waitpid(pid, 0)
        exit_code = os.WEXITSTATUS(status)
        sys.exit(exit_code)
    # first fork
    (pid, status) = os.waitpid(pid, 0)
    exit_code = os.WEXITSTATUS(status)
    sys.exit(exit_code)

if __name__ == '__main__':
    main()
