#!/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:
# - David Cameron, <david.cameron@cern.ch>, 2015
# - Cedric Serfon, <cedric.serfon@cern.ch>, 2016
#
# Get country groups from VOMS, get DNs with production role,
# map to Rucio account and set country-<country>: admin
# Get all DNs in country and set country-<country>: user
# Get all DNs in phys/perf group with production role and set group-<group>: admin
# Get all names in cloud support e-group and set cloud-<cloud>: admin

import os
import re
import sys
import shlex
import subprocess

from rucio.api.account import list_accounts, add_account_attribute, del_account_attribute, list_account_attributes
from rucio.common.config import config_get
from rucio.common.exception import RucioException, Duplicate, AccountNotFound

from VOMSAdmin.VOMSCommands import VOMSAdminProxy

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

result = OK


def set_account_attributes(dns):

    if not dns:
        print 'No user'
        return

    account_mapping = {}
    for dn in dns:
        # Get Rucio account
        accounts = list_accounts(filter={'identity': dn, 'account_type': 'USER'})
        if not accounts:
            print "Warning: no user accounts for %s" % dn
            continue
        account_type, role = dns[dn].split(':')
        for acc in accounts:
            account = acc['account']
            if not (account, account_type) in account_mapping:
                account_mapping[(account, account_type)] = role
            elif account_mapping[(account, account_type)] != role and role == 'admin':
                print 'Will promote account %s to admin role' % (account)
                account_mapping[(account, account_type)] = role
            else:
                pass
                # print 'Do nothing for %s:%s DN : %s' % (account, account_type, dn)
    for account, account_type in account_mapping:
        role = account_mapping[(account, account_type)]
        print 'Set %s role for %s for %s' % (role, account_type, account)
        try:
            attrs = dict([(dummy['key'], dummy['value']) for dummy in list_account_attributes(account)])
            if account_type in attrs and attrs[account_type] == role:
                print 'Account %s already has the role %s for %s' % (account, role, account_type)
                continue
            elif account_type in attrs:
                print 'Removing attribute %s from account %s' % (account_type, account)
                del_account_attribute(account_type, account, 'root')
            print 'Adding attribute %s for %s from account %s' % (role, account_type, account)
            add_account_attribute(account_type, role, account, 'root')
        except Duplicate:
            pass
        except RucioException, error:
            print "Failed to add account attribute: %s" % str(error)
            result = WARNING
    return result


def add_cloud_admins():

    for cloud in ['ca', 'de', 'es', 'fr', 'it', 'ng', 'nl', 'ru', 't0', 'tw', 'uk', 'us']:
        egroup = 'atlas-support-cloud-%s' % cloud
        if cloud == 'ru':
            egroup = 'atlas-adc-cloud-ru'
        cmd = "/usr/bin/ldapsearch -x -h xldap.cern.ch -b 'CN=%s,OU=e-groups,OU=Workgroups,DC=cern,DC=ch' member" % egroup

        proc = subprocess.Popen(shlex.split(cmd), stderr=subprocess.STDOUT, stdout=subprocess.PIPE)
        (stdout, stderr) = proc.communicate()
        if proc.returncode != 0:
            print "Warning: Failed to run ldapsearch: %s" % stdout
            result = WARNING
            continue

        # Always exceptions...
        cloud = cloud.replace('t0', 'cern')
        cloud = cloud.replace('ng', 'nd')
        for line in stdout.split('\n'):
            match = re.match('member: CN=(\w+),OU=Users', line)
            if match:
                account = match.group(1)
                print 'Set admin role for cloud-%s for %s' % (cloud, account)
                try:
                    add_account_attribute('cloud-%s' % cloud, 'admin', account, 'root')
                except Duplicate:
                    pass
                except AccountNotFound:
                    print 'Warning: Account %s not in Rucio' % account
                    result = WARNING
                except RucioException, error:
                    print "Failed to add account attribute: %s" % str(error)
                    result = WARNING
    return result

if __name__ == '__main__':
    result = OK
    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)

    admin = VOMSAdminProxy(vo='atlas', host='voms2.cern.ch', port=8443,
                           user_cert=CERT, user_key=KEY)
    res = admin.call_method('list-sub-groups', '/atlas')
    if not res:
        print 'Could not list VOMS groups'
        sys.exit(CRITICAL)

    for group in res:
        match_pattern = re.match('/atlas/(\w\w)$', group) or re.match('/atlas/(cern)$', group) or re.match('/atlas/(usatlas)$', group)
        if not match_pattern:
            match_pattern = re.match('/atlas/(.*-.*)$', group)
            if not match_pattern:
                continue
            vomsgroup = match_pattern.group(0)
            phys_group = match_pattern.group(1)
            dns = {}
            print 'Working on group-%s' % (phys_group)
            list_dns = admin.call_method('list-users-with-role', vomsgroup, 'Role=production') or []
            for dn in list_dns:
                dns[dn._DN] = 'group-%s:admin' % (phys_group)
            set_account_attributes(dns)
            continue

        vomsgroup = match_pattern.group(0)
        country = match_pattern.group(1)
        if country == 'usatlas':
            country = 'us'

        print 'Working on country-%s' % (country)
        dns = {}
        # Get DNs in country group
        list_dns = admin.call_method('list-members', vomsgroup) or []
        for dn in list_dns:
            dns[dn._DN] = 'country-%s:user' % (country)

        # Get DNs with production role in each country (upgrades the role from user to admin if applicable)
        list_dns = admin.call_method('list-users-with-role', vomsgroup, 'Role=production') or []
        for dn in list_dns:
            dns[dn._DN] = 'country-%s:admin' % (country)
        set_account_attributes(dns)

    # Add cloud admins from ldap egroups
    add_cloud_admins()

    sys.exit(result)
