#!/usr/bin/python
# Copyright (C) 2014-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/>.

'''Preparatory checks for Morph 'kvm' write extension'''

import os
import re
import urlparse

import writeexts


class KvmPlusSshCheckExtension(writeexts.WriteExtension):

    location_pattern = '^/(?P<guest>[^/]+)(?P<path>/.+)$'

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

        self.require_btrfs_in_deployment_host_kernel()

        upgrade = self.get_environment_boolean('UPGRADE')
        if upgrade:
            raise writeexts.ExtensionError(
                'Use the `ssh-rsync` write extension to deploy upgrades to an '
                'existing remote system.')

        location = args[0]
        ssh_host, vm_name, vm_path = self.check_and_parse_location(location)

        self.check_ssh_connectivity(ssh_host)
        self.check_can_create_file_at_given_path(ssh_host, vm_path)
        self.check_no_existing_libvirt_vm(ssh_host, vm_name)
        self.check_extra_disks_exist(ssh_host, self.parse_attach_disks())
        self.check_virtual_networks_are_started(ssh_host)
        self.check_host_has_virtinstall(ssh_host)

    def check_and_parse_location(self, location):
        '''Check and parse the location argument to get relevant data.'''

        x = urlparse.urlparse(location)

        if x.scheme != 'kvm+ssh':
            raise writeexts.ExtensionError(
                'URL schema must be kvm+ssh in %s' % location)

        m = re.match(self.location_pattern, x.path)
        if not m:
            raise writeexts.ExtensionError(
                'Cannot parse location %s' % location)

        return x.netloc, m.group('guest'), m.group('path')

    def check_no_existing_libvirt_vm(self, ssh_host, vm_name):
        try:
            writeexts.ssh_runcmd(ssh_host,
                ['virsh', '--connect', 'qemu:///system', 'domstate', vm_name])
        except writeexts.ExtensionError as e:
            pass
        else:
            raise writeexts.ExtensionError(
                'Host %s already has a VM named %s. You can use the ssh-rsync '
                'write extension to deploy upgrades to existing machines.' %
                (ssh_host, vm_name))

    def check_can_create_file_at_given_path(self, ssh_host, vm_path):

        def check_can_write_to_given_path():
            try:
                writeexts.ssh_runcmd(ssh_host, ['touch', vm_path])
            except writeexts.ExtensionError as e:
                raise writeexts.ExtensionError(
                    "Can't write to location %s on %s" % (vm_path, ssh_host))
            else:
                writeexts.ssh_runcmd(ssh_host, ['rm', vm_path])

        try:
            writeexts.ssh_runcmd(ssh_host, ['test', '-e', vm_path])
        except writeexts.ExtensionError as e:
            # vm_path doesn't already exist, so let's test we can write
            check_can_write_to_given_path()
        else:
            raise writeexts.ExtensionError('%s already exists on %s'
                                           % (vm_path, ssh_host))

    def check_extra_disks_exist(self, ssh_host, filename_list):
        for filename in filename_list:
            try:
                writeexts.ssh_runcmd(ssh_host, ['ls', filename])
            except writeexts.ExtensionError as e:
                raise writeexts.ExtensionError(
                    'Did not find file %s on host %s' % (filename, ssh_host))

    def check_virtual_networks_are_started(self, ssh_host):

        def check_virtual_network_is_started(network_name):
            cmd = ['virsh', '-c', 'qemu:///system', 'net-info', network_name]
            net_info = writeexts.ssh_runcmd(ssh_host, cmd).split('\n')

            def pretty_concat(lines):
                return '\n'.join(['\t%s' % line for line in lines])

            for line in net_info:
                m = re.match('^Active:\W*(\w+)\W*', line)
                if m:
                    break
            else:
                raise writeexts.ExtensionError(
                        "Got unexpected output parsing output of `%s':\n%s"
                        % (' '.join(cmd), pretty_concat(net_info)))

            network_active = m.group(1) == 'yes'

            if not network_active:
                raise writeexts.ExtensionError("Network '%s' is not started"
                                               % network_name)

        def name(nic_entry):
            if ',' in nic_entry:
                # network=NETWORK_NAME,mac=12:34,model=e1000...
                return nic_entry[:nic_entry.find(',')].lstrip('network=')
            else:
                return nic_entry.lstrip('network=')    # NETWORK_NAME

        if 'NIC_CONFIG' in os.environ:
            nics = os.environ['NIC_CONFIG'].split()

            for n in nics:
                if not (n.startswith('network=')
                        or n.startswith('bridge=')
                        or n == 'user'):
                    raise writeexts.ExtensionError(
                        "malformed NIC_CONFIG: %s\n"
                        " (expected 'bridge=BRIDGE' 'network=NAME'"
                        " or 'user')" % n)

            # --network bridge= is used to specify a bridge
            # --network user is used to specify a form of NAT
            # (see the virt-install(1) man page)
            networks = [name(n) for n in nics if not n.startswith('bridge=')
                                                and not n.startswith('user')]
        else:
            networks = ['default']

        for network in networks:
            check_virtual_network_is_started(network)

    def check_host_has_virtinstall(self, ssh_host):
        try:
            writeexts.ssh_runcmd(ssh_host, ['which', 'virt-install'])
        except writeexts.ExtensionError:
            raise writeexts.ExtensionError(
                'virt-install does not seem to be installed on host %s'
                % ssh_host)


KvmPlusSshCheckExtension().run()