#!/usr/bin/env python

# Copyright European Organization for Nuclear Research (CERN)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Mario Lassnig, <mario.lassnig@cern.ch>, 2012-2015
# - Vincent Garonne, <vincent.garonne@cern.ch>, 2012-2015
# - Thomas Beermann, <thomas.beermann@cern.ch>, 2012-2013
# - Wen Guan, <wguan@cern.ch>, 2014
# - Martin Barisits, <martin.barisits@cern.ch>, 2014
# - David Cameron, <david.cameron@cern.ch>, 2014
# - Cedric Serfon, <cedric.serfon@cern.ch>, 2014-2016
# - Cheng-Hsi Chao, <cheng-hsi.chao@cern.ch>, 2014
# - Joaquin Bogado, <joaquin.bogado@cern.ch>, 2014-2015

"""
Rucio Administration Utility
"""

import argcomplete
import argparse
import json
import logging
import os
import signal
import sys
import tabulate
import time
import traceback

from ConfigParser import NoOptionError, NoSectionError
from functools import wraps


from rucio.client import Client
from rucio.common.config import config_get
from rucio.common.exception import (AccountNotFound, DataIdentifierAlreadyExists, AccessDenied, DataIdentifierNotFound, InvalidObject, ReplicaNotFound,
                                    RSENotFound, RSEOperationNotSupported, InvalidRSEExpression, DuplicateContent, RuleNotFound, CannotAuthenticate)
from rucio.common.utils import chunks, construct_surl, sizefmt
from rucio import version
from rucio.rse import rsemanager as rsemgr

# If ../lib/rucio/__init__.py exists, add ../ to Python search path, so that
# it will override what happens to be installed in /usr/(local/)lib/python...
possible_topdir = os.path.normpath(os.path.join(os.path.abspath(sys.argv[0]), os.pardir, os.pardir))
if os.path.exists(os.path.join(possible_topdir, 'lib/rucio', '__init__.py')):
    sys.path.insert(0, possible_topdir)


SUCCESS = 0
FAILURE = 1
DEFAULT_PORT = 443


logger = logging.getLogger("user")

tablefmt = 'psql'


def setup_logger(logger):
    logger.setLevel(logging.DEBUG)
    hdlr = logging.StreamHandler()

    def emit_decorator(fcn):
        def func(*args):
            if True:
                formatter = logging.Formatter("%(message)s")
            else:
                levelno = args[0].levelno
                if levelno >= logging.CRITICAL:
                    color = '\033[31;1m'
                elif levelno >= logging.ERROR:
                    color = '\033[31;1m'
                elif levelno >= logging.WARNING:
                    color = '\033[33;1m'
                elif levelno >= logging.INFO:
                    color = '\033[32;1m'
                elif levelno >= logging.DEBUG:
                    color = '\033[36;1m'
                else:
                    color = '\033[0m'
                formatter = logging.Formatter('{0}%(asctime)s %(levelname)s [%(message)s]\033[0m'.format(color))
            hdlr.setFormatter(formatter)
            return fcn(*args)
        return func
    hdlr.emit = emit_decorator(hdlr.emit)
    logger.addHandler(hdlr)

setup_logger(logger)


def signal_handler(signal, frame):
    logger.warning('You pressed Ctrl+C! Exiting gracefully')
    sys.exit(1)

signal.signal(signal.SIGINT, signal_handler)


def exception_handler(function):
    @wraps(function)
    def new_funct(*args, **kwargs):
        try:
            return function(*args, **kwargs)
        except InvalidObject, error:
            logger.error(error)
            return FAILURE
        except DataIdentifierNotFound, error:
            logger.error(error)
            logger.debug('This means that the Data IDentifier you provided is not known by Rucio.')
            return FAILURE
        except AccessDenied, error:
            logger.error(error)
            logger.debug('This error is a permission issue. You cannot run this command with your account.')
            return FAILURE
        except DataIdentifierAlreadyExists, error:
            logger.error(error)
            logger.debug('This means that the data IDentifier you try to add is already registered in Rucio.')
            return FAILURE
        except RSENotFound, error:
            logger.error(error)
            logger.debug('This means that the Rucio Storage Element you provided is not known by Rucio.')
            return FAILURE
        except InvalidRSEExpression, error:
            logger.error(error)
            logger.debug('This means the RSE expression you provided is not syntactically correct.')
            return FAILURE
        except DuplicateContent, error:
            logger.error(error)
            logger.debug('This means that the DID you want to attach is already in the target DID.')
            return FAILURE
        except TypeError, error:
            logger.error(error)
            logger.debug('This means the parameter you passed has a wrong type.')
            return FAILURE
        except RuleNotFound, error:
            logger.error(error)
            logger.debug('This means the rule you specified does not exist.')
            return FAILURE
        except AccountNotFound, error:
            logger.error(error)
            logger.debug('This means that the specified account cannot be found.')
            return FAILURE
        except NotImplementedError, error:
            logger.error(error)
            logger.debug('This means that the method is not implemented yet.')
            return FAILURE
        except Exception, error:
            logger.error(error)
            logger.error('Rucio exited with an unexpected/unknown error, please provide the traceback below to the developers.')
            logger.debug(traceback.format_exc())
            return FAILURE
    return new_funct


