#!/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.config import config_get
from rucio.common.exception import Duplicate

from VOMSAdmin.VOMSCommands import VOMSAdminProxy

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


if __name__ == '__main__':
    try:
        proxy = config_get('nagios', 'proxy')
        os.environ["X509_USER_PROXY"] = proxy
        cert, key = os.environ['X509_USER_PROXY'], os.environ['X509_USER_PROXY']
    except Exception as e:
        print "Failed to get proxy from rucio.cfg"
        sys.exit(CRITICAL)
    account_map = {'Role=pilot': 'pilot', 'Role=production': 'pilot'}
    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
        attempts = 0
        totattemps = 3
        for attempts in xrange(0, totattemps):
            res = admin.call_method('list-users-with-role', '/atlas', account)
            if isinstance(res, types.ListType) and (attempts < totattemps - 1):
                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:
                            pass
                        except Exception, e:
                            print e
                    except:
                        print 'ERROR getting info for %s' % (user._DN)
                        status = WARNING
                break
            else:
                sys.exit(CRITICAL)
        print '%i users extracted from VOMS with %s' % (nbusers, account)

    sys.exit(status)
