#!/usr/bin/env python
"""
 Copyright European Organization for Nuclear Research (CERN) 2013

 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, 2017
"""

import json
import sys

import requests

from rucio.core.rse import list_rses, update_rse

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

DRYRUN = '-n' in sys.argv

if __name__ == '__main__':
    status = OK
    URL = 'http://atlas-agis-api.cern.ch/request/ddmendpointstatus/query/list/?json'
    resp = requests.get(url=URL)
    DATA = json.loads(resp.content)
    mapping = {'r': 'availability_read', 'w': 'availability_write', 'd': 'availability_delete'}
    blacklisted_sites = {'availability_read': [], 'availability_write': [], 'availability_delete': []}
    for site in DATA:
        for activity in DATA[site]:
            if activity in ['r', 'w', 'd']:
                try:
                    if DATA[site][activity]['status']['value'] == 'OFF':
                        blacklisted_sites[mapping[activity]].append(site)
                except KeyError as error:
                    print error

    URL = 'http://atlas-agis-api.cern.ch/request/ddmendpoint/query/list/?json&state=ACTIVE&site_state=ACTIVE'
    resp = requests.get(url=URL)
    DATA = json.loads(resp.content)
    agis_rses = []
    for rse in DATA:
        if (rse['name'] not in agis_rses) and not (rse['state'] == 'DISABLED' or rse['site_state'] == 'DISABLED'):
            agis_rses.append(rse['name'])

    # Blacklisting
    for activity in blacklisted_sites:
        for site in list_rses({activity: True}):
            if site['rse'] in blacklisted_sites[activity]:
                print '%s will be blacklisted for %s' % (site['rse'], activity)
                try:
                    if not DRYRUN:
                        update_rse(site['rse'], {activity: False})
                except Exception as error:
                    print error
                    if status == OK:
                        status = WARNING

    # Whitelisting
    for activity in blacklisted_sites:
        for site in list_rses({activity: False}):
            if site['rse'] not in agis_rses:
                print '%s does not exist in AGIS' % (site['rse'])
                continue
            if site['rse'] not in blacklisted_sites[activity]:
                print '%s will be whitelisted for %s' % (site['rse'], activity)
                try:
                    if not DRYRUN:
                        update_rse(site['rse'], {activity: True})
                except Exception as error:
                    print error
                    if status == OK:
                        status = WARNING
    sys.exit(status)