def get_client(args):
    """
    Returns a new client object.
    """
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}
    else:
        creds = None

    try:
        client = Client(rucio_host=args.host, auth_host=args.auth_host,
                        account=args.issuer,
                        auth_type=args.auth_strategy, creds=creds,
                        ca_cert=args.ca_certificate, timeout=args.timeout)
    except CannotAuthenticate, error:
        logger.error(error)
        if not args.auth_strategy:
            if 'RUCIO_AUTH_TYPE' in os.environ:
                auth_type = os.environ['RUCIO_AUTH_TYPE']
            else:
                try:
                    auth_type = config_get('client', 'auth_type')
                except (NoOptionError, NoSectionError):
                    logger.error('Cannot get AUTH_TYPE')
                    sys.exit(FAILURE)
        if auth_type == 'x509_proxy':
            logger.error('Please verify that your proxy is still valid and renew it if needed.')
        sys.exit(FAILURE)
    return client


def extract_scope(did):
    # Try to extract the scope from the DSN
    if did.find(':') > -1:
        scope, name = did.split(':')[0], did.split(':')[1]
        if name.endswith('/'):
            name = name[:-1]
        return scope, name
    else:
        scope = did.split('.')[0]
        if did.startswith('user') or did.startswith('group'):
            scope = ".".join(did.split('.')[0:2])
        if did.endswith('/'):
            did = did[:-1]
        return scope, did


@exception_handler
def add_account(args):
    """
    %(prog)s add [options] <field1=value1 field2=value2 ...>

    Adds a new account. Specify metadata fields as arguments.

    """
    client = get_client(args)
    client.add_account(account=args.account, type=args.accounttype, email=args.accountemail)
    print 'Added new account: %s' % args.account
    return SUCCESS


@exception_handler
def delete_account(args):
    """
    %(prog)s disable [options] <field1=value1 field2=value2 ...>

    Delete account.

    """
    client = get_client(args)
    client.delete_account(args.acnt)
    print 'Deleted account: %s' % args.acnt
    return SUCCESS


@exception_handler
def ban_account(args):
    """
    %(prog)s ban [options] <field1=value1 field2=value2 ...>

    Ban an account.

    """
    client = get_client(args)
    client.set_account_status(account=args.account, status='SUSPENDED')
    print 'Account %s banned' % args.account
    return SUCCESS


@exception_handler
def unban_account(args):
    """
    %(prog)s unban [options] <field1=value1 field2=value2 ...>

    Unban a banned account.

    """
    client = get_client(args)
    client.set_account_status(account=args.account, status='ACTIVE')
    print 'Account %s unbanned' % args.account
    return SUCCESS


@exception_handler
def list_accounts(args):
    """
    %(prog)s list [options] <field1=value1 field2=value2 ...>

    List accounts.

    """
    client = get_client(args)
    accounts = client.list_accounts(identity=args.identity, account_type=args.account_type)
    for account in accounts:
        print account['account']
    return SUCCESS


@exception_handler
def info_account(args):
    """
    %(prog)s show [options] <field1=value1 field2=value2 ...>

    Show extended information of a given account

    """
    client = get_client(args)
    info = client.get_account(account=args.account)
    for k in info:
        print k.ljust(10) + ' : ' + str(info[k])
    return SUCCESS


@exception_handler
def list_identities(args):
    """
    %(prog)s list-identities [options] <field1=value1 field2=value2 ...>

    List all identities on an account.
    """
    client = get_client(args)
    identities = client.list_identities(account=args.account)
    for identity in identities:
        print 'Identity: %(identity)s,\ttype: %(type)s' % identity
    return SUCCESS


@exception_handler
def set_limits(args):
    """
    %(prog)s set [options] <field1=value1 field2=value2 ...>

    Set account limit for an account and rse.
    """
    client = get_client(args)
    client.set_account_limit(account=args.account, rse=args.rse, bytes=args.bytes)
    print 'Set account limit for account %s on RSE %s: %s' % (args.account, args.rse, sizefmt(args.bytes, True))
    return SUCCESS


@exception_handler
def get_limits(args):
    """
    %(prog)s get-limits [options] <field1=value1 field2=value2 ...>

    Grant an identity access to an account.

    """
    client = get_client(args)
    limits = client.get_account_limit(account=args.account, rse=args.rse)
    for rse in limits:
        print 'Quota on %s for %s : %s' % (rse, args.account, sizefmt(limits[rse], True))
    return SUCCESS


@exception_handler
def delete_limits(args):
    """
    %(prog)s delete [options] <field1=value1 field2=value2 ...>

    Delete account limit for an account and rse.
    """
    client = get_client(args)
    client.delete_account_limit(account=args.account, rse=args.rse)
    print 'Deleted account limit for account %s and RSE %s' % (args.account, args.rse)
    return SUCCESS


