#!/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:
# - Vincent Garonne, <vincent.garonne@cern.ch>, 2012-2013
# - Cedric Serfon, <cedric.serfon@cern.ch>, 2014-2016
# - David Cameron, <david.cameron@cern.ch>, 2015

import commands
import os
import re
import sys
import time
import types


from rucio.db.sqla.session import get_session
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

BADCHAR = 1
CASEPROB = 2
NOCERNACC = 3
SERVICEACC = 4
MAILMISMATCH = 5


def get_accounts_identities():
    session = get_session()
    query = '''select b.identity, a.account, a.email from atlas_rucio.accounts a,  atlas_rucio.account_map b where a.account=b.account and identity_type='X509' and account_type='USER' '''
    dns = {}
    try:
        result = session.execute(query)
        for dn, account, email in result:
            dns[dn] = (account, email)
        return dns
    except Exception as error:
        print error


def validate_nickname(nickname, email, dn, debug=False):
    # Check bad characters
    bad_characters = ['.', ' ', '@', '#', ' ']
    for bad_char in bad_characters:
        if bad_char in nickname:
            reason = 'Bad Characters in nickname %(nickname)s for %(dn)s' % locals()
            if debug:
                print reason
            return False, BADCHAR

    # Check lowercase
    if nickname.lower() != nickname:
        reason = 'VOMS nickname %(nickname)s should be lowercase for %(dn)s' % locals()
        if debug:
            print reason
        return False, CASEPROB

    # Check CERN account
    cmd = "/usr/bin/ldapsearch -x -h xldap.cern.ch -b 'OU=Users,OU=Organic Units,DC=cern,DC=ch' '(&(objectClass=user) (CN=%(nickname)s))'" % locals()
    status, output = commands.getstatusoutput(cmd)
    ad_email = re.search('mail: (.*)\n', output)
    ad_name = re.search('displayName: (.*)\n', output)
    ad_dn = re.search('dn: (.*)\n', output)

    if not ad_email:
        reason = "VOMS nickname %(nickname)s for %(dn)s doesn't exist as CERN account" % locals()
        if debug:
            print reason
        return False, NOCERNACC

    # Check mail address
    ad_email = ad_email.groups()[0]
    ad_name = ad_name.groups()[0].replace(' ', '.') + '@cern.ch'
    is_service_account = 'memberOf: CN=cern-accounts-service,OU=e-groups,OU=Workgroups,DC=cern,DC=ch' in output

    if is_service_account:
        reason = "The %(nickname)s is a service account" % locals()
        if debug:
            print reason
        return False, SERVICEACC

    if not is_service_account and email.lower().replace('@mail.cern.ch', '@cern.ch') not in (ad_email.lower(), ad_name.lower(), nickname + '@cern.ch'):
        reason = "Mail mismatch for VOMS nickname %(nickname)s-%(dn)s: %(email)s vs. %(ad_email)s " % locals()
        if debug:
            print reason
            print 'If you want to add manually the identity to the account : rucio-admin identity add --account %s --type X509 --id "%s" --email %s' % (nickname, dn, email)
        return False, MAILMISMATCH

    return True, OK


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 error:
        print "Failed to get proxy from rucio.cfg"
        sys.exit(CRITICAL)
    starttime = time.time()
    status = OK
    nbusers = 0
    nonicknames = []
    client = Client()
    admin = VOMSAdminProxy(vo='atlas', host='voms2.cern.ch', port=8443,
                           user_cert=cert, user_key=key)
    res = admin.call_method('list-users')
    accounts = [i['account'] for i in client.list_accounts()]
    scopes = [i for i in client.list_scopes()]
    dns = get_accounts_identities()
    voms_identities = {}
    if isinstance(res, types.ListType):
        for user in res:
            valid_nickname = True
            dn = user._DN
            scope = None
            if dn in dns:
                # Check if scope exists
                account = dns[dn][0]
                scope = 'user.' + account
            else:
                nbusers += 1
                attempts = 0
                totattemps = 3
                for attempts in xrange(0, totattemps):
                    if attempts < totattemps - 1:
                        try:
                            dn = user._DN
                            ca = user._CA
                            email = user._mail
                            result = admin.call_method('list-user-attributes', dn, ca)
                            if result is None:
                                print "Failed to list-user-attributes for dn: %(dn)" % locals()
                                continue
                            nickname = None
                            try:
                                nickname = result[0]._value
                            except TypeError as error:
                                print 'ERROR : Failed to process DN: %s' % dn
                            if nickname:
                                if nickname in accounts:
                                    if nickname not in voms_identities:
                                        voms_identities[nickname] = [identity['identity'] for identity in client.list_identities(account=nickname) if identity['type'] == 'X509']
                                    if dn not in voms_identities[nickname]:
                                        try:
                                            client.add_identity(account=nickname, identity=dn, authtype='X509', email=email, default=True)
                                            print 'Identity %(dn)s added' % locals()
                                        except Duplicate:
                                            pass
                                    scope = 'user.' + account
                                    break
                                else:
                                    valid_nickname, vstatus = validate_nickname(nickname=nickname, email=email, dn=dn, debug=True)
                                    if valid_nickname:
                                        account = nickname
                                        if account not in accounts:
                                            try:
                                                client.add_account(account=account, type='USER', email=email)
                                                client.add_identity(account=account, identity=dn, authtype='X509', email=email, default=True)
                                                scope = 'user.' + account
                                                print 'Account %(account)s added' % locals()
                                            except Exception:
                                                pass
                                    else:
                                        if nickname not in accounts and vstatus in [MAILMISMATCH, ]:
                                            print 'Account %s does not exist. To create it : rucio-admin account add --type USER --email %s %s' % (nickname, email, nickname)
                                        break
                            elif user._DN not in nonicknames:
                                nonicknames.append(user._DN)
                        except Exception as error:
                            print error
                    else:
                        print 'ERROR getting info for %s' % (user._DN)
                        status = WARNING
            if scope and scope not in scopes and valid_nickname:
                try:
                    client.add_scope(account, scope)
                    print 'Scope %(scope)s added' % locals()
                except Duplicate:
                    pass
    else:
        sys.exit(CRITICAL)
    print '%i users extracted from VOMS' % nbusers
    if nonicknames != []:
        print 'Users with no nickname :'
        print nonicknames

    sys.exit(status)
