check_voms

#!/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-2018
# - David Cameron, <david.cameron@cern.ch>, 2015
# - Dimitrios Christidis, <dimitrios.christidis@cern.ch>, 2019
#
# PY3K COMPATIBLE

from __future__ import print_function

import os
import sys
import time

import ldap  # pylint: disable=import-error

import requests
import requests.auth

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


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

LDAP_HOSTS = ['ldaps://xldap.cern.ch']
LDAP_OPTIONS = [
    # default configuration comes from /etc/openldap/ldap.conf
    # (ldap.OPT_X_TLS_CACERTDIR, '/etc/pki/tls/certs'),
    # (ldap.OPT_X_TLS_CACERTFILE, '/etc/pki/tls/certs/ca-bundle.crt'),
    # (ldap.OPT_X_TLS_REQUIRE_CERT, ldap.OPT_X_TLS_NEVER),
    (ldap.OPT_X_TLS_REQUIRE_CERT, ldap.OPT_X_TLS_DEMAND),
    (ldap.OPT_REFERRALS, 0),
]
LDAP_BASE = 'OU=Users,OU=Organic Units,DC=cern,DC=ch'
LDAP_SCOPE = ldap.SCOPE_SUBTREE
LDAP_FILTER = '(objectClass=user)'
LDAP_ATTRS = ['cn', 'mail', 'proxyAddresses', 'cernExternalMail', 'cernAccountType']
LDAP_PAGE_SIZE = 1000


def get_accounts_identities_fast():
    from rucio.db.sqla.session import get_session
    session = get_session()
    query = '''select b.identity, a.account, a.email, a.account_type, b.identity_type from atlas_rucio.accounts a,  atlas_rucio.account_map b where a.account=b.account'''
    dns = {}
    account2identities = {}
    try:
        result = session.execute(query)
        for identity, account, email, atype, itype in result:
            account2identities.setdefault(account, []).append((itype, identity))
            if itype not in ['X509', 'OIDC']:
                continue
            dns.setdefault(atype, {})[identity] = (account, email)
        return dns, account2identities
    except Exception as error:
        print(error)


def get_accounts_identities_slow():
    dns = {}
    account2identities = {}
    client = Client()
    for account in client.list_accounts():
        #if account['type'] != 'USER':
        #    continue
        for identity in client.list_identities(account['account']):
            account2identities.setdefault(account['account'], []).append((identity['type'], identity['identity']))
            if identity['type'] not in ['X509', 'OIDC']:
                continue
            if account['type'] == 'USER' and identity['identity'] in dns.get(account['type'], {}):
                print("Duplicate identity for users {} and {}: {}".format(dns[account['type']][identity['identity']][0], account['account'], identity['identity']))
            dns.setdefault(account['type'], {})[identity['identity']] = (account['account'], account['email'])
    return dns, account2identities


get_accounts_identities = get_accounts_identities_slow


def get_ldap_identities():
    """Get user account info from CERN AD/LDAP"""
    for opt_key, opt_val in LDAP_OPTIONS:
        ldap.set_option(opt_key, opt_val)

    conn = ldap.initialize(",".join(LDAP_HOSTS))
    conn.simple_bind_s()

    paged_serverctrls = []
    old_paged_search = [int(x) for x in ldap.__version__.split('.')] < [2, 4, 0]
    if old_paged_search:
        paged_serverctrls.append(ldap.controls.SimplePagedResultsControl(ldap.LDAP_CONTROL_PAGE_OID, True, (LDAP_PAGE_SIZE, '')))
    else:
        paged_serverctrls.append(ldap.controls.SimplePagedResultsControl(True, size=LDAP_PAGE_SIZE, cookie=''))

    ret = {}
    while True:
        msgid = conn.search_ext(LDAP_BASE, LDAP_SCOPE, filterstr=LDAP_FILTER, attrlist=LDAP_ATTRS, serverctrls=paged_serverctrls)
        rtype, rdata, rmsgid, serverctrls = conn.result3(msgid=msgid)
        for dn, attrs in rdata:
            cn = attrs['cn'][0].decode('utf-8')
            user = {
                'mails': [],
                # 'x509': [],
            }

            for pmail in [x.decode('utf-8') for x in attrs.get('proxyAddresses', [])]:
                if pmail.lower().startswith('smtp:'):
                    mail = pmail[len('smtp:'):]
                    if mail.lower() not in [umail.lower() for umail in user['mails']]:
                        user['mails'].append(mail)

            for attr in ['mail', 'cernExternalMail']:
                if attr not in attrs:
                    continue

                mail = attrs[attr][0].decode('utf-8')
                user[attr] = mail

                if mail.lower() not in [umail.lower() for umail in user['mails']]:
                    user['mails'].append(mail)

            user['type'] = attrs.get('cernAccountType', [b'Unknown'])[0].decode('utf-8')