@exception_handler
def identity_add(args):
    """
    %(prog)s del [options] <field1=value1 field2=value2 ...>

    Grant an identity access to an account.

    """
    client = get_client(args)
    if args.email == "":
        print 'Error: --email argument can\'t be an empty string. Failed to grant an identity access to an account'
        return FAILURE
    client.add_identity(account=args.account, identity=args.identity, authtype=args.authtype, email=args.email)
    print 'Added new identity to account: %s-%s' % (args.identity, args.account)
    return SUCCESS


@exception_handler
def identity_delete(args):
    """
    %(prog)s delete [options] <field1=value1 field2=value2 ...>

    Revoke an identity's access to an account.

    """
    client = get_client(args)
    client.del_identity(args.account, args.identity, authtype=args.authtype)
    print 'Deleted identity: %s' % args.identity
    return SUCCESS


@exception_handler
def add_rse(args):
    """
    %(prog)s add [options] <field1=value1 field2=value2 ...>

    Adds a new rse. Specify metadata fields as arguments.

    """
    client = get_client(args)
    client.add_rse(args.rse)
    print 'Added new RSE: %s' % args.rse
    return SUCCESS


@exception_handler
def disable_rse(args):
    """
    %(prog)s del [options] <field1=value1 field2=value2 ...>

    Disable rse.

    """
    client = get_client(args)
    client.delete_rse(args.rse)
    return SUCCESS


@exception_handler
def list_rses(args):
    """
    %(prog)s list [options] <field1=value1 field2=value2 ...>

    List rses.

    """
    client = get_client(args)
    rses = client.list_rses()
    for rse in rses:
        print '%(rse)s' % rse
    return SUCCESS


@exception_handler
def info_rse(args):
    """
    %(prog)s info [options] <field1=value1 field2=value2 ...>

    Show extended information of a given rse

    """
    client = get_client(args)
    rseinfo = client.get_rse(rse=args.rse)
    attributes = client.list_rse_attributes(rse=args.rse)
    usage = client.get_rse_usage(rse=args.rse)
    print 'Settings:'
    print '========='
    for key in rseinfo:
        if key != 'protocols':
            print '  ' + key + ': ' + str(rseinfo[key])
    print 'Attributes:'
    print '==========='
    for attribute in attributes:
        print '  ' + attribute + ': ' + str(attributes[attribute])
    print 'Protocols:'
    print '=========='
    for protocol in rseinfo['protocols']:
        print '  ' + protocol['scheme']
        for item in protocol:
            print '    ' + item + ': ' + str(protocol[item])
    print 'Usage:'
    print '======'
    for elem in usage:
        print '  ' + elem['source']
        for item in elem:
            print '    ' + item + ': ' + str(elem[item])

    return SUCCESS


@exception_handler
def set_attribute_rse(args):
    """
    %(prog)s set-attribute [options] <field1=value1 field2=value2 ...>

    setattr RSE.

    """
    client = get_client(args)
    client.add_rse_attribute(rse=args.rse, key=args.key, value=args.value)
    print 'Added new RSE attribute for %s: %s-%s ' % (args.rse, args.key, args.value)
    return SUCCESS


@exception_handler
def get_attribute_rse(args):
    """
    %(prog)s get-attribute [options] <field1=value1 field2=value2 ...>

    getattr RSE.

    """
    client = get_client(args)
    attributes = client.list_rse_attributes(rse=args.rse)
    for k in attributes:
        print k + ': ' + str(attributes[k])

    return SUCCESS


@exception_handler
def delete_attribute_rse(args):
    """
    %(prog)s delete-attribute [options] <field1=value1 field2=value2 ...>

    setattr RSE.

    """
    client = get_client(args)
    client.delete_rse_attribute(rse=args.rse, key=args.key)
    print 'Deleted RSE attribute for %s: %s-%s ' % (args.rse, args.key, args.value)
    return SUCCESS


@exception_handler
def add_scope(args):
    """
    %(prog)s add [options] <field1=value1 field2=value2 ...>

    Add scope.

    """
    client = get_client(args)
    client.add_scope(account=args.account, scope=args.scope)
    print 'Added new scope to account: %s-%s' % (args.scope, args.account)
    return SUCCESS


@exception_handler
def list_scopes(args):
    """
    %(prog)s list [options] <field1=value1 field2=value2 ...>

    List scopes.

    """
    client = get_client(args)
    scopes = client.list_scopes()
    for scope in scopes:
        print scope
    return SUCCESS


@exception_handler
def get_config(args):
    """
    %(prog)s get [options] <field1=value1 field2=value2 ...>

    Get the configuration. Either everything, or matching the given section/option.
    """
    client = get_client(args)
    res = client.get_config(section=args.section, option=args.option)
    if not isinstance(res, dict):
        print '[%s]\n%s=%s' % (args.section, args.option, str(res))
    else:
        print_header = True
        for i in res.keys():
            if print_header:
                if args.section is not None:
                    print '[%s]' % args.section
                else:
                    print '[%s]' % i
            if not isinstance(res[i], dict):
                print '%s=%s' % (i, str(res[i]))
                print_header = False
            else:
                for j in res[i].keys():
                    print '%s=%s' % (j, str(res[i][j]))
    return SUCCESS


