#!/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:
# - Cedric Serfon, <cedric.serfon@cern.ch>, 2014

import os
import sys
import types


from rucio.client import Client
from rucio.common.exception import Duplicate

from VOMSAdmin.VOMSCommands import VOMSAdminProxy


os.environ['X509_USER_PROXY'] = '/opt/rucio/etc/ddmadmin.proxy.nagios'
cert, key = os.environ['X509_USER_PROXY'], os.environ['X509_USER_PROXY']

UNKNOWN = 3
CRITICAL = 2
WARNING = 1
OK = 0


if __name__ == '__main__':
    account_map = {'Role=pilot': 'pilot', 'Role=production': 'pilot'}
    # account_map = {'Role=pilot': 'pilot', 'Role=production': 'ddmusr01'}
    status = OK
    nbusers = 0
    nonicknames = []
    c = Client()
    admin = VOMSAdminProxy(vo='atlas', host='voms.cern.ch', port=8443,
                           user_cert=cert, user_key=key)
    for account in account_map:
        nbusers = 0
        res = admin.call_method('list-users-with-role', '/atlas', account)
        if isinstance(res, types.ListType):
            for user in res:
                nbusers += 1
                try:
                    dn = user._DN
                    ca = user._CA
                    email = user._mail
                    print account_map[account], dn, ca, email
                    try:
                        c.add_identity(account=account_map[account], identity=dn, authtype='X509', email=email, default=True)
                        print 'Identity %(dn)s added' % locals()
                    except Duplicate:
                        # print 'Identity %(account)s already added' % locals()
                        pass
                except:
                    print 'ERROR getting info for %s' % (user._DN)
                    status = WARNING
        else:
            sys.exit(CRITICAL)
        print '%i users extracted from VOMS with %s' % (nbusers, account)

    sys.exit(status)
