osc

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright © 2015  Daniel Dehennin <daniel.dehennin@baby-gnu.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""OpenNebula Sunstone REST API Client

"""

import gevent
import gevent.socket
import gevent.subprocess
from gevent.socket import AF_UNIX
from gevent.socket import SOCK_STREAM

# Make standard library gevent friendly
from gevent import monkey
monkey.patch_all()

import sys

import requests
from bs4 import BeautifulSoup

from websocket import create_connection
from urllib.parse import urlparse

from os import access
from os import R_OK
from os import unlink
from os import environ

from os.path import join

from argparse import ArgumentParser
from argparse import RawDescriptionHelpFormatter

import re

import logging

VM_STATE = {'0': 'INIT',
            '1': 'PENDING',
            '2': 'HOLD',
            '3': 'ACTIVE',
            '4': 'STOPPED',
            '5': 'SUSPENDED',
            '6': 'DONE',
            '7': 'SHOULD NOT EXIST',
            '8': 'POWEROFF',
            '9': 'UNDEPLOYED'}

VM_LCM_STATE = {'0': 'LCM_INIT',
                '1': 'PROLOG',
                '2': 'BOOT',
                '3': 'RUNNING',
                '4': 'MIGRATE',
                '5': 'SAVE_STOP',
                '6': 'SAVE_SUSPEND',
                '7': 'SAVE_MIGRATE',
                '8': 'PROLOG_MIGRATE',
                '9': 'PROLOG_RESUME',
                '10': 'EPILOG_STOP',
                '11': 'EPILOG',
                '12': 'SHUTDOWN',
                '15': 'CLEANUP_RESUBMIT',
                '16': 'UNKNOWN',
                '17': 'HOTPLUG',
                '18': 'SHUTDOWN_POWEROFF',
                '19': 'BOOT_UNKNOWN',
                '20': 'BOOT_POWEROFF',
                '21': 'BOOT_SUSPENDED',
                '22': 'BOOT_STOPPED',
                '23': 'CLEANUP_DELETE',
                '24': 'HOTPLUG_SNAPSHOT',
                '25': 'HOTPLUG_NIC',
                '26': 'HOTPLUG_SAVEAS',
                '27': 'HOTPLUG_SAVEAS_POWEROFF',
                '28': 'HOTPLUG_SAVEAS_SUSPENDED',
                '29': 'SHUTDOWN_UNDEPLOY',
                '30': 'EPILOG_UNDEPLOY',
                '31': 'PROLOG_UNDEPLOY',
                '32': 'BOOT_UNDEPLOY',
                '33': 'HOTPLUG_PROLOG_POWEROFF',
                '34': 'HOTPLUG_EPILOG_POWEROFF',
                '35': 'BOOT_MIGRATE',
                '36': 'BOOT_FAILURE',
                '37': 'BOOT_MIGRATE_FAILURE',
                '38': 'PROLOG_MIGRATE_FAILURE',
                '39': 'PROLOG_FAILURE',
                '40': 'EPILOG_FAILURE',
                '41': 'EPILOG_STOP_FAILURE',
                '42': 'EPILOG_UNDEPLOY_FAILURE',
                '43': 'PROLOG_MIGRATE_POWEROFF',
                '44': 'PROLOG_MIGRATE_POWEROFF_FAILURE',
                '45': 'PROLOG_MIGRATE_SUSPEND',
                '46': 'PROLOG_MIGRATE_SUSPEND_FAILURE',
                '47': 'BOOT_UNDEPLOY_FAILURE',
                '48': 'BOOT_STOPPED_FAILURE',
                '49': 'PROLOG_RESUME_FAILURE',
                '50': 'PROLOG_UNDEPLOY_FAILURE',
                '51': 'DISK_SNAPSHOT_POWEROFF',
                '52': 'DISK_SNAPSHOT_REVERT_POWEROFF',
                '53': 'DISK_SNAPSHOT_DELETE_POWEROFF',
                '54': 'DISK_SNAPSHOT_SUSPENDED',
                '55': 'DISK_SNAPSHOT_REVERT_SUSPENDED',
                '56': 'DISK_SNAPSHOT_DELETE_SUSPENDED',
                '57': 'DISK_SNAPSHOT',
                '58': 'DISK_SNAPSHOT_REVERT',
                '59': 'DISK_SNAPSHOT_DELETE'}

POOL_FILTER = {
    'INFO_GROUP'         : -1,
    'INFO_ALL'           : -2,
    'INFO_MINE'          : -3,
    'INFO_PRIMARY_GROUP' : -4}

def main():
    opts = parse_cmdline()

    try:
        main_logger = init_logging('OpenNebula Sunstone REST API Client', opts)
        osc_logger = init_logging('ONESunstoneClient', opts)

        get_one_credentials(opts)
        get_one_sunstone_url(opts)

        osc = ONESunstoneClient(opts)

        osc.loop()

    except Exception as err:
        main_logger.error(f'{err}',
                          exc_info=main_logger.isEnabledFor(logging.DEBUG))
        sys.exit(1)


def parse_cmdline():
    desc = """Run commands on VMs throught Sunstone REST API.

