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 json
if not __debug__:
    import traceback

gettext.install('YunoHost')

try:
    from yunohost import YunoHostError, YunoHostLDAP, str_to_func, colorize, pretty_print_dict
except ImportError:
    sys.stderr.write('Require YunoHost lib')
    sys.exit(1)

"""
Category/actions/arguments dictionnary


Except for general_arguments, this dictionary contains 3 levels
as in this sample command line :

    yunohost   monitor    info    --cpu --ram
       ^          ^         ^          ^
    (script) | category | action | parameters


Above example will lead to the function 'monitor_info(args)'
in the file 'yunohost_monitor.py' with 'cpu' and 'ram'
stored in args dictionnary.

Usage:
    You can add a category at the first level, action at the second one,
    and arguments at the third one.
    If a connexion is needed for the action, don't forget to add it to
    the action parameters (ldap, repo, dns or firewall).

Documentation:
    You can see all arguments settings at the argparse documentation:
    http://docs.python.org/dev/library/argparse.html
    #argparse.ArgumentParser.add_argument

    Don't forget to turn argument as dictionnary ('setting' : 'value')

"""
action_map = { 
    #############################
    #       General args        #
    #############################
    'general_arguments' : {
        '-v' : {
            'full'      : '--version',
            'help'      : _("Display %(prog)s version"),
            'action'    : 'version',
            'version'   : '%(prog)s ' + __version__,
        },
    },
    #############################
    #           User            #
    #############################
    'user' : {
        'category_help' : _("Manage users"),
        'actions'       : {
            ### user_list()
            'list' : {
                'action_help'   : _("List users"),
                'connections'   : ['ldap'],
                'arguments'     : {
                   '--fields' : {
                        'help'      : _("fields to fetch"),
                        'nargs'     : '+',
                   },
                   '-f' : {
                        'full'      : '--filter',
                        'help'      : _("LDAP filter used to search"),
                   },
                   '-l' : {
                        'full'      : '--limit',
                        'help'      : _("Maximum number of user fetched"), 
                   },
                   '-o' : {
                        'full'      : '--offset',
                        'help'      : _("Starting number for user fetching"), 
                   },
                }
            },
            ### user_create()
            'create' : {
                'action_help'   : _("Create user"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    '-u' : {
                        'full'      : '--username',
                        'help'      : _("Must be unique"),
                    },
                    '-f' : {
                        'full'      : '--firstname',
                    },
                    '-l' : {
                        'full'      : '--lastname',
                    },
                    '-m' : {
                        'full'      : '--mail',
                        'help'      : _("Main mail address, must be unique"),
                    },
                    '-p' : {
                        'full'      : '--password',
                    }, 
                }
            },
            ### user_delete()
            'delete' : {
                'action_help'   : _("Delete user"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'users' : {
                        'help'      : _("Username of users to delete"),
                        'nargs'     : '+',
                    },
                }
            },
            ### user_update()
            'update' : {
                'action_help'   : _("Update user informations"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'user' : {
                        'help'      : _("Username of user to update"),
                    },
                    '-f' : {
                        'full'      : '--firstname',
                    },
                    '-l' : {
                        'full'      : '--lastname',
                    },
                    '-m' : {
                        'full'      : '--mail',
                    },
                    '-cp' : {
                        'full'      : '--change-password',
                        'help'      : _("New password to set"),
                        'metavar'   : 'PASSWORD',
                    },
                    '--add-mailforward' : {
                        'help'      : _("Mailforward addresses to add"),
                        'nargs'     : '+',
                        'metavar'   : 'MAIL',
                    },
                    '--remove-mailforward' : {
                        'help'      : _("Mailforward addresses to remove"),
                        'nargs'     : '+',
                        'metavar'   : 'MAIL',
                    },
                    '-add-mailalias' : {
                        'help'      : _("Mail aliases to add"),
                        'nargs'     : '+',
                        'metavar'   : 'MAIL',
                    },
                    '-remove-mailalias' : {
                        'help'      : _("Mail aliases to remove"),
                        'nargs'     : '+',
                        'metavar'   : 'MAIL',
                    },
                }
            },
            ### user_info()
            'info'      : {
                'action_help'   : _("Get user informations"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'user' : {
                        'nargs'     : '?',
                    }, 
                    '-m' : {
                        'full'      : '--mail',
                    }, 
                    '-cn' : {
                        'full'      : '--fullname',
                    }, 
                }
            },
        }
    },
    #############################
    #          Domain           #
    #############################
    'domain' : {
        'category_help' : _("Manage domains"),
        'actions'       : {
            ### domain_list()
            'list' : {
                'action_help'   : _("List domains"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    '-f' : {
                        'full'      : '--filter',
                        'help'      : _("LDAP filter used to search"),
                    },
                    '-l' : {
                        'full'      : '--limit',
                        'help'      : _("Maximum number of domain fetched"), 
                    },
                    '-o' : {
                        'full'      : '--offset',
                        'help'      : _("Starting number for domain fetching"), 
                   },
                }
            },
            ### domain_create()
            'create' : {
                'action_help'   : _("Create a custom domain"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'domain name' : {
                        'help'      : _("Domain name to create"),
                    }
                }
            },
            ### domain_delete()
            'delete' : {
                'action_help'   : _("Delete domains"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'domain' : {
                        'help'      : _("Domain(s) to delete"),
                        'nargs'     : '+',
                    },
                }
            },
            ### domain_info()
            'info' : {
                'action_help'   : _("Get domain informations"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'domain-name' : {}
                }
            },
            ### domain_renewcert()
            'renewcert' : {
                'action_help'   : _("Renew domain certificate"),
                'arguments'     : {
                    'domain-name' : {}
                }
            },
        }
    },
    #############################
    #            App            #
    #############################
    'app' : {
        'category_help' : _("Manage apps"),
        'actions'       : {
            ### app_list()
            'list' : {
                'action_help'   : _("List apps"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    '--fields' : {
                        'help'      : _("fields to fetch"),
                        'nargs'     : '+',
                    },
                    '-f' : {
                        'full'      : '--filter',
                        'help'      : _("LDAP filter used to search"),
                    },
                    '-l' : {
                        'full'      : '--limit',
                        'help'      : _("Maximum number of app fetched"), 
                    },
                    '-o' : {
                        'full'      : '--offset',
                        'help'      : _("Starting number for app fetching"), 
                   },
                }
            },
            ### app_install() TODO: Write help
            'install' : {
                'action_help'   : _("Install apps"),
                'connections'   : ['ldap', 'dns'],
                'arguments'     : {
                    'appname' : {
                        'nargs'     : '+',
                    },
                    '-d' : {
                        'full'      : '--domain',
                    },
                    '-p' : {
                        'full'      : '--path',
                    },
                    '-l' : {
                        'full'      : '--label',
                    },
                    '--public' : {
                        'action'    : 'store_true',
                    },
                    '--protected' : {
                        'action'    : 'store_true',
                    },
                }
            },
            ### app_remove() TODO: Write help
            'remove' : {
                'action_help'   : _("Remove app"),
                'connections'   : ['ldap', 'dns'],
                'arguments'     : {
                    'app' : {
                        'help'      : _("App(s) to delete"),
                        'nargs'     : '+',
                    },
                }
            },
            ### app_upgrade()
            'upgrade' : {
                'action_help'   : _("Upgrade app"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'app' : {
                        'help'      : _("App(s) to upgrade"),
                        'nargs'     : '+',
                    },
                }
            },
            ### app_info() TODO: Write help
            'info' : {
                'action_help'   : _("Get app informations"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'app' : {},
                }
            },
            ### app_addaccess() TODO: Write help
            'addaccess' : {
                'action_help'   : _("Grant access right to users (everyone by default)"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'app' : {
                        'nargs'     : '+',
                    },
                    '-u' : {
                        'full'      : '--user',
                        'nargs'     : '+',
                    },
                }
            },
            ### app_removeaccess() TODO: Write help
            'removeaccess' : {
                'action_help'   : _("Revoke access right to users (everyone by default)"),
                'connections'   : ['ldap'],
                'arguments'     : {
                    'app' : {
                        'nargs'     : '+',
                    },
                    '-u' : {
                        'full'      : '--user',
                        'nargs'     : '+',
                    },
                }
            },
        }
    },
    #############################
    #        Repository         #
    #############################
    'repo' : {
        'category_help' : _("Manage app repositories"),
        'actions'       : {
            ### repo_list()
            'list' : {
                'action_help'   : _("List repositories"),
                'connections'   : ['repo'],
                'arguments'     : {
                    '-f' : {
                        'full'      : '--filter',
                        'help'      : _("LDAP filter used to search"),
                    },
                    '-l' : {
                        'full'      : '--limit',
                        'help'      : _("Maximum number of repository fetched"), 
                    },
                    '-o' : {
                        'full'      : '--offset',
                        'help'      : _("Starting number for repository fetching"), 
                   },
                }
            },
            'add' : {
                'action_help'   : _("Add app repository"),
                'connections'   : ['repo'],
                'arguments'     : {
                    'url' : {
                        'help'      : _("URL of the repository"),
                    },
                    '-n' : {
                        'full'      : '--name',
                        'help'      : _("Unique name of the repository"),
                    },
                }
            },
            'remove' : {
                'action_help'   : _("Remove repository"),
                'connections'   : ['repo'],
                'arguments'     : {
                    'repo' : {
                        'help'      : _("Name or URL of the repository"),
                    },
                }
            },
            'update' : {
                'action_help'   : _("Update app list from the repositories"),
                'connections'   : ['repo'],
            },
        }
    },
    #############################
    #          Monitor          #
    #############################
    'monitor' : {
        'category_help' : _("Monitoring functions"),
        'actions'       : {
            'info': {
                'action_help'   : _("Check System"),
                'connections'   : ['ldap'],
                'arguments' : {
                    '-m' : {
                        'full'      : '--memory',
                        'help'      : _("Check Memory"),
                        'action'    : 'store_true',
                        },
		            '-c' : {
                        'full'      : '--cpu',
                        'help'      : _("Check CPU"),
	                    'action'    : 'store_true',		
                        },
                    '-d' : {
                        'full'      : '--disk',
                        'help'      : _("Check Disk"),
                        'action'    : 'store_true',
                        },
                    '-i' : {
                        'full'      : '--ifconfig',
                        'help'      : _("Ifconfig"),
                        'action'    : 'store_true',
                        },
                    '-u' : {
                        'full'      : '--uptime',
                        'help'      : _("Show Uptime"),
                        'action'    : 'store_true',
                        },
                    }
                },
            }
        },
    },
    #############################
    #         Firewall          #
    #############################
    'firewall' : {
        'category_help' : _("Manage firewall rules"),
        'actions'       : {
            'list' : {
                'action_help'   : _("List all firewall rules"),
                'connections'   : ['firewall'],
            },
            'allow' : {
                'action_help'   : _("Allow connection port/protocol"),
                'connections'   : ['firewall'],
                'arguments'     : {
                    'port' : {
                        'help'      : _("Port to open"),
                    },
                    'protocol' : {
                        'help'      : _("Protocol associated with port"),
                        'choices'   : ['UDP', 'TCP', 'Both'],
                    },
                    'name' : {
                        'help'      : _("Reference name of the rule"),
                    },
                }
            },
            'disallow' : {
                'action_help'   : _("Disallow connection"),
                'connections'   : ['firewall'],
                'arguments'     : {
                    'name' : {
                        'help'      : _("Reference name of the rule to delete"),
                    },
                }
            },
        }
    },
    #############################
    #           Tools           #
    #############################
    'tools' : {
        'category_help' : _("Specific tools"),
        'actions'       : {}
    },
}

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()

    # Add general arguments
    for arg_name, arg_params in action_map['general_arguments'].items():
        if arg_params['full']:
            arg_fullname = arg_params['full']
            del arg_params['full']
            parsers['general'].add_argument(arg_name, arg_fullname, **arg_params)
        else:
            parsers['general'].add_argument(arg_name, **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():
                        if 'full' in arg_params:
                            arg_fullname = arg_params['full']
                            del arg_params['full']
                            parsers[category + '_' + action].add_argument(arg_name, arg_fullname, **arg_params)
                        else:
                            parsers[category + '_' + action].add_argument(arg_name, **arg_params)

    return parsers['general'].parse_args()

def display_error(error):
    """
    Nice error displaying

    """
    if not __debug__ :
        traceback.print_exc()
    if os.isatty(1):
        print('\n' + colorize(_("Error: "), 'red') + error.message)
    else:
        print(json.dumps({ 'error' : error.message }))

def connect_services(action_map):
    """
    Connect to different services needed by the action

    Keyword arguments:
        action_map -- Map of actions

    Returns:
        Dict -- openned connections or error code

    """
    action_dict = action_map[sys.argv[1]]['actions'][sys.argv[2]]
    connections = {}
    required_connections = []

    if 'connections' in action_dict:
        required_connections = action_dict['connections']
    
    try:
        # Connect to different services if the action is requiring it
        if 'ldap' in required_connections:
            connections['ldap'] = YunoHostLDAP()
        if 'firewall' in required_connections:
            connections['firewall'] = open('/etc/init.d/iptables', 'w')
        # TODO: Add other services connections
    except YunoHostError, error:
        display_error(error)
        sys.exit(error.code)
    else:
        return connections

def disconnect_services(connections):
    """
    Disconnect openned connections

    Keyword arguments:
        connections -- Dictionnary of openned connections

    Returns:
        Boolean

    """
    try:
        if 'ldap' in connections:
            connections['ldap'].disconnect()
        if 'firewall' in connections:
            connections['firewall'].close()
        # TODO: Add other services deconnections 
    except YunoHostError, error:
        display_error(error)
        sys.exit(error.code)
    else:
        return True
    
def main(action_map):
    """ 
    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
    
    """ 
    args = parse_dict(action_map)
    connections = connect_services(action_map)

    try:
        if connections:
            result = args.func(vars(args), connections)
        else:
            result = args.func(vars(args))
    except TypeError:
        print(_("Not (yet) implemented function"))
        return 1
    except YunoHostError, error:
        display_error(error)
        return error.code
    else: 
        if os.isatty(1):
            pretty_print_dict(result)
        else:
            print(json.dumps(result))
    finally:
        disconnect_services(connections)

    return 0        

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