#            user['type'] = {
#                'Primary': 'USER',
#                'Secondary': 'USER',
#                'Service': 'SERVICE',
#            }.get(utype, 'USER')

            ret[cn] = user

        cookie = None
        for serverctrl in serverctrls:
            if old_paged_search:
                if serverctrl.controlType == ldap.LDAP_CONTROL_PAGE_OID:
                    unused_est, cookie = serverctrl.controlValue
                    if cookie:
                        serverctrl.controlValue = (LDAP_PAGE_SIZE, cookie)
                    break
            else:
                if serverctrl.controlType == ldap.controls.SimplePagedResultsControl.controlType:
                    cookie = serverctrl.cookie
                    if cookie:
                        serverctrl.size = LDAP_PAGE_SIZE
                    break

        if not cookie:
            break

        paged_serverctrls = serverctrls
    return ret


def get_access_token(token_url, client_id, client_secret):
    data = {
        'grant_type': 'client_credentials'
    }
    basic_auth = requests.auth.HTTPBasicAuth(client_id, client_secret)
    response = requests.post(token_url, data=data, auth=basic_auth)

    if response.status_code != 200:
        print("Failed to get access token")
        sys.exit(CRITICAL)

    token_info = response.json()
    return token_info['access_token']


def get_users(scim_url, access_token):
    ret = []

    with requests.Session() as session:
        session.headers = {
            'Authorization': "Bearer {}".format(access_token),
        }

        start = 1
        page_size = 100
        while True:
            response = session.get("{}/Users".format(scim_url), params={'startIndex': start, 'count': page_size})
            if response.status_code != 200:
                print("Failed to get users (offset={})".format(start))
                sys.exit(CRITICAL)
            d = response.json()
            if d['itemsPerPage'] == 0: break
            for user in d['Resources']:
                if not user['active']: continue
                certs = []
                for certificate in user.get('urn:indigo-dc:scim:schemas:IndigoUser', {}).get('certificates', []):
                    certs.append(certificate['subjectDn'])
                groups = [x['display'] for x in user.get('groups', [])]
                ret.append((user['id'], user['userName'], user['active'], user['emails'][0]['value'], certs, groups))
            start += page_size

    return ret


