#!/usr/bin/python
# Copyright (C) 2013-2015  Codethink Limited
#
# 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 2 of the License.
#
# 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/>.


'''A Morph deployment write extension for upgrading systems over ssh.'''


import contextlib
import os
import subprocess
import sys
import tempfile
import time

import writeexts


def ssh_runcmd_ignore_failure(location, command, **kwargs):
    try:
        return writeexts.ssh_runcmd(location, command, **kwargs)
    except writeexts.ExtensionError:
        pass


class SshRsyncWriteExtension(writeexts.WriteExtension):

    '''See ssh-rsync.write.help for documentation'''
    

    def find_root_disk(self, location):
        '''Read /proc/mounts on location to find which device contains "/"'''

        self.status(msg='Finding device that contains "/"')
        contents = writeexts.ssh_runcmd(location,
                                        ['cat', '/proc/mounts'])
        for line in contents.splitlines():
            line_words = line.split()
            if (line_words[1] == '/' and line_words[0] != 'rootfs'):
                return line_words[0]

    @contextlib.contextmanager
    def _remote_mount_point(self, location):
        self.status(msg='Creating remote mount point')
        remote_mnt = writeexts.ssh_runcmd(location,
                                          ['mktemp', '-d']).strip()
        try:
            yield remote_mnt
        finally:
            self.status(msg='Removing remote mount point')
            writeexts.ssh_runcmd(location, ['rmdir', remote_mnt])

    @contextlib.contextmanager
    def _remote_mount(self, location, root_disk, mountpoint):
        self.status(msg='Mounting root disk')
        writeexts.ssh_runcmd(location, ['mount', root_disk, mountpoint])
        try:
            yield
        finally:
            self.status(msg='Unmounting root disk')
            writeexts.ssh_runcmd(location, ['umount', mountpoint])

    @contextlib.contextmanager
    def _created_version_root(self, location, remote_mnt, version_label):
        version_root = os.path.join(remote_mnt, 'systems', version_label)
        self.status(msg='Creating %(root)s', root=version_root)
        writeexts.ssh_runcmd(location, ['mkdir', version_root])
        try:
            yield version_root
        except BaseException as e:
            # catch all, we always want to clean up
            self.status(msg='Cleaning up %(root)s', root=version_root)
            ssh_runcmd_ignore_failure(location, ['rmdir', version_root])
            raise

    def get_old_orig(self, location, remote_mnt):
        '''Identify which subvolume to snapshot from'''

        # rawdisk upgrades use 'default'
        return os.path.join(remote_mnt, 'systems', 'default', 'orig')

    @contextlib.contextmanager
    def _created_orig_subvolume(self, location, remote_mnt, version_root):
        self.status(msg='Creating "orig" subvolume')
        old_orig = self.get_old_orig(location, remote_mnt)
        new_orig = os.path.join(version_root, 'orig')
        writeexts.ssh_runcmd(location, ['btrfs', 'subvolume', 'snapshot',
                                        old_orig, new_orig])
        try:
            yield new_orig
        except BaseException as e:
            ssh_runcmd_ignore_failure(
                    location, ['btrfs', 'subvolume', 'delete', new_orig])
            raise

    def populate_remote_orig(self, location, new_orig, temp_root):
        '''Populate the subvolume version_root/orig on location'''

        self.status(msg='Populating "orig" subvolume')
        subprocess.check_call(['rsync', '-as', '--checksum', '--numeric-ids',
                               '--delete', temp_root + os.path.sep,
                               '%s:%s' % (location, new_orig)])

    @contextlib.contextmanager
    def _deployed_version(self, location, version_label,
                          system_config_sync, system_version_manager):
        self.status(msg='Calling system-version-manager to deploy upgrade')
        deployment = os.path.join('/systems', version_label, 'orig')
        writeexts.ssh_runcmd(location,
                ['env', 'BASEROCK_SYSTEM_CONFIG_SYNC='+system_config_sync,
                 system_version_manager, 'deploy', deployment])
        try:
            yield deployment
        except BaseException as e:
            self.status(msg='Cleaning up failed version installation')
            writeexts.ssh_runcmd(location,
                    [system_version_manager, 'remove', version_label])
            raise

    def upgrade_remote_system(self, location, temp_root):
        root_disk = self.find_root_disk(location)
        uuid = writeexts.ssh_runcmd(location,
            ['blkid', '-s', 'UUID', '-o', 'value', root_disk]).strip()

        self.complete_fstab_for_btrfs_layout(temp_root, uuid)

        version_label = os.environ['VERSION_LABEL']
        autostart = self.get_environment_boolean('AUTOSTART')

        with self._remote_mount_point(location) as remote_mnt, \
             self._remote_mount(location, root_disk, remote_mnt), \
             self._created_version_root(location, remote_mnt,
                                        version_label) as version_root, \
             self._created_orig_subvolume(location, remote_mnt,
                                          version_root) as orig:
                self.populate_remote_orig(location, orig, temp_root)
                system_root = os.path.join(remote_mnt, 'systems',
                                           version_label, 'orig')
                config_sync = os.path.join(system_root, 'usr', 'bin',
                                           'baserock-system-config-sync')
                version_manager = os.path.join(system_root, 'usr', 'bin',
                                               'system-version-manager')
                with self._deployed_version(location, version_label,
                                            config_sync, version_manager):
                    self.status(msg='Setting %(v)s as the new default system',
                                v=version_label)
                    writeexts.ssh_runcmd(location,
                        [version_manager, 'set-default', version_label])

        if autostart:
            self.status(msg="Rebooting into new system ...")
            ssh_runcmd_ignore_failure(location, ['reboot'])

    def process_args(self, args):
        if len(args) != 2:
            raise writeexts.ExtensionError(
                'Wrong number of command line args')

        temp_root, location = args

        self.upgrade_remote_system(location, temp_root)


SshRsyncWriteExtension().run()