Skip to content
Snippets Groups Projects
envsetup.py 7.28 KiB
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
Environment setup script for MediaServer
'''
import importlib.util
import os
import subprocess
import sys
import traceback

import utils
from utils import log


class EnvSetup():
    USAGE = '''%s [-d] [-h] [<action id>]
    -d: debug mode (can be started with non root users).
    -h: show this message.
    action id: specify which action should be started (non interactive).''' % __file__
    PY_SETUP_NAME = '0_setup.py'
    BASH_SETUP_NAME = '0_setup.sh'

    def __init__(self, *args):
        self.display_header()
        args = list(args)
        # Check if help is required
        if '-h' in args:
            log('USAGE: ' + self.USAGE)
            sys.exit(0)
        # Check current dir
        root_dir = utils.get_dir(__file__)
        if root_dir != '':
            os.chdir(root_dir)
        self.root_dir = root_dir
        # Add to python path
        if root_dir not in sys.path:
            sys.path.append(root_dir)
        # Get available actions
        self.actions = self.discover_actions()
        if not self.actions:
            log('No action available.')
            sys.exit(1)
        # Check that this script is run by root
        self.debug = '-d' in args
        if self.debug:
            args.remove('-d')
        whoami = subprocess.check_output(['whoami']).decode('utf-8').strip()
        if whoami != 'root' and not self.debug:
            log('This script should be run as root user.')
            sys.exit(1)
        # Load conf
        conf = utils.load_conf()
        if not conf:
            log('No configuration loaded.')
            sys.exit(1)
        if args:
            # Run command
            for arg in args:
                self.run(arg, interactive=False)
        else:
            # Open main menu
            self.menu()

    def display_header(self):
        log('\033[96m----------------------------------\033[0m')
        log('\033[96m- UbiCast environment setup tool -\033[0m')
        log('\033[96m----------------------------------\033[0m')

    def discover_actions(self):
        actions = list()
        for section_name in os.listdir(self.root_dir):
            section_path = os.path.join(self.root_dir, section_name)
            if not os.path.isdir(section_path):
                continue
            try:
                section_index = int(section_name.split('.')[0])
            except ValueError:
                continue
            section_label = section_name[len(str(section_index)) + 1:].strip().replace('_', ' ')
            if not section_label:
                log('No label found for dir %s.' % section_name)
                continue
            actions.append(dict(index=int(str(section_index) + '0'), label=section_label, path=section_path, fct=None))

            for name in os.listdir(section_path):
                path = os.path.join(section_path, name)
                if not os.path.isdir(path):
                    continue
                try:
                    index = int(str(section_index) + name.split('.')[0])
                except ValueError:
                    continue
                label = name[len(str(index)) + 1 - len(str(section_index)):].strip().replace('_', ' ')
                if not label:
                    log('No label found for dir %s.' % name)
                    continue

                if os.path.isfile(os.path.join(path, self.PY_SETUP_NAME)):
                    spec = importlib.util.spec_from_file_location('setup_%s' % name, os.path.join(path, self.PY_SETUP_NAME))
                    setup_module = importlib.util.module_from_spec(spec)
                    spec.loader.exec_module(setup_module)
                    actions.append(dict(index=index, label=label, path=path, fct=setup_module.setup))
                elif os.path.isfile(os.path.join(path, self.BASH_SETUP_NAME)):
                    actions.append(dict(index=index, label=label, path=path, fct='bash -e "%s"' % os.path.join(path, self.BASH_SETUP_NAME)))
        actions.sort(key=lambda a: a['index'])
        return actions

    def menu(self):
        # Show main menu
        log('Actions:')
        for action in self.actions:
            if action['fct']:
                log('  %s: %s' % (action['index'], action['label']))
            else:
                log('  \033[1;94m%s\033[0m' % (action['label']))
        log('')
        log('  t: Run tests')
        log('  c: Configuration status')
        log('  e: Exit\n')
        log('Info:')
        log('\033[0;36m  To setup a system entirely for a determined purpose (Worker, MS, CM, ...), you should use the launcher:\033[0m')
        log('\033[0;36m  bash /root/envsetup/launcher.sh\033[0m')
        log('\nWhat action do you want to start ?')
        try:
            target = input('---> ').strip()
        except (KeyboardInterrupt, EOFError):
            log('')
            target = 'e'
        self.run(target)

    def run(self, target, interactive=True):
        if target == 'e':
            log('Exit')
            sys.exit(0)
        exit_code = 0
        if target == 't':
            # Run tests
            args = [os.path.join(self.root_dir, 'tester.py'), 'tester.py']
            if self.debug:
                args.append('-d')
            os.execl(*args)
        elif target == 'c':
            # Display current configuration
            log('Configuration status:')
            override = utils.get_conf('_override')
            if not override:
                log('Configuration status not available.')
            else:
                log('Is default | Name | Value')
                for name, is_overriden in override.items():
                    is_default = '\033[93m no ' if is_overriden else '\033[94m yes'
                    log('%s\033[0m | \033[95m%s\033[0m | \033[96m%s\033[0m' % (is_default, name, utils.get_conf(name)))
        else:
            # Run an action
            found = False
            for action in self.actions:
                if target == str(action['index']) and action['fct']:
                    found = True
                    log('Starting action %s: %s setup.' % (action['index'], action['label']))
                    try:
                        os.chdir(action['path'])
                        if isinstance(action['fct'], str):
                            utils.run_commands([action['fct']])
                        else:
                            action['fct'](interactive)
                    except Exception as e:
                        exit_code = 1
                        if isinstance(action['fct'], str):
                            log(action['fct'])
                        else:
                            log(traceback.format_exc())
                        log('Unable to setup %s:\n%s\n' % (action['label'], e))
                    else:
                        log('%s setup complete.\n' % action['label'])
                    os.chdir(self.root_dir)
                    break
            if not found:
                exit_code = 1
                log('Invalid action requested: "%s".' % target)
        if interactive:
            try:
                input('Press enter to continue.')
            except (KeyboardInterrupt, EOFError):
                log('')
                sys.exit(exit_code)
            self.display_header()
            self.menu()
        else:
            sys.exit(exit_code)


if __name__ == '__main__':
    EnvSetup(*sys.argv[1:])