~cprov/tanuki-agent/doc-fix

#!/usr/bin/env python3

# © Canonical 2015
# Licensed under AGPLv3

"""
Tanuki Agent

Usage:
  agent.py [-vq] [--once] [--drain] [--register] [--debug] <configfile>
  agent.py [-vq] [--debug] --local-test <test_opportunity_json_file> <configfile>
  agent.py --default-config [<configfile>]
  agent.py --interactive-login [<configfile>]

Options:
  -h --help             Show this screen
  --version             Show version
  --default-config      Show default configuration.
                        If configfile is given, write default configuration to file.
  --once                Check for tests once and stop
  --drain               Check for tests and stop once the queue is empty.
  --register            Register with SPPI and exit immediately
  --debug               Run tests in a "debug" working dir, do not remove
                        it when done, implies -v
  -v                    Verbose
  -q                    Quiet
  --interactive-login   Get and display SSO credentials.
                        If configfile is given, write SSO credentials to file.
  --local-test          Process a test from a local test opportunity json file instead
                        of polling the server.
"""

from configparser import ConfigParser
import codecs
import datetime
import getpass
import json
import logging
import os
import shutil
import shlex
import subprocess  # NOQA
import sys
import time
import urllib.parse

import docopt
import requests
import requests_oauthlib

import version_info

# Uncomment this to debug OAuth issues
# logging.getLogger("requests_oauthlib.oauth1_auth").setLevel(logging.DEBUG)

DEFAULT_CONFIG = {
    'server': 'https://spi.canonical.com/',
    'id': 'agent-queue',
    'interval': 10,
    'logging_basedir': os.path.join(os.path.dirname(__file__), 'logs'),
    'execution_basedir': os.path.join(os.path.dirname(__file__), 'run'),
    'network_timeout': 10,
}

DEFAULT_COMMANDS = {
    'provcommand': "/bin/echo provisioning now",
    'setupcommand': "/bin/echo setting up now",
    'testcommand': "/bin/echo -e '{\"foo\": \"bar\"}' > spi_test_result.json \\\c",
}

if 'BuiltinImporter' in str(__loader__):
    nuitka = True
else:
    nuitka = False