@exception_handler
def set_config_option(args):
    """
    %(prog)s set [options] <field1=value1 field2=value2 ...>

    Set the configuration value for a matching section/option. Missing section/option will be created.
    """
    client = get_client(args)
    client.set_config_option(section=args.section, option=args.option, value=args.value)
    print 'Set configuration: %s.%s=%s' % (args.section, args.option, args.value)
    return SUCCESS


@exception_handler
def delete_config_option(args):
    """
    %(prog)s delete [options] <field1=value1 field2=value2 ...>

    Delete a configuration option from a section
    """
    client = get_client(args)
    if client.delete_config_option(section=args.section, option=args.option):
        print 'Deleted section \'%s\' option \'%s\'' % (args.section, args.option)
    else:
        print 'Section \'%s\' option \'%s\' not found' % (args.section, args.option)
    return SUCCESS


@exception_handler
def add_subscription(args):
    """
    %(prog)s add [options] name filter replication_rules

    Add subscription.

    """
    client = get_client(args)
    if args.subs_account:
        account = args.subs_account
    elif args.issuer:
        account = args.issuer
    else:
        account = client.account
    subscription_id = client.add_subscription(name=args.name, account=account, filter=json.loads(args.filter), replication_rules=json.loads(args.replication_rules),
                                              comments=args.comments, lifetime=args.lifetime, retroactive=False, dry_run=False, priority=args.priority)
    print 'Subscription added %s' % (subscription_id)
    return SUCCESS


@exception_handler
def list_subscriptions(args):
    """
    %(prog)s list [options] [name]

    List subscriptions.

    """
    client = get_client(args)
    if args.subs_account:
        account = args.subs_account
    elif args.issuer:
        account = args.issuer
    else:
        account = client.account
    subs = client.list_subscriptions(name=args.name, account=account)
    for sub in subs:
        if args.long:
            print '\n'.join('%s: %s' % (str(k), str(v)) for (k, v) in sub.items())
            print
        else:
            print "%s: %s %s\n  priority: %s\n  filter: %s\n  rules: %s\n  comments: %s" % (sub['account'], sub['name'], sub['state'], sub['policyid'], sub['filter'], sub['replication_rules'], sub.get('comments', ''))
    return SUCCESS


@exception_handler
def update_subscription(args):
    """
    %(prog)s update [options] name filter replication_rules

    Update a subscription.

    """
    client = get_client(args)
    if args.subs_account:
        account = args.subs_account
    elif args.issuer:
        account = args.issuer
    else:
        account = client.account
    client.update_subscription(name=args.name, account=account, filter=json.loads(args.filter), replication_rules=json.loads(args.replication_rules),
                               comments=args.comments, lifetime=args.lifetime, retroactive=False, dry_run=False, priority=args.priority)
    return SUCCESS


@exception_handler
def reevaluate_did_for_subscription(args):
    """
    %(prog)s reevaulate [options] dids

    Reevaluate a list of DIDs against all active subscriptions.

    """
    client = get_client(args)
    for did in args.dids.split(','):
        scope, name = extract_scope(did)
        client.set_metadata(scope, name, 'is_new', True)
    return SUCCESS


@exception_handler
def list_account_attributes(args):
    """
    %(prog)s show [options] <field1=value1 field2=value2 ...>

    List the attributes for an account.

    """
    client = get_client(args)
    account = args.account or client.account
    attributes = client.list_account_attributes(account).next()
    table = []
    for attr in attributes:
        table.append([attr['key'], attr['value']])
    print tabulate.tabulate(table, tablefmt=tablefmt, headers=['Key', 'Value'])
    return SUCCESS


@exception_handler
def add_account_attribute(args):
    """
    %(prog)s show [options] <field1=value1 field2=value2 ...>

    Add attribute for an account.

    """
    client = get_client(args)
    client.add_account_attribute(account=args.account, key=args.key, value=args.value)
    return SUCCESS


@exception_handler
def delete_account_attribute(args):
    """
    %(prog)s show [options] <field1=value1 field2=value2 ...>

    Delete attribute for an account.

    """
    client = get_client(args)
    client.delete_account_attribute(account=args.account, key=args.key)
    return SUCCESS


@exception_handler
def declare_bad_file_replicas(args):
    """
    %(prog)s show [options] <field1=value1 field2=value2 ...>

    Declare a list of bad replicas.

    """
    client = get_client(args)
    bad_files = []
    if args.inputfile:
        with open(args.inputfile) as infile:
            for line in infile:
                bad_file = line.rstrip('\n')
                if bad_file != '':
                    bad_files.append(bad_file)
    else:
        bad_files = args.listbadfiles
    for chunk in chunks(bad_files, 500):
        non_declared = client.declare_bad_file_replicas(pfns=chunk, reason=args.reason)
        for rse in non_declared:
            for pfn in non_declared[rse]:
                print '%s : PFN %s cannot be declared.' % (rse, pfn)
    return SUCCESS