if __name__ == '__main__':
    TEST = True if 'TEST' in os.environ and os.environ['TEST'].lower() in ['1', 'y', 'yes', 'true'] else False
    SKIP_REMOVE = True if 'SKIP_REMOVE' in os.environ and os.environ['SKIP_REMOVE'].lower() in ['1', 'y', 'yes', 'true'] else False
    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")
        try:
            CERT = config_get('client', 'client_cert')
            KEY = config_get('client', 'client_key')
        except Exception as error:
            print("Failed to get also client certificate and key from rucio.cfg")
            sys.exit(CRITICAL)
    ISSUER = 'https://atlas-auth.cern.ch/'
    TOKEN_URL = '{}/token'.format(ISSUER)
    SCIM_URL = '{}/scim'.format(ISSUER)
    try:
        CLIENT_ID = os.environ['CLIENT_ID']
        CLIENT_SECRET = os.environ['CLIENT_SECRET']
    except Exception:
        print("Failed to get token client from environment")
        sys.exit(CRITICAL)
    starttime = time.time()
    status = OK
    nbusers = 0
    syncusers = 0
    client = Client()
    accounts = {account['account']: account for account in client.list_accounts()}
    scopes = [_ for _ in client.list_scopes()]
    dns, account2identities = get_accounts_identities()
    ldap_accounts = get_ldap_identities()

    # synchronize accounts, identities and scopes
    synced_accounts = []
    access_token = get_access_token(TOKEN_URL, CLIENT_ID, CLIENT_SECRET)
    for userid, nickname, active, email, certs, groups in get_users(SCIM_URL, access_token):
        if 'atlas' not in groups:
            print("Skipping user {0} because it is not member of ATLAS groups".format(nickname))
            continue
        if TEST:
            print("Processing user {0} ({1}, {2}, {3})".format(nickname, email, certs, groups))

        syncusers += 1
        atype = accounts.get(nickname, {}).get('type', 'USER')
        if nickname not in accounts:
            if ldap_accounts.get(nickname, {}).get('type', 'Unknown') in ['Primary', 'Secondary']:
                atype = 'USER'
            else:
                print("Skipping user {0} because ldap type is {1}".format(nickname, ldap_accounts.get(nickname, {}).get('type', 'Unknown')))
                continue
        if atype == 'USER':
            # prefer existing user accounts discovered by associated dn
            new_nickname = None
            for dn in certs:
                if dn not in dns.get('USER', []):
                    continue
                found = False
                for dn_nickname in dns['USER'][dn]:
                    if nickname == dn_nickname:
                        found = True
                if found:
                    break
                new_nickname = dns['USER'][dn][0]
            if new_nickname is not None and new_nickname != nickname:
                if TEST:
                    print("Using discovered user {0} instead of {1}, certs {2}".format(new_nickname, nickname, certs))
                nickname = new_nickname

        synced_accounts.append(nickname)

        scope = None
        if nickname in accounts:
            # rucio account exists
            # a) synchronize email address
            aemail = accounts[nickname]['email']
            if aemail is None or aemail.lower() != email.lower():
                if not TEST:
                    try:
                        client.update_account(account=nickname, key='email', value=email)
                    except Exception as error:
                        print(error)
                else:
                    print('CMD: rucio-admin account update --account {0} --key email --value {1}'.format(nickname, email))
            # b) add all X.509 identities
            for dn in certs:
                if dn in dns.get(atype, []):
                    continue
                try:
                    if not TEST:
                        client.add_identity(account=nickname, identity=dn, authtype='X509', email=email, default=True)
                    else:
                        print("CMD: rucio-admin identity add --account {0} --type X509 --id '{1}' --email '{2}'".format(nickname, dn, email))
                    nbusers += 1
                    print('Identity {0} added to {2} account {1}'.format(dn, nickname, atype))
                except Duplicate:
                    pass
            # c) add OIDC identity
            oidc_identity = "SUB={}, ISS={}".format(userid, ISSUER)
            if oidc_identity not in dns.get(atype, []):
                try:
                    if not TEST:
                        client.add_identity(account=nickname, identity=oidc_identity, authtype='OIDC', email=email, default=True)
                    else:
                        print("CMD: rucio-admin identity add --account {0} --type OIDC --id '{1}' --email '{2}'".format(nickname, oidc_identity, email))
                    nbusers += 1
                    print('Identity {0} added to {2} account {1}'.format(userid, nickname, atype))
                except Duplicate:
                    pass
            if atype == 'USER':
                scope = 'user.' + nickname

        elif active and nickname in ldap_accounts:
            # no rucio account exists for this nickname
            if email.lower() not in [mail.lower() for mail in ldap_accounts[nickname]['mails']]:
                print('Account %s (%s) without matching email in LDAP does not exist. To create it : rucio-admin account add --type USER --email %s %s' % (nickname, ldap_accounts[nickname]['type'], email, nickname))
            else:
                # found matching user account and email in the CERN LDAP
                try:
                    if not TEST:
                        client.add_account(account=nickname, type_='USER', email=email)
                    else:
                        print('CMD: rucio-admin account add --type USER --email {0} {1}'.format(email, nickname))
                    print('Account {0} added'.format(nickname))
                    for dn in certs:
                        if not TEST:
                            client.add_identity(account=nickname, identity=dn, authtype='X509', email=email, default=True)
                        else:
                            print("CMD: rucio-admin identity add --account {0} --type X509 --id '{1}' --email '{2}'".format(nickname, dn, email))
                        nbusers += 1
                        print('Identity {0} added to USER account {1}'.format(dn, nickname))
                    oidc_identity = "SUB={}, ISS={}".format(userid, ISSUER)
                    if not TEST:
                        client.add_identity(account=nickname, identity=oidc_identity, authtype='OIDC', email=email, default=True)
                    else:
                        print("CMD: rucio-admin identity add --account {0} --type OIDC --id '{1}' --email '{2}'".format(nickname, oidc_identity, email))
                    nbusers += 1
                    print('Identity {0} added to {2} account {1}'.format(userid, nickname, atype))
                    scope = 'user.' + nickname
                except Exception:
                    print('Failed to add new account %s (%s)' % (nickname, ldap_accounts[nickname]['type']))

        else:
            print("No new account: nickname %s, active %s, certs: %s, scope: %s" % (nickname, active, certs, scope))

        if scope and scope not in scopes and atype == 'USER':
            try:
                if not TEST:
                    client.add_scope(nickname, scope)
                else:
                    print("CMD: rucio-admin scope add --scope user.{0} --account {0}".format(nickname))
                print('Scope {0} added to {1}'.format(scope, nickname))
            except Duplicate:
                pass

    # list rucio user accounts that can't be matched with any VO account
    delusers = 0
    for account in accounts.values():
        if account['type'] != 'USER':
            continue
        aname = account['account']
        if aname not in synced_accounts:
            delusers += 1
            identities = account2identities.get(aname, [])
            if TEST:
                print('User account {0} is no longer member VO, details {1}, identities({2}) {3}'.format(aname, account, len(identities), identities))

    print('%i users extracted from VOMS, %i users updated, %i users not in VO' % (syncusers, nbusers, delusers))

    sys.exit(status)