The authentication is performed with ~/.one/one_auth file as describe in the documentation[1].

The Sunstone URL is retrieve in the following order:

1. from command line parameter --url
2. from environment variable ONE_SUNSTONE

[1] http://docs.opennebula.org/stable/administration/users_and_groups/manage_users.html#shell-environment
"""

    parser = ArgumentParser(description='OpenNebula Sunstone REST API Client',
                            epilog=desc,
                            formatter_class=RawDescriptionHelpFormatter)

    parser.add_argument('--vm', metavar='VMID',
                        help='ID of a VM')

    parser.add_argument('--startvnc', action='store_true',
                        help='Start VNC client ssvncviewer')

    parser.add_argument('--url', help='URL to OpenNebula Sunstone frontend')

    parser.add_argument('--ws-port', default='29876', help='Websocket port')

    parser.add_argument('--ws-path', default='/', help='Websocket PATH')

    parser.add_argument('--auth', help='ONE_AUTH file')

    group = parser.add_argument_group('logging')

    group.add_argument('-l', '--log-level',
                       choices=['debug', 'info', 'warning', 'error', 'critical'],
                       default='error', help='Log level')

    group.add_argument('-v', '--verbose',
                       action='store_true',
                       help='Verbose mode, equivalent to -l info')

    group.add_argument('-d', '--debug',
                       action='store_true',
                       help='Debug mode, equivalent to -l debug')

    opts = parser.parse_args()

    if opts.verbose:
        opts.log_level = 'info'
    if opts.debug:
        opts.log_level = 'debug'
                    
    opts.log_level = getattr(logging, opts.log_level.upper())

    return opts


def init_logging(name, opts):
    """Initialize logging

    """
    logger = logging.getLogger(name)
    logger.setLevel(opts.log_level)
    log_handler = logging.StreamHandler()
    log_handler.setLevel(opts.log_level)
    if logger.isEnabledFor(logging.DEBUG):
        log_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    else:
        log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    log_handler.setFormatter(log_formatter)
    logger.addHandler(log_handler)
    return logger


def get_one_credentials(opts):
    """Get user and password credential

    """
    user = None
    password = None
    one_auth = join(environ['HOME'], '.one', 'one_auth')

    if opts.auth is not None and access(opts.auth, R_OK):
        one_auth = opts.auth
    elif 'ONE_AUTH' in environ:
        one_auth = environ['ONE_AUTH']

    with open(one_auth, 'r') as one_auth_fh:
        user, password = one_auth_fh.readline().strip().split(':')

    opts.creds = (user, password)


def get_one_sunstone_url(opts):
    """Get the URL to OpenNebula Sunstone frontend

    """
    url = None
    if opts.url is not None:
        url = opts.url
    elif 'ONE_SUNSTONE' in environ:
        url = environ['ONE_SUNSTONE']
    else:
        raise SystemError('No Sunstone URL provided')

    opts.url = url


class ONESunstoneClient(object):
    """OpenNebula Sunstone Client object

    """

    csrftoken_regex = re.compile(r'''csrftoken\s*=\s*['"]([^'"]+)['"]''')

    def __init__(self, opts):
        """Initialize the Sunstone client

        """
        self.opts = opts
        self.logger = logging.getLogger('ONESunstoneClient')
        self.logger.setLevel(opts.log_level)

        self.client = requests.session()
        self.client_opts = {'headers': {'Referer': self.opts.url},
                            'csrftoken': None}

        self.logged = False
        self.threads = []
        self.websocket = None
        self.connection = None
        self.socket = None
        self.socket_path = None


    def loop(self):
        """OpenNebula Sunstone client loop

        """
        self.logger.debug("Start main loop")
        self.login()
        vm = self.ask_vm()

        action = self.ask_action(vm)
        if action:
            getattr(self, action)(vm)

        # Clean command line arguments
        self.opts.vm = None
        self.opts.startvnc = False
            
        return True


    def ask_vm(self):
        """Display OpenNebula Sunstone client menu

        """
        self.logger.debug('Ask VM to user')
        if self.opts.vm is not None:
            try:
                vm = self.get_vm(self.opts.vm)
                return vm
            except SystemError as err:
                print(err)

        vms = self.get_vm()
        if not vms:
            raise SystemError('No virtual machine are running, nothing to do')

        vm_info = {}

        print('List of current VMs:')
        for vm in vms:
            vmid = vm['ID']
            vm_info[vmid] = vm
            state = VM_STATE[vm['STATE']]
            lcm_state = VM_LCM_STATE[vm['LCM_STATE']]
            if state != 'ACTIVE':
                vm_msg = f'{vmid} - {vm["NAME"]} ({state})'
            else:
                vm_msg = f'{vmid} - {vm["NAME"]} ({state}/{lcm_state})'

            print(vm_msg)

        print('Select a virtual machine id or q|quit to exit')
        choice = input('[q]: ')

        vm_ids = vm_info.keys()
        if choice == '':
            choice = 'q'
        elif choice in vm_ids:
            return vm_info[choice]
        elif re.match(r'^q(uit)?', choice):
            print('Voluntary stay of proceedings')
            sys.exit(0)
        else:
            print(f'Invalid response: {choice}')
            return False


    def ask_action(self, vm):
        """Ask action to user

        """
        self.logger.debug('Ask VM action to user')
        actions = ['startvnc']

        if self.opts.startvnc:
            return actions[0]

        if vm['STATE'] != '3' or vm['LCM_STATE'] != '3':
            actions.remove('startvnc')

        if len(actions) == 0:
            vmid = vm['ID']
            state = VM_STATE[vm['STATE']]
            lcm_state = VM_LCM_STATE[vm['LCM_STATE']]
            if state != 'ACTIVE':
                vm_msg = f'{vmid} - {vm["NAME"]} ({state})'
            else:
                vm_msg = f'{vmid} - {vm["NAME"]} ({state}/{lcm_state})'

            msg = f'No action available for {vm_msg}'
            raise SystemError(msg)

        print('List of available actions:')
        for aid, action in enumerate(actions):
            print(f'{aid} - {action}')

        msg = f'Select an action to perform on {vm["ID"]} - {vm["NAME"]}, or q|quit to exit'
        print(msg)

        choice = input('[startvnc]: ')

        if choice == '':
            choice = 'startvnc'

        try:
            choice = int(choice)
            return actions[choice]
        except ValueError:
            if re.match(r'^q(uit)?', choice):
                print('Voluntary stay of proceedings')
                sys.exit(0)
            elif choice in actions:
                return choice
            else:
                raise ValueError(f'Invalid response: {choice}')

    def _find_csrftoken(self, lines):
        """Find the csrftoken in the text

        """
        csrftoken = None

        for line in lines:
            m = self.csrftoken_regex.search(line)
            if m is not None:
                csrftoken = m.group(1)
                break

        return csrftoken


    def login(self):
        """Log in OpenNebula Sunstone

        """
        user = self.opts.creds[0]
        url = self.opts.url

        self.logger.info(f'Login as {user} to {url}')

        login = self.client.post(f'{url}/login',
                                 auth=self.opts.creds,
                                 headers=self.client_opts['headers'])

        if not login.ok:
            reason = login.reason
            raise SystemError(f'Unable to login as {user}: {reason}')
        else:
            self.logger.debug('Loging ok')

        root = self.client.get(f'{url}/',
                               headers=self.client_opts['headers'])

        if not root.ok:
            raise SystemError("Unable to get root page after login")

        soup = BeautifulSoup(root.content, 'html.parser')

        csrftoken = None
        for script in soup.find_all('script'):
            csrftoken = self._find_csrftoken(script)
            if csrftoken is not None:
                break

        self.client_opts['csrftoken'] = {'csrftoken': csrftoken}
        self.logged = True


    def get_vm(self, vmid=None):
        """Get a VM or list of VMs

        """
        if vmid is not None:
            msg = f'Get informations of vm {vmid}'
        else:
            msg = 'Get informations of all user VMs'

        self.logger.info(msg)

        if not self.logged:
            self.login()

        vm_url = f'{self.opts.url}/vm'
        params = {}
        params.update(self.client_opts['csrftoken'])

        if vmid is not None:
            vm_url += '/{vmid}'.format(vmid=vmid)
        else:
            params['timeout'] = False
            params['pool_filter'] = POOL_FILTER['INFO_ALL']
            
        reply = self.client.get(vm_url, params=params)
        if not reply.ok:
            reason = reply.reason
            if vmid is not None:
                msg = f'Unable to get virtual machine {vmid}: {reason}'
            else:
                msg = f'Unable to get virtual machines: {reason}'

            raise SystemError(msg)

        if vmid is None:
            return reply.json()['VM_POOL']['VM']
        else:
            return reply.json()['VM']


    def startvnc(self, vm):
        """Start a VNC client for a VM

        """
        vmid = vm['ID']

        self.logger.info(f'Start VNC client for {vmid}')

        startvnc_url = f'{self.opts.url}/vm/{vmid}/startvnc'
        reply = self.client.post(startvnc_url,
                                 params=self.client_opts['csrftoken'])

        vnc = reply.json()

        vnc_host = urlparse(startvnc_url).netloc
        ws_port = self.opts.ws_port
        ws_path = self.opts.ws_path
        ws_url = f'ws://{vnc_host}:{ws_port}{ws_path}?token={vnc["token"]}'

        socket_name = f'one-sunstone-client-{vnc["token"]}.socket'
        self.socket_path = join('/tmp', socket_name)

        if access(self.socket_path, R_OK):
            unlink(self.socket_path)

        try:
            self.socket = gevent.socket.socket(AF_UNIX, SOCK_STREAM)
            self.socket.bind(self.socket_path)
            self.socket.listen(1)
        except gevent.socket.error as err:
            print(f'Failed to create socket. Error code: {err[0]}, Error message : {err[1]}')
            sys.exit(1)

        self.logger.debug(f'Create a websocket connection to {ws_url}')
        self.websocket = create_connection(ws_url,
                                           header=[f'Origin: {self.opts.url}',
                                                   'Sec-WebSocket-Protocol: binary'])

        self.threads = [gevent.spawn(self._from_ws),
                        gevent.spawn(self._from_socket),
                        gevent.spawn(self._start_viewer)]

        gevent.joinall(self.threads)


    def terminate(self):
        """Terminate everything

        """
        self.logger.debug('Terminate gevent threads')

        if self.websocket.connected:
            self.websocket.close()

        if not self.socket.closed:
            self.socket.close()

        if access(self.socket_path, R_OK):
            unlink(self.socket_path)

        for thread in self.threads:
            gevent.kill(thread)


    def _start_viewer(self):
        try:
            cmd = ['ssvncviewer', self.socket_path]

            self.logger.debug(f'Start VNC viewer: {" ".join(cmd)}')

            gevent.subprocess.call(cmd)
        except Exception as err:
            print(f'Fail to run ssvncviewer: {err}')

        self.terminate()


    def _from_ws(self):
        self.logger.debug('Start passing data from WebSocket to VNC client')
        while self.websocket.connected and not self.socket.closed:
            if self.connection is None:
                gevent.sleep(2)
                continue

            data = self.websocket.recv()
            try:
                self.connection.sendall(data)
            except gevent.socket.error as err:
                print(f'Unable to write to local socket: {err}')
                self.terminate()


    def _from_socket(self):
        self.logger.debug('Start passing data from VNC client to WebSocket')
        conn, _ = self.socket.accept()
        self.connection = conn
        while not self.connection.closed and self.websocket.connected:

            data = self.connection.recv(4096)
            try:
                self.websocket.send(data)
            except gevent.socket.error as err:
                print(f'Unable to write to WebSocket: {err}')
                self.terminate()


if __name__ == '__main__':
    sys.exit(main())