@exception_handler
def list_pfns(args):
    """
    %(prog)s list [options] <field1=value1 field2=value2 ...>

    List the possible PFN for a file at a site.

    """
    client = get_client(args)
    dids = args.dids.split(',')
    rse = args.rse
    protocol = args.protocol
    for input_did in dids:
        scope, name = extract_scope(input_did)
        replicas = [rep for rep in client.list_replicas([{'scope': scope, 'name': name}, ], schemes=[protocol, ])]
        if rse in replicas[0]['rses'] and replicas[0]['rses'][rse]:
            print replicas[0]['rses'][rse][0]
        else:
            logger.warning('The file has no replica on the specified RSE')
            rse_info = rsemgr.get_rse_info(rse)
            proto = rsemgr.create_protocol(rse_info, 'read', scheme=protocol)
            try:
                pfn = proto.lfns2pfns(lfns={'scope': scope, 'name': name})
                result = pfn.values()[0]
            except ReplicaNotFound as error:
                result = error
            if isinstance(result, RSEOperationNotSupported) or isinstance(result, ReplicaNotFound):
                if not rse_info['deterministic']:
                    logger.warning('This is a non-deterministic site, so the real PFN might be different from the on suggested')
                    rse_attr = client.list_rse_attributes(rse)
                    naming_convention = rse_attr.get('naming_convention', None)
                    parents = [did for did in client.list_parent_dids(scope, name)]
                    if len(parents) > 1:
                        logger.warning('The file has multiple parents')
                    for did in parents:
                        if did['type'] == 'DATASET':
                            path = construct_surl(did['name'], name, naming_convention=naming_convention)
                            pfn = ''.join([proto.attributes['scheme'],
                                           '://',
                                           proto.attributes['hostname'],
                                           ':',
                                           str(proto.attributes['port']),
                                           proto.attributes['prefix'],
                                           path if not path.startswith('/') else path[1:]])
                            print pfn
                else:
                    logger.error('Unexpected error')
                    return FAILURE
            else:
                print result
    return SUCCESS

