#!/usr/bin/env python
# Copyright 2012-2018 CERN for the benefit of the ATLAS collaboration.
#
# 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
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors:
#  - Cedric Serfon, <cedric.serfon@cern.ch>, 2014-2019
#
# PY3K COMPATIBLE

from __future__ import print_function

import json
import os
import sys

import requests

from rucio.common.config import config_get
from rucio.core.rse import list_rses, update_rse, get_rse_protocols
from rucio.transfertool.fts3 import FTS3Transfertool


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

DRYRUN = '-n' in sys.argv

if __name__ == '__main__':

    STATUS = 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(error)
        print("Failed to get proxy from rucio.cfg")
        sys.exit(CRITICAL)

    try:
        FTSHOSTS = config_get('conveyor', 'ftshosts')
        FTSHOSTS = FTSHOSTS.split(',')
    except Exception as error:
        print("Failed to get ftshosts")
        STATUS = WARNING

    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': []}
    DOWNTIMES = {}
    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)
                        if DATA[site][activity]['status']['provider'] == 'AGIS' and DATA[site][activity]['status']['reason'].find('https://goc.egi.eu/portal/index.php?Page_Type=Downtime') > -1:
                            DOWNTIMES[site] = DATA[site][activity]['status']['reason']
                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})
                        if activity in ['availability_write']:
                            for rse in BLACKLISTED_SITES[activity]:
                                if rse in DOWNTIMES:
                                    protocols = get_rse_protocols(rse)
                                    for prot in protocols['protocols']:
                                        if 'third_party_copy' in prot['domains']['wan']:
                                            se = '%s://%s' % (prot['scheme'], prot['hostname'])
                                            for fts_host in FTSHOSTS:
                                                print("FTS3Transfertool(external_host='%s').set_se_status(storage_element='%s', message='%s', ban=True, timeout=60)" % (fts_host, se, DOWNTIMES[rse]))
                                                FTS3Transfertool(external_host=fts_host).set_se_status(storage_element=se, message=DOWNTIMES[rse], ban=True, timeout=60)
                except Exception as error:
                    print(error)
                    if STATUS == OK:
                        STATUS = WARNING

    # Whitelisting
    for activity in BLACKLISTED_SITES:
        for site in list_rses({activity: False}):
            rse = site['rse']
            if rse not in AGIS_RSES:
                print('%s does not exist in AGIS' % (rse))
                continue
            if rse not in BLACKLISTED_SITES[activity]:
                print('%s will be whitelisted for %s' % (rse, activity))
                try:
                    if not DRYRUN:
                        update_rse(rse, {activity: True})
                        protocols = get_rse_protocols(rse)
                        for prot in protocols['protocols']:
                            if 'third_party_copy' in prot['domains']['wan']:
                                se = '%s://%s' % (prot['scheme'], prot['hostname'])
                                for fts_host in FTSHOSTS:
                                    print("FTS3Transfertool(external_host='%s').set_se_status(storage_element='%s', message='%s', ban=False, timeout=60)" % (fts_host, se, DOWNTIMES[rse]))
                                    FTS3Transfertool(external_host=fts_host).set_se_status(storage_element=se, message=DOWNTIMES[rse], ban=False, timeout=60)
                except Exception as error:
                    print(error)
                    if STATUS == OK:
                        STATUS = WARNING

    FTS_BAN_SES = {}
    for fts_host in FTSHOSTS:
        try:
            ban_ses = FTS3Transfertool(external_host=fts_host).list_se_status()
            FTS_BAN_SES[fts_host] = ban_ses
        except Exception as error:
            print(error)
            if STATUS == OK:
                STATUS = WARNING

    SES = []
    for rse in DOWNTIMES:
        protocols = get_rse_protocols(rse)
        for prot in protocols['protocols']:
            if 'third_party_copy' in prot['domains']['wan']:
                se = '%s://%s' % (prot['scheme'], prot['hostname'])
                if se not in SES:
                    SES.append(se)
                    for fts_host in FTS_BAN_SES:
                        if se not in [storage['se'] for storage in FTS_BAN_SES[fts_host]]:
                            try:
                                print("FTS3Transfertool(external_host='%s').set_se_status(storage_element='%s', message='%s', ban=True, timeout=60)" % (fts_host, se, DOWNTIMES[rse]))
                                FTS3Transfertool(external_host=fts_host).set_se_status(storage_element=se, message=DOWNTIMES[rse], ban=True, timeout=60)
                            except Exception as error:
                                print(error)

    for fts_host in FTS_BAN_SES:
        ban_ses = FTS_BAN_SES[fts_host]
        for banned_se in ban_ses:
            if banned_se['se'] not in SES and banned_se['admin_dn'].find('ddmadmin') > -1:
                try:
                    print('%s is banned but is not in downtime. Removing the ban' % (banned_se))
                    FTS3Transfertool(external_host=fts_host).set_se_status(storage_element=banned_se['se'], message='', ban=False, timeout=60)
                except Exception as error:
                    print(error)
                    if STATUS == OK:
                        STATUS = WARNING

    sys.exit(STATUS)