class Agent(object):

    def __init__(self, arguments):
        # Parse CLI
        self.config = ConfigParser()
        self.config.optionxform = str
        self.confpath = arguments['<configfile>']
        self.register_only = arguments['--register']
        self.once = arguments['--once'] or self.register_only
        self.drain = arguments['--drain']
        self.debug_mode = arguments['--debug']

        self.local_test = arguments['--local-test']
        if self.local_test:
            self.local_test_file = arguments['<test_opportunity_json_file>']
            self.once = True

        # Load config
        if self.confpath is not None:
            self.config.read(self.confpath)
        if 'config' not in self.config.sections():
            self.config['config'] = {}

        self.id = self.config.get('config', 'id', fallback=DEFAULT_CONFIG['id'])
        self.log_dir = self.config.get(
            'config', 'logging_basedir', fallback=DEFAULT_CONFIG['logging_basedir'])

        # Setup logging
        logging.basicConfig(format='%(asctime)s [%(name)s] %(message)s')
        self.log = logging.getLogger(self.id)
        if nuitka:
            self.log.exception = self.log.error
        if arguments['-v'] or arguments['--debug']:
            self.log.setLevel(logging.DEBUG)
        elif arguments['-q']:
            self.log.setLevel(logging.WARNING)
        else:
            self.log.setLevel(logging.INFO)
        makedirs(self.log_dir)
        fh = logging.FileHandler(os.path.join(self.log_dir, 'agent.log'))
        self.log.addHandler(fh)

        # Read the rest of the configuration

        self.execution_basedir = self.config.get(
            'config', 'execution_basedir', fallback=DEFAULT_CONFIG['execution_basedir'])
        self.interval = self.config.getfloat(
            'config', 'interval', fallback=DEFAULT_CONFIG['interval'])
        self.nw_timeout = self.config.getfloat(
            'config', 'network_timeout', fallback=DEFAULT_CONFIG['network_timeout'])
        self.server = self.config.get('config', 'server', fallback=DEFAULT_CONFIG['server'])

        if not self.config.has_section('provkit'):
            self.log.warning('Deprecation: The provisioning kit commands '
                             'should be moved to the [provkit] section.')

        self.testcommand = self.config.get(
            'provkit', 'testcommand', fallback=self.config.get(
                'config', 'testcommand', fallback=DEFAULT_COMMANDS['testcommand']))
        self.setupcommand = self.config.get(
            'provkit', 'setupcommand', fallback=self.config.get(
                'config', 'setupcommand', fallback=DEFAULT_COMMANDS['setupcommand']))
        self.provcommand = self.config.get(
            'provkit', 'provcommand', fallback=self.config.get(
                'config', 'provcommand', fallback=DEFAULT_COMMANDS['provcommand']))

        if 'auth' in self.config:  # Has authentication configured
            self.auth_data = {
                'consumer_key': self.config.get('auth', 'consumer_key'),
                'consumer_secret': self.config.get('auth', 'consumer_secret'),
                'token_key': self.config.get('auth', 'token_key'),
                'token_secret': self.config.get('auth', 'token_secret'),
            }
            self.session = requests_oauthlib.OAuth1Session(
                self.auth_data['consumer_key'],
                client_secret=self.auth_data['consumer_secret'],
                resource_owner_key=self.auth_data['token_key'],
                resource_owner_secret=self.auth_data['token_secret'],
            )
            self.log.debug('Using oauth, consumer_key=%s', self.auth_data['consumer_key'])
        else:
            self.log.warn('Authentication not used')
            self.session = requests.Session()

        self.session.headers.update({'X-Bzr-Revision-Number': version_info.revno})
        self.timestamps = {}

        register_path = "labs/{}".format(self.id)
        self.register_url = urllib.parse.urljoin(self.server, register_path)
        self.testresult_url = None
        self.queue_url = None

        if self.confpath and not os.path.isfile(self.confpath):
            self.log.info('Creating config file: %s', self.confpath)
            self.do_default_config()  # Store default config
        self.warn_empty_config_values()

    def warn_empty_config_values(self):
        """Check configuration for empty values."""
        for section in self.config:
            for key in self.config[section]:
                if self.config.get(section, key) == '':
                    self.log.warn('Key %s in section [%s] is set to empty.', key, section)

    def do_default_config(self):
        """Print default configuration."""
        self.config = ConfigParser()
        self.config['config'] = DEFAULT_CONFIG
        self.config['environment'] = {}
        self.config['provkit'] = DEFAULT_COMMANDS
        if self.confpath is None:
            self.config.write(sys.stdout)
        else:
            with open(self.confpath, 'w+', encoding='utf8') as fp:
                self.config.write(fp)

    def do_interactive_login(self):
        """Login to SOO and produce token."""

        base_sso_url = os.environ.get('SSO_BASE_URL',
                                      'https://login.ubuntu.com')

        def _real_login(email, password, otp=None):
            sso_url = urllib.parse.urljoin(base_sso_url, '/api/v2/tokens/oauth')
            payload = {
                'email': email,
                'password': password,
                'token_name': self.id,
            }
            if otp is not None:
                payload['otp'] = otp
            res = requests.post(
                sso_url,
                json=payload,
                timeout=self.nw_timeout,
            )

            if res.status_code == 401 and res.json().get('code') == 'TWOFACTOR_REQUIRED':
                return _real_login(email, password, input('Enter 2-factor authentication code: '))
            if res.status_code > 299:
                msg = res.json()['message']
                formatted_extra = ';'.join('{}: {}'.format(k, ','.join(v))
                                           for k, v in res.json().get('extra', {}).items())
                self.log.error(
                    'Error %d getting SSO token %s: %s, %s',
                    res.status_code,
                    sso_url,
                    msg, formatted_extra)
                return {}
            return res.json()

        print('Logging into {}\n'.format(base_sso_url))
        email = input('Enter email address: ')
        password = getpass.getpass('Enter password: ')
        res = _real_login(email, password)
        if not res:
            return 1
        auth_data = {
            'consumer_key': res['consumer_key'],
            'consumer_secret': res['consumer_secret'],
            'token_key': res['token_key'],
            'token_secret': res['token_secret'],
        }
        if self.confpath:
            self.config['auth'] = auth_data
            if not self.config['config']:
                self.config['config'] = DEFAULT_CONFIG
            with open(self.confpath, 'w+', encoding='utf8') as fp:
                self.config.write(fp)
        else:
            print('Add this to your configuration file:\n\n')
            print('[auth]')
            print('consumer_key = {}'.format(res['consumer_key']))
            print('consumer_secret = {}'.format(res['consumer_secret']))
            print('token_key = {}'.format(res['token_key']))
            print('token_secret = {}'.format(res['token_secret']))
        return 0

    def register(self):
        """Register with the Message Hub."""
        payload = {
            'agent_id': self.id,
        }

        try:
            r = self.retry_it(
                self.session.post,
                self.register_url,
                json=payload,
                timeout=self.nw_timeout)
            if r.status_code > 299:
                self.log.error(
                    'Error %d registering at %s: %s',
                    r.status_code,
                    self.register_url,
                    r.text)
                return False
            self.log.info('Registered at %s', self.register_url)
            self.log.info('Available products: %s', r.json().get('packages', []))
            try:
                href = r.json()['_links']['dequeue']['href']
                self.queue_url = urllib.parse.urljoin(self.server, href)
                return True
            except KeyError:
                self.log.error("Malformed registration response: %s", r.json())
                return False
        except requests.exceptions.ConnectionError as err:
            self.log.error('Exception: %s', err)
            return False

    def poll(self):
        """Get available tests from Message Hub."""
        if self.queue_url is None:
            self.log.error('Polling without registering correctly')
            exit(1)
        try:
            r = self.retry_it(
                self.session.post,
                self.queue_url,
                timeout=self.nw_timeout)
            if r.status_code == 200 and not r.json():
                # No tests waiting
                self.log.debug('No tests at %s', self.queue_url)
                return False
            elif r.status_code > 299:
                self.log.error(
                    'Error %d polling at %s: %s',
                    r.status_code,
                    self.queue_url,
                    r.text)
                return False
            # Have a test!
            test_data = r.json()
            self.log.info('Found tests at %s: %s', self.queue_url, test_data)
            try:
                links = test_data.pop('_links')
                self.testresult_url = urllib.parse.urljoin(
                    self.server, links['result']['href'])
                return test_data
            except KeyError:
                self.testresult_url = None
                self.log.error("Malformed poll response: %s", test_data)
                return False
        except requests.exceptions.ConnectionError as err:
            self.log.error('Exception: %s', err)
            return False

    def dispatch(self, test_path, command, log_path, working_dir):
        """Dispatch test opportunity to a command, logs in log_path.

        * stderr and stdout will be JOINED in log_path.
        * test_path is passed as argument to command.
        * command is executed  with CWD set to working_dir
        """
        with open(log_path, 'wb+') as out_log:
            cmdline = ' '.join([command, shlex.quote(test_path)])
            self.log.info('Running "%s"', cmdline)
            env = os.environ.copy()
            # remove these, as set by nuitka they confuse descendant pythons
            env.pop('PYTHONHOME', None)
            env.pop('PYTHONPATH', None)
            # add agent config file, used by provkit to get creds for private packages
            env['AGENT_CONFIG_FILE'] = os.path.abspath(self.confpath)
            if 'environment' in self.config:
                env.update(self.config['environment'])
            proc = subprocess.Popen(
                cmdline,
                stdin=1,  # have to use the FD for stdin
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                shell=True,
                env=env,
                cwd=working_dir)

            # Not using communicate() because of python docs warning
            while proc.poll() is None:
                line = proc.stdout.readline()
                if line:
                    out_log.write(line)
                    print('> ', line.decode('utf-8'), end='')
            rest = proc.stdout.read()
            if rest:
                out_log.write(rest)
                for line in rest.decode('utf-8').splitlines():
                    print('> ', line)
            if proc.returncode == 0:
                return True
            else:
                self.log.error('Error calling %s: %d', command, proc.returncode)
                return False

    def post_result(self, result):
        """Send the result of a test to the message hub."""
        self.log.info('Dispatching result: %r', result)
        if self.testresult_url is None:
            self.log.error('post_result called without testresult_url set')
            return False
        try:
            r = self.retry_it(
                self.session.post,
                self.testresult_url,
                json=result,
                timeout=self.nw_timeout)
            if r.status_code > 299:
                self.log.error(
                    'Error %d posting result at %s: %s',
                    r.status_code,
                    self.testresult_url,
                    r.text)
                return False
            self.log.info('Posted result at %s', self.testresult_url)
            return True
        except requests.exceptions.ConnectionError as err:
            self.log.error('Exception: %s', err)
            return False

    def process_test(self, test):
        """Process a single test opportunity, return a test result."""
        test_id = '{}_{}'.format(
            test['id'], datetime.datetime.utcnow().isoformat())

        try:
            if self.debug_mode:
                working_dir = os.path.join(self.execution_basedir, 'debug')
            else:
                working_dir = os.path.join(self.execution_basedir, test_id)
            self.log.info('Running tests in %s', working_dir)
            test_logdir = os.path.join(working_dir, 'logs')
            makedirs(working_dir)
            makedirs(test_logdir)
        except Exception as err:
            self.log.exception('Error creating working directories: %s', err)
            return None

        prov_log = os.path.join(test_logdir, 'provisioning.log')
        setup_log = os.path.join(test_logdir, 'setup.log')
        test_log = os.path.join(test_logdir, 'testing.log')

        # Store test opportunity in known location.
        test_path = 'spi_test_opportunity.json'
        try:
            with open(os.path.join(working_dir, test_path),
                      'w+', encoding='utf-8') as outf:
                json.dump(test, outf)
        except Exception as err:
            self.log.exception('Error saving test opportunity: %s', err)
            return None

        # Run everything.
        test_status = 'PROVFAIL'
        self.timestamps[
            'provisioning_start'] = datetime.datetime.utcnow().isoformat()
        if self.dispatch(test_path, self.provcommand, prov_log, working_dir):
            self.timestamps[
                'setup_start'] = datetime.datetime.utcnow().isoformat()
            test_status = 'SETUPFAIL'
            if self.dispatch(test_path, self.setupcommand,
                             setup_log, working_dir):
                self.timestamps[
                    'testing_start'] = datetime.datetime.utcnow().isoformat()
                test_status = 'FAILED'
                if self.dispatch(test_path, self.testcommand,
                                 test_log, working_dir):
                    test_status = 'PASSED'
        self.timestamps['finished'] = datetime.datetime.utcnow().isoformat()
        # Load results provided by the test runner
        payload = {}
        try:
            result_path = os.path.join(working_dir, 'spi_test_result.json')
            with open(result_path, 'r+', encoding='utf-8') as inf:
                payload = json.load(inf)
        except Exception as err:
            self.log.exception('Error loading test results: %s', err)

        result = {
            'test_opportunity_id': test['id'],
            'agent_id': self.id,
            'test_status': test_status,
            'result_payload': payload,
            'primary_snap_name': test['product'].get('primary_snap_name'),
            'agent': {
                'id': self.id,
                'timing': self.timestamps,
            }
        }

        _result = test.copy()
        _result.pop('id')
        _result.update(result)

        # The working directory is transient, the provisioning kit
        # will have saved any interesting artifacts / logs
        if not self.debug_mode:
            try:
                shutil.rmtree(working_dir)
            except Exception as err:
                self.log.exception('Exception: %s', err)
        return _result

    def start(self):
        """Poll the message hub, and dispatch tests.

        Return False if there was an error or no
        tests were available.
        """

        if not self.queue_url and not self.local_test:
            if not self.register():
                return False
        if self.register_only:
            return True
        if self.local_test:
            self.log.info('local_test: %r', self.local_test_file)
            with codecs.open(self.local_test_file, 'r', encoding='utf8') as fh:
                test = json.load(fh)
            self.log.info('test_opp: %r', test)
        else:
            test = self.poll()
        if test:
            result = self.process_test(test)
            if result:
                if self.local_test:
                    pass  # TODO: write out
                else:
                    for attempt in range(3):
                        if self.post_result(result):
                            return True
                        time.sleep(self.interval)
        return False

    def main(self):
        """Run start and handle exceptions."""
        try:
            wait = not self.start()
        except KeyboardInterrupt:
            self.log.info('Quitting (keyboard interrupt)')
            exit(0)
        except Exception as err:
            self.log.exception('Exception: %s', err)
            wait = True
        if wait and self.drain:
            exit(0)
        if wait and not self.once:
            self.log.debug('Sleeping %f seconds', self.interval)
            time.sleep(self.interval)

    def retry_it(self, func, *a, **kw):
        for i in range(3):
            try:
                return func(*a, **kw)
            except:
                if i < 2:
                    self.log.warn('Retrying...')
                else:
                    raise


def makedirs(path):
    """Create a folder."""
    if os.path.exists(path):
        if not os.path.isdir(path):
            raise OSError(
                'Path {0} already exists and is not a folder.'.format(path))
        else:
            return
    os.makedirs(path)


if __name__ == "__main__":
    vstr = 'SPI Agent r{}'.format(version_info.revno)
    arguments = docopt.docopt(__doc__, version=vstr)
    agent = Agent(arguments)
    agent.log.info(vstr)
    if arguments['--interactive-login']:
        exit(agent.do_interactive_login())
    elif arguments['--default-config']:
        agent.do_default_config()
        exit(0)
    else:
        while True:
            agent.main()
            if agent.once:
                break