moulinette/parse_args

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

__credits__ = """
    Copyright (C) 2012 YunoHost

    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
    """
__author__  = 'Kload <kload@kload.fr>'
__version__ = '2.0 beta1'

import os
import sys
import argparse
import gettext
import getpass
try:
    import yaml
except ImportError:
    sys.stderr.write('Error: Yunohost CLI Require yaml lib\n')
    sys.stderr.write('apt-get install python-yaml\n')
    sys.exit(1)
import json
if not __debug__:
    import traceback

gettext.install('YunoHost')

try:
    from yunohost import YunoHostError, YunoHostLDAP, str_to_func, colorize, pretty_print_dict, display_error, validate, win
except ImportError:
    sys.stderr.write('Error: Yunohost CLI Require YunoHost lib\n')
    sys.exit(1)


def parse_dict(action_map):
    """
    Turn action dictionnary to parser, subparsers and arguments

    Keyword arguments:
        action_map -- Multi-level dictionnary of categories/actions/arguments list

    Returns:
        Namespace of args

    """
    # Intialize parsers
    parsers = subparsers_category = subparsers_action = {}
    parsers['general'] = argparse.ArgumentParser()
    subparsers = parsers['general'].add_subparsers()
    new_args = []
    patterns = {}

    # Add general arguments
    for arg_name, arg_params in action_map['general_arguments'].items():
        if 'full' in arg_params:
             arg_names = [arg_name, arg_params['full']]
             arg_fullname = arg_params['full']
             del arg_params['full']
        else: arg_names = [arg_name]
        parsers['general'].add_argument(*arg_names, **arg_params)

    del action_map['general_arguments']

    # Split categories into subparsers
    for category, category_params in action_map.items():
        if 'category_help' not in category_params: category_params['category_help'] = ''
        subparsers_category[category] = subparsers.add_parser(category, help=category_params['category_help'])
        subparsers_action[category] = subparsers_category[category].add_subparsers()
        # Split actions
        if 'actions' in category_params:
            for action, action_params in category_params['actions'].items():
                if 'action_help' not in action_params: action_params['action_help'] = ''
                parsers[category + '_' + action] = subparsers_action[category].add_parser(action, help=action_params['action_help'])
                # Set the action s related function
                parsers[category + '_' + action].set_defaults(
                    func=str_to_func('yunohost_' + category + '.'
                                     + category + '_' + action))
                # Add arguments
                if 'arguments' in action_params:
                    for arg_name, arg_params in action_params['arguments'].items():
                        arg_fullname = False

                        if 'password' in arg_params:
                            if arg_params['password']: is_password = True
                            del arg_params['password']
                        else: is_password = False

                        if 'full' in arg_params:
                            arg_names = [arg_name, arg_params['full']]
                            arg_fullname = arg_params['full']
                            del arg_params['full']
                        else: arg_names = [arg_name]

                        if 'ask' in arg_params:
                            require_input = True
                            if '-h' in sys.argv or '--help' in sys.argv:
                                require_input = False
                            if (category != sys.argv[1]) or (action != sys.argv[2]):
                                require_input = False
                            for name in arg_names:
                                if name in sys.argv[2:]: require_input = False

                            if require_input:
                                if is_password:
                                    if os.isatty(1):
                                        pwd1 = getpass.getpass(colorize(arg_params['ask'] + ': ', 'cyan'))
                                        pwd2 = getpass.getpass(colorize('Retype ' + arg_params['ask'][0].lower() + arg_params['ask'][1:] + ': ', 'cyan'))
                                        if pwd1 != pwd2:
                                            raise YunoHostError(22, _("Passwords don't match"))
                                            sys.exit(1)
                                    else:
                                        raise YunoHostError(22, _("Missing arguments") + ': ' + arg_name)
                                    if arg_name[0] == '-': arg_extend = [arg_name, pwd1]
                                    else: arg_extend = [pwd1]
                                else:
                                    if os.isatty(1):
                                        arg_value = raw_input(colorize(arg_params['ask'] + ': ', 'cyan'))
                                    else:
                                        raise YunoHostError(22, _("Missing arguments") + ': ' + arg_name)
                                    if arg_name[0] == '-': arg_extend = [arg_name, arg_value]
                                    else: arg_extend = [arg_value]
                                new_args.extend(arg_extend)
                            del arg_params['ask']

                        if 'pattern' in arg_params:
                            if (category == sys.argv[1]) and (action == sys.argv[2]):
                                if 'dest' in arg_params: name = arg_params['dest']
                                elif arg_fullname: name = arg_fullname[2:]
                                else: name = arg_name
                                patterns[name] = arg_params['pattern']
                            del arg_params['pattern']

                        parsers[category + '_' + action].add_argument(*arg_names, **arg_params)

    args = parsers['general'].parse_args(sys.argv.extend(new_args))
    args_dict = vars(args)
    for key, value in patterns.items():
        validate(value, args_dict[key])

    return args


def main():
    """
    Main instructions

    Parse the action_dict and execute the action-specific function,
    then print json or pretty result if executed in a tty :)

    Returns:
        int -- 0 or error code

    """
    if len(sys.argv) < 2:
        sys.argv.append('-h')

    with open('action_map.yml') as f:
        action_map = yaml.load(f)

    admin_password_provided = False
    json_print = False

    for key, arg in enumerate(sys.argv):
        if arg == '--admin-password':
            admin_password_provided = True
            admin_password = sys.argv[key+1]
            sys.argv.pop(key)
            sys.argv.pop(key)
        if arg == '--json':
            json_print = True
            sys.argv.pop(key)

    try:
        args = parse_dict(action_map)
        args_dict = vars(args).copy()
        for key in args_dict.keys():
            sanitized_key = key.replace('-', '_')
            if sanitized_key is not key:
                args_dict[sanitized_key] = args_dict[key]
                del args_dict[key]
        del args_dict['func']
        if admin_password_provided:
            with YunoHostLDAP(password=admin_password):
                result = args.func(**args_dict)
        else:
            result = args.func(**args_dict)
    #except TypeError, error:
        #if not __debug__ :
            #traceback.print_exc()
        #print(_("Not (yet) implemented function"))
        #return 1
    except YunoHostError, error:
        display_error(error, json_print)
        return error.code
    else:
        if json_print or not os.isatty(1):
            if result is None:
                result = {}
            if len(win) > 0:
                result['success'] = []
                for msg in win:
                    result['success'].append(msg)
            print(json.dumps(result))
        elif result is not None:
            pretty_print_dict(result)
        else:
            pass

    return 0

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