if __name__ == '__main__':
    usage = """
%(prog)s

 <command> [options] [args]

Commands:

    help <command>  Output help for one of the commands below

"""
    oparser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]), add_help=True)
    subparsers = oparser.add_subparsers()

    # Main arguments
    oparser.add_argument('--version', action='version', version='%(prog)s ' + version.version_string())
    oparser.add_argument('--verbose', '-v', default=False, action='store_true', help="Print more verbose output")
    oparser.add_argument('-H', '--host', dest="host", metavar="ADDRESS", help="The Rucio API host")
    oparser.add_argument('--auth_host', dest="auth_host", metavar="ADDRESS", help="The Rucio Authentication host")
    oparser.add_argument('-a', '--account', dest="issuer", metavar="ACCOUNT", help="Rucio account to use")
    oparser.add_argument('-S', '--auth-strategy', dest="auth_strategy", default=None, help="Authentication strategy (userpass or x509 or ...)")
    oparser.add_argument('-T', '--timeout', dest="timeout", type=float, default=None, help="Set all timeout values to SECONDS")

    # Options for the userpass auth_strategy
    oparser.add_argument('-u', '--user', dest='username', default=None, help='username')
    oparser.add_argument('-pwd', '--password', dest='password', default=None, help='password')

    # Options for the x509  auth_strategy
    oparser.add_argument('--certificate', dest='certificate', default=None, help='Client certificate file')
    oparser.add_argument('--ca-certificate', dest='ca_certificate', default=None, help='CA certificate to verify peer against (SSL)')

    # The account subparser
    account_parser = subparsers.add_parser('account', help='Account methods')
    account_subparser = account_parser.add_subparsers()

    # The list_accounts command
    list_account_parser = account_subparser.add_parser('list', help='List Rucio accounts')
    list_account_parser.add_argument('--type', dest='account_type', action='store', help='Account Type (USER, GROUP, SERVICE)')
    list_account_parser.add_argument('--id', dest='identity', action='store', help='Identity (e.g. DN)')
    list_account_parser.set_defaults(which='list_accounts')

    # The list_account_attributes command
    list_attr_parser = account_subparser.add_parser('list-attributes', help='List attributes for an account')
    list_attr_parser.add_argument('account', action='store', help='Account name')
    list_attr_parser.set_defaults(which='list_account_attributes')

    # The add_account_attribute command
    add_attr_parser = account_subparser.add_parser('add-attribute', help='Add attribute for an account')
    add_attr_parser.add_argument('account', action='store', help='Account name')
    add_attr_parser.add_argument('--key', dest='key', action='store', help='Attribute key', required=True)
    add_attr_parser.add_argument('--value', dest='value', action='store', help='Attribute value', required=True)
    add_attr_parser.set_defaults(which='add_account_attribute')

    # The delete_account_attribute command
    delete_attr_parser = account_subparser.add_parser('delete-attribute', help='Delete attribute for an account')
    delete_attr_parser.add_argument('account', action='store', help='Account name')
    delete_attr_parser.add_argument('--key', dest='key', action='store', help='Attribute key', required=True)
    delete_attr_parser.set_defaults(which='delete_account_attribute')

    # The add_account command
    add_account_parser = account_subparser.add_parser('add', help='Add Rucio account')
    add_account_parser.set_defaults(which='add_account')
    add_account_parser.add_argument('account', action='store', help='Account name')
    add_account_parser.add_argument('--type', dest='accounttype', default='USER', help='Account Type')
    add_account_parser.add_argument('--email', dest='accountemail', action='store',
                                    help='Email address associated with the account')

    # The disable_account command
    delete_account_parser = account_subparser.add_parser('delete', help='Delete Rucio account')
    delete_account_parser.set_defaults(which='delete_account')
    delete_account_parser.add_argument('acnt', action='store', help='Account name')

    # The info_account command
    info_account_parser = account_subparser.add_parser('info', help='Show detailed information about an account')
    info_account_parser.set_defaults(which='info_account')
    info_account_parser.add_argument('account', action='store', help='Account name')

    # The list_account_identities command
    list_account_identities_parser = account_subparser.add_parser('list-identities', help='List all identities on an account')
    list_account_identities_parser.set_defaults(which='list_identities')
    list_account_identities_parser.add_argument('account', action='store', help='Account name')

    # The set-limits command
    set_account_limits_parser = account_subparser.add_parser('set-limits', help='Set the limits for the provided account')
    set_account_limits_parser.set_defaults(which='set_limits')
    set_account_limits_parser.add_argument('account', action='store', help='Account name')
    set_account_limits_parser.add_argument('rse', action='store', help='RSE boolean expression')
    set_account_limits_parser.add_argument('bytes', action='store', type=int, help='The total number of bytes that can be stored')

    # The account-limit subparser
    get_account_limits_parser = account_subparser.add_parser('get-limits', help='To get the account limits on an RSE')
    get_account_limits_parser.set_defaults(which='get_limits')
    get_account_limits_parser.add_argument('account', action='store', help='Account name')
    get_account_limits_parser.add_argument('rse', action='store', help='The RSE name')

    # The delete_quota command
    delete_account_limits_parser = account_subparser.add_parser('delete-limits', help='Delete quota for an account')
    delete_account_limits_parser.set_defaults(which='delete_limits')
    delete_account_limits_parser.add_argument('account', action='store', help='Account name')
    delete_account_limits_parser.add_argument('rse', action='store', help='RSE name')

    # Ban/unban operations not implemented yet
    ban_account_limits_parser = account_subparser.add_parser('ban', help='Ban an account')
    ban_account_limits_parser.set_defaults(which='ban_account')
    ban_account_limits_parser.add_argument('--account', dest='account', action='store', help='Account name', required=True)

    unban_account_limits_parser = account_subparser.add_parser('unban', help='Unban a banned account')
    unban_account_limits_parser.set_defaults(which='unban_account')
    unban_account_limits_parser.add_argument('--account', dest='account', action='store', help='Account name', required=True)

    # The identity subparser
    identity_parser = subparsers.add_parser('identity', help='Identity methods')
    identity_subparser = identity_parser.add_subparsers()

    # The identity_add command
    identity_add_parser = identity_subparser.add_parser('add', help='Grant an identity access to an account')
    identity_add_parser.set_defaults(which='identity_add')
    identity_add_parser.add_argument('--account', dest='account', action='store', help='Account name', required=True)
    identity_add_parser.add_argument('--type', dest='authtype', action='store', choices=['X509', 'GSS', 'USERPASS'], help='Authentication type [X509|GSS|USERPASS]', required=True)
    identity_add_parser.add_argument('--id', dest='identity', action='store', help='Identity', required=True)
    identity_add_parser.add_argument('--email', dest='email', action='store', help='Email address associated with the identity', required=True)

    # The identity_delete command
    identity_delete_parser = identity_subparser.add_parser('delete', help="Revoke an identity's access to an account")
    identity_delete_parser.set_defaults(which='identity_delete')
    identity_delete_parser.add_argument('--account', dest='account', action='store', help='Account name', required=True)
    identity_delete_parser.add_argument('--type', dest='authtype', action='store', help='Authentivation type(e.g. X509, gss)', required=True)
    identity_delete_parser.add_argument('--id', dest='identity', action='store', help='Identity', required=True)

    # The RSE subparser
    rse_parser = subparsers.add_parser('rse', help='RSE(Rucio Storage Element) methods')
    rse_subparser = rse_parser.add_subparsers()

    # The list_rses command
    list_rse_parser = rse_subparser.add_parser('list', help='List Rucio RSEs')
    list_rse_parser.set_defaults(which='list_rses')

    # The add_rse command
    add_rse_parser = rse_subparser.add_parser('add', help='Add RSE')
    add_rse_parser.set_defaults(which='add_rse')
    add_rse_parser.add_argument('rse', action='store', help='RSE name')

    # The info_rse command
    info_rse_parser = rse_subparser.add_parser('info', help='Information about RSE')
    info_rse_parser.set_defaults(which='info_rse')
    info_rse_parser.add_argument('rse', action='store', help='RSE name')

    # The set_attribute_rse command
    set_attribute_rse_parser = rse_subparser.add_parser('set-attribute', help='Add RSE attribute(key-value pair)')
    set_attribute_rse_parser.set_defaults(which='set_attribute_rse')
    set_attribute_rse_parser.add_argument('--rse', dest='rse', action='store', help='RSE name', required=True)
    set_attribute_rse_parser.add_argument('--key', dest='key', action='store', help='Attribute key', required=True)
    set_attribute_rse_parser.add_argument('--value', dest='value', action='store', help='Attribute value', required=True)

    # The delete_attribute_rse command
    delete_attribute_rse_parser = rse_subparser.add_parser('delete-attribute', help='Delete a RSE attribute(key-value pair)')
    delete_attribute_rse_parser.set_defaults(which='delete_attribute_rse')
    delete_attribute_rse_parser.add_argument('--rse', dest='rse', action='store', help='RSE name', required=True)
    delete_attribute_rse_parser.add_argument('--key', dest='key', action='store', help='Attribute key', required=True)
    delete_attribute_rse_parser.add_argument('--value', dest='value', action='store', help='Attribute value', required=True)

    # The get_attribute_rse command
    get_attribute_rse_parser = rse_subparser.add_parser('get-attribute', help='List RSE attributes')
    get_attribute_rse_parser.set_defaults(which='get_attribute_rse')
    get_attribute_rse_parser.add_argument(dest='rse', action='store', help='RSE name')

    # The disable_location command
    disable_rse_parser = rse_subparser.add_parser('delete', help='Disable RSE')
    disable_rse_parser.set_defaults(which='disable_rse')
    disable_rse_parser.add_argument('rse', action='store', help='RSE name')

    # The scope subparser
    scope_parser = subparsers.add_parser('scope', help='Scope methods')
    scope_subparser = scope_parser.add_subparsers()

    # The add_scope command
    add_scope_parser = scope_subparser.add_parser('add', help='Add scope')
    add_scope_parser.set_defaults(which='add_scope')
    add_scope_parser.add_argument('--account', dest='account', action='store', help='Account name', required=True)
    add_scope_parser.add_argument('--scope', dest='scope', action='store', help='Scope name', required=True)

    # The list_scope command
    list_scope_parser = scope_subparser.add_parser('list', help='List scopes')
    list_scope_parser.set_defaults(which='list_scopes')
    list_scope_parser.add_argument('--account', dest='account', action='store', help='Account name')

    # The config subparser
    config_parser = subparsers.add_parser('config', help='Configuration methods')
    config_subparser = config_parser.add_subparsers()

    # The get_config command
    get_config_parser = config_subparser.add_parser('get', help='Get matching configuration')
    get_config_parser.set_defaults(which='get_config')
    get_config_parser.add_argument('--section', dest='section', action='store', help='Section name', required=False)
    get_config_parser.add_argument('--option', dest='option', action='store', help='Option name', required=False)

    # The set_config_option command
    set_config_parser = config_subparser.add_parser('set', help='Set matching configuration')
    set_config_parser.set_defaults(which='set_config_option')
    set_config_parser.add_argument('--section', dest='section', action='store', help='Section name', required=True)
    set_config_parser.add_argument('--option', dest='option', action='store', help='Option name', required=True)
    set_config_parser.add_argument('--value', dest='value', action='store', help='String-encoded value', required=True)

    # The delete_config_option command
    delete_config_parser = config_subparser.add_parser('delete', help='Delete matching configuration')
    delete_config_parser.set_defaults(which='delete_config_option')
    delete_config_parser.add_argument('--section', dest='section', action='store', help='Section name', required=True)
    delete_config_parser.add_argument('--option', dest='option', action='store', help='Option name', required=True)

    # The subscription parser
    subs_parser = subparsers.add_parser('subscription', help='Subscription methods')
    subs_subparser = subs_parser.add_subparsers()

    # The add-subscription command
    add_sub_parser = subs_subparser.add_parser('add', help='Add subscription')
    add_sub_parser.set_defaults(which='add_subscription')
    add_sub_parser.add_argument(dest='name', action='store', help='Subscription name')
    add_sub_parser.add_argument(dest='filter', action='store', help='DID filter (eg \'{"scope": ["tests"], "project": ["data12_8TeV"]}\')')
    add_sub_parser.add_argument(dest='replication_rules', action='store', help='Replication rules (eg \'[{"copies": 2, "rse_expression": "tier=2", "lifetime": 3600, "weight": "mou"}]\')')
    add_sub_parser.add_argument(dest='comments', action='store', help='Comments on subscription')
    add_sub_parser.add_argument('--lifetime', dest='lifetime', action='store', type=int, help='Subscription lifetime (in days)')
    add_sub_parser.add_argument('--account', dest='subs_account', action='store', help='Account name')
    add_sub_parser.add_argument('--priority', dest='priority', action='store', help='The priority of the subscription')
    # retroactive and dry_run hard-coded for now

    # The list-subscriptions command
    list_sub_parser = subs_subparser.add_parser('list', help='List subscriptions')
    list_sub_parser.set_defaults(which='list_subscriptions')
    list_sub_parser.add_argument('--account', dest='subs_account', action='store', help='Account name')
    list_sub_parser.add_argument('--long', dest='long', action='store_true', help='Long listing')
    list_sub_parser.add_argument(dest='name', nargs='?', action='store', help='Subscription name')

    # The update-subscription command
    update_sub_parser = subs_subparser.add_parser('update', help='Update subscription')
    update_sub_parser.set_defaults(which='update_subscription')
    update_sub_parser.add_argument(dest='name', action='store', help='Subscription name')
    update_sub_parser.add_argument(dest='filter', action='store', help='DID filter (eg \'{"scope": ["tests"], "project": ["data12_8TeV"]}\')')
    update_sub_parser.add_argument(dest='replication_rules', action='store', help='Replication rules (eg \'[{"copies": 2, "rse_expression": "tier=2", "lifetime": 3600, "weight": "mou"}]\')')
    update_sub_parser.add_argument(dest='comments', action='store', help='Comments on subscription')
    update_sub_parser.add_argument('--lifetime', dest='lifetime', action='store', type=int, help='Subscription lifetime (in days)')
    update_sub_parser.add_argument('--account', dest='subs_account', action='store', help='Account name')
    update_sub_parser.add_argument('--priority', dest='priority', action='store', help='The priority of the subscription')
    # subscription policy, retroactive and dry_run hard-coded for now

    # The reevaluate command
    reevaluate_did_for_subscription_parser = subs_subparser.add_parser('reevaluate', help='Reevaluate a list of DIDs against all active subscriptions')
    reevaluate_did_for_subscription_parser.set_defaults(which='reevaluate_did_for_subscription')
    reevaluate_did_for_subscription_parser.add_argument(dest='dids', action='store', help='List of DIDs (coma separated)')

    # The replica parser
    rep_parser = subparsers.add_parser('replicas', help='Replica methods')
    rep_subparser = rep_parser.add_subparsers()

    # The declare-bad command
    declare_bad_file_replicas_parser = rep_subparser.add_parser('declare-bad', help='Declare bad file replicas')
    declare_bad_file_replicas_parser.set_defaults(which='declare_bad_file_replicas')
    declare_bad_file_replicas_parser.add_argument(dest='listbadfiles', action='store', nargs='*', help='The list of bad files')
    declare_bad_file_replicas_parser.add_argument('--reason', dest='reason', required=True, action='store', help='Reason')
    declare_bad_file_replicas_parser.add_argument('--inputfile', dest='inputfile', nargs='?', action='store', help='Reason')

    # The list-pfns command
    list_pfns_parser = rep_subparser.add_parser('list-pfns', help='List the possible PFN for a file at a site.')
    list_pfns_parser.set_defaults(which='list_pfns')
    list_pfns_parser.add_argument(dest='dids', action='store', help='List of DIDs (coma separated)')
    list_pfns_parser.add_argument(dest='rse', action='store', help='RSE')
    list_pfns_parser.add_argument(dest='protocol', action='store', default='srm', help='The protocol, by default srm, can be one of [root|srm|http(s)].')

    commands = {'add_account': add_account,
                'list_accounts': list_accounts,
                'list_account_attributes': list_account_attributes,
                'add_account_attribute': add_account_attribute,
                'delete_account_attribute': delete_account_attribute,
                'delete_account': delete_account,
                'info_account': info_account,
                'ban_account': ban_account,
                'unban_account': unban_account,
                'get_limits': get_limits,
                'set_limits': set_limits,
                'delete_limits': delete_limits,
                'list_identities': list_identities,
                'identity_add': identity_add,
                'identity_delete': identity_delete,
                'add_rse': add_rse,
                'set_attribute_rse': set_attribute_rse,
                'get_attribute_rse': get_attribute_rse,
                'delete_attribute_rse': delete_attribute_rse,
                'list_rses': list_rses,
                'disable_rse': disable_rse,
                'add_scope': add_scope,
                'list_scopes': list_scopes,
                'info_rse': info_rse,
                'get_config': get_config,
                'set_config_option': set_config_option,
                'delete_config_option': delete_config_option,
                'add_subscription': add_subscription,
                'list_subscriptions': list_subscriptions,
                'update_subscription': update_subscription,
                'reevaluate_did_for_subscription': reevaluate_did_for_subscription,
                'declare_bad_file_replicas': declare_bad_file_replicas,
                'list_pfns': list_pfns}

    argcomplete.autocomplete(oparser)

    if len(sys.argv) == 1:
        oparser.print_help()
        sys.exit(FAILURE)

    args = oparser.parse_args(sys.argv[1:])

    try:
        start_time = time.time()
        command = commands.get(args.which)
        result = command(args)
        end_time = time.time()
        if args.verbose:
            print "Completed in %-0.4f sec." % (end_time - start_time)
        sys.exit(result)
    except (RuntimeError, NotImplementedError), error:
        print >> sys.stderr, "ERROR: ", error
        sys.exit(FAILURE)
