#!/usr/bin/env python
##############################################################################
# Copyright (c) CERN, 2012.
#
# 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.
##############################################################################

"""Ginfo - Developped by Ivan Calvet for CERN - ivan.calvet@cern.ch"""

import ldap
import sys
import getopt
try:
	import json
except:
	import simplejson as json
import os
import re
import urllib2
import signal

VERSION = '1.0.1'
TIMEOUT = 60

# configuration dictionnary that will be filled with the submitted options
CONF = {'object': None, 'flag': {}, 'filter': {}, 'attribute': []}

# long_flag: [short_flag, parameter_needed]
FLAGS = {
    'host': [None, True],
    'bind': ["b", True],
    'list': ["l", True],
    'clean': [None, False],
    'strict': ["s", False],
    'csv': ["c", False],
    'json': ["j", False],
    'timeout': ["t", True],
    'verbose': ["v", False],
    'version': ["V", False],
    'help': ["h", False],
    }

# {'object': {'attribute': 'attribute name's for the bdii', ...}, ...}
ELTS = {
'AdminDomain':
    {'ID': 'DomainID',
     'Description': 'DomainDescription'},
'Location':
    {'ID': 'LocationID',
     'Country': 'LocationCountry',
     'Latitude': 'LocationLatitude',
     'Longitude': 'LocationLongitude'},
'Service':
    {'ID': 'ServiceID',
     'Capability': 'ServiceCapability',
     'Type': 'ServiceType',
     'QualityLevel': 'ServiceQualityLevel',
     'StatusInfo': 'ServiceStatusInfo',
     'AdminDomainID': 'ServiceAdminDomainForeignKey'},
'Endpoint':
    {'ID': 'EndpointID',
     'URL': 'EndpointURL',
     'Capability': 'EndpointCapability',
     'InterfaceName': 'EndpointInterfaceName',
     'InterfaceVersion': 'EndpointInterfaceVersion',
     'Implementor': 'EndpointImplementor',
     'ImplementationVersion': 'EndpointImplementationVersion',
     'QualityLevel': 'EndpointQualityLevel',
     'HealthState': 'EndpointHealthState',
     'ServingState': 'EndpointServingState',
     'ServiceForeignKey': 'EndpointServiceForeignKey'},
'ComputingShare':
    {'ID': 'ShareID',
     'MaxCPUTime': 'ComputingShareMaxCPUTime',
     'MaxWallTime': 'ComputingShareMaxWallTime',
     'ServingState': 'ComputingShareServingState',
     'ExecutionEnvironmentForeignKey': 'ComputingShareExecutionEnvironmentForeignKey',
     'RunningJobs': 'ComputingShareRunningJobs',
     'WaitingJobs': 'ComputingShareWaitingJobs'},
'MappingPolicy':
    {'ID': 'PolicyID',
     'Scheme': 'PolicyScheme',
     'Rule': 'PolicyRule',
     'ComputingShareID': 'MappingPolicyShareForeignKey'},
'GlueCESEBindGroupCEUniqueID':
    {},
'ExecutionEnvironment':
    {'ID': 'ResourceID',
     'OSName': 'ExecutionEnvironmentOSName',
     'ConnectivityOut': 'ExecutionEnvironmentConnectivityOut',
     'MainMemorySize': 'ExecutionEnvironmentMainMemorySize',
     'VirtualMemorySize': 'ExecutionEnvironmentVirtualMemorySize'},
'ComputingManager':
    {'ID': 'ManagerID',
     'ProductName': 'ManagerProductName',
     'ProductVersion': 'ManagerProductVersion',
     'ServiceID': 'ComputingManagerComputingServiceForeignKey'},
'ToComputingService':
    {},
}

def main(argv):
    """Main function that launches the other functions""" 

    parse_options(argv)
    validate_conf()
    if 'list' in CONF['flag']:
        result = list_attributes() # option --list
    else:
        result = list_object() # get all the informations about the specified object
    if 'verbose' in CONF['flag']:
        print ''
    # result = clean(result)
    print serialize_output(result)
    sys.exit()

def parse_options(argv):
    """Parse the selected options and put them in a config dictionnary""" 

    # build the good sequence to parse the flags with getopt
    short_flags = ''
    long_flags = [[], []] # Long flags with a short flag, or without
    for i in FLAGS:
        j = 1
        if FLAGS[i][0]:
            j = 0
            short_flags += FLAGS[i][0]
            if FLAGS[i][1]:
                short_flags += ':'
        long_flags[j].append(i)
        if FLAGS[i][1]:
            long_flags[j][-1] += '='
    long_flags = long_flags[0] + long_flags[1]
    # identify flags and put them in the config dictionnary
    try:
        flags, args = getopt.getopt(argv, short_flags, long_flags)
    except getopt.error, err:
        sys.exit(usage())
    for flag, arg in flags:
        flag = flag[flag.rfind('-')+1:]
        for i in FLAGS:
            if flag in (i, FLAGS[i][0]):
                if i not in CONF['flag']:
                    if FLAGS[i][1]:
                        CONF['flag'][i] = arg
                    else:
                        CONF['flag'][i] = True
                    break 
                else:
                    sys.exit("Error: Don't use a flag more than once.")
    
    for arg in args:
        # identify filters and put them in the config dictionnary
        if '=' in arg:
            filter = arg[:arg.find('=')]
            CONF['filter'][filter] = arg[arg.find('=')+1:]
        # identify the object and put it in the config dictionnary
        elif arg in ELTS.keys() and CONF['object'] is None:
            CONF['object'] = arg
        # all other elements are attributes and are put in the config dictionnary
        else:
            CONF['attribute'].append(arg)

def validate_conf():
    """Prints verbose messages and checks for errors""" 
        
    # options
    if 'help' in CONF['flag']:
        print usage()
        sys.exit()
    if 'version' in CONF['flag']:
        print os.path.basename(sys.argv[0]) +' V'+VERSION
        sys.exit()
    if 'verbose' in CONF['flag']:
        print 'Verbose mode enabled'
    if 'csv' in CONF['flag'] and 'json' in CONF['flag'] and 'list' not in CONF['flag']:
        sys.exit('Error: choose between csv and json.')
    elif 'csv' in CONF['flag'] and 'list' not in CONF['flag']:
        if 'verbose' in CONF['flag']:
            print 'Output in csv formating'
    elif 'json' in CONF['flag'] and 'list' not in CONF['flag']:
        if 'verbose' in CONF['flag']:
            print 'Output in json formating'
    if 'host' in CONF['flag']:
        if 'verbose' in CONF['flag']:
            print 'The following host will be used:', CONF['flag']['host']
    else:
        if 'LCG_GFAL_INFOSYS' in os.environ:
            CONF['flag']['host'] = os.environ['LCG_GFAL_INFOSYS']
            if 'verbose' in CONF['flag']:
                print 'The following host will be used:', CONF['flag']['host']
        else:
            sys.exit(usage())
    if 'bind' in CONF['flag']:
        if 'verbose' in CONF['flag']:
            print 'The following binding will be used:', CONF['flag']['bind']
    else:
        CONF['flag']['bind'] = 'o=glue'
        if 'verbose' in CONF['flag']:
            print 'The default binding will be used:', CONF['flag']['bind']
    if 'timeout' in CONF['flag']:
        if 'verbose' in CONF['flag']:
            print 'Ldap timeout has been set to '+CONF['flag']['timeout']+' second(s).'

    # Object
    if not CONF['object']:
        sys.exit('Error: Please specify an object.')
    else:
        if 'verbose' in CONF['flag']:
            print 'The specified object is '+CONF['object']+'.'
    if 'list' in CONF['flag']:
        if CONF['flag']['list'] not in ELTS[CONF['object']].keys():
            sys.exit('Error: '+CONF['flag']['list']+' is a wrong attribute.')
        if 'verbose' in CONF['flag']:
            print 'List all the possible values for the following attribute:', CONF['flag']['list']


    # Filters
    if CONF['filter']:
        for filter in CONF['filter']:
            if filter not in ELTS[CONF['object']]:
                sys.exit('Error: '+filter+' is not a valid filter.')
        if 'verbose' in CONF['flag']:
            for filt in CONF['filter']:
                print 'Filter results by the following '+filt+':', CONF['filter'][filt]

    # Attributes
    if CONF['attribute']:
        for att in CONF['attribute']:
            if att not in ELTS[CONF['object']]:
                sys.exit('Error: '+att+' is not a valid attribute.')
    else:
        CONF['attribute'] = ELTS[CONF['object']].keys()
    if 'verbose' in CONF['flag']:
        print 'The following attribute(s) will be displayed:', ', '.join(CONF['attribute'])

def usage():
    """Returns the usage message""" 
    
    return '''Usage: ginfo [options] Object [attribute_to_filter='value of the attribute'] [attribute_to_display]

    List attributes corresponding to an object. By default, all the attributes of an object are
    displayed.

    [OPTIONS]
        --host      host        Specify a host to query. By default the
                                environmental variable LCG_GFAL_INFOSYS will be
                                used.
    -b, --bind      binding     Specify the binding (o=glue by default).
    -l, --list      attribute   List all the possible values of the specified
                                attribute.
    -c, --csv                   Output in CSV format
    -j, --json                  Output in JSON format
    -t, --timeout               Change the ldap timeout (15 seconds by default).
    -v, --verbose               Enable verbose mode
    -V, --version               Prints the version of ginfo
    -h, --help                  Prints this helpful message

    [OBJECTS AND CORRESPONDING ATTRIBUTES]
        AdminDomain:
            ID, Description.
        ComputingManager:
            ID, ProductName, ProductVersion, ServiceID.
        ComputingShare:
            ID, MaxCPUTime, MaxWallTime, ServingState, 
            ExecutionEnvironmentForeignKey, RunningJobs, WaitingJobs.
        Endpoint:
            ID, URL, Capability, InterfaceName, InterfaceVersion, Implementor, 
            ImplementationVersion, QualityLevel, HealthState, ServingState, 
            ServiceForeignKey.
        ExecutionEnvironment:
            ID, OSName, ConnectivityOut, MainMemorySize, VirtualMemorySize.
        Location:
            ID, Country, Latitude, Longitude.
        MappingPolicy:
            ID, Scheme, Rule, ComputingShareID.
        Service:
            ID, Capability, Type, QualityLevel, StatusInfo, AdminDomainID.
'''

def handler(signum, frame):
    sys.exit('Error: Timeout to contact the LDAP server.')

def request(filter=None):
    """Returns the result of the ldap request with the filter given"""

    try:
        t = int(CONF['flag']['timeout'])
    except (ValueError, KeyError):
        t = TIMEOUT
    signal.signal(signal.SIGALRM, handler)
    signal.alarm(t)
    if 'host' in CONF['flag']:
        try:
            port = ''
            if ':' not in CONF['flag']['host']:
                port = ':2170'
            con = ldap.initialize('ldap://'+CONF['flag']['host']+port)
            print filter
            if filter:
                result = con.result(con.search(CONF['flag']['bind'], ldap.SCOPE_SUBTREE, filter))[1]
            else:
                result = con.result(con.search(CONF['flag']['bind'], ldap.SCOPE_SUBTREE))[1]
        except ldap.SERVER_DOWN:
            sys.exit('Error: Can\'t contact the LDAP server. Please check your host.')
    return result

def list_object():
    """Returns a dictionary of filtered results from a ldap request""" 

    # Construct the filter
    filter = ''
    for f in CONF['filter']:
        filter += '(GLUE2'+CONF['object'] + f + '=' + CONF['filter'][f] + ')'
    
    # Main loop using Endpoint object
    result = request('(&(objectClass=GLUE2' + CONF['object'] + ')' + filter + ')')
    dic = {} # Final results
    for res in result:
        id = res[1]['GLUE2' + ELTS[CONF['object']]['ID']][0]
        dic[id] = {}
        for att in CONF['attribute']:
            id2 = 'GLUE2' + ELTS[CONF['object']][att]
            if id2 in res[1]:
                dic[id][att] = res[1][id2]
            else:
                dic[id][att] = None
    return dic

def list_attributes():
    """Returns a list of values for a given attribute""" 
    
    id = 'GLUE2' + ELTS[CONF['object']][CONF['flag']['list']]
    result = request('objectClass=GLUE2' + CONF['object'])
    attr_list = []
    for res in result:
        if id in res[1]:
            for att in res[1][id]:
                if att not in attr_list:
                    attr_list.append(att)
        else:
            if 'None' not in attr_list:
                attr_list.append('None')
    attr_list.sort()
    return attr_list

"""
def clean(result):
    ""Clean all the results""
    values = {
        'ql': [
            'development',
            'pre-production',
            'production',
            'testing'
            ],
        'cap': [
            'data.access.flatfiles',
            'data.access.relational',
            'data.access.xml',
            'data.management.replica',
            'data.management.storage',
            'data.management.transfer',
            'data.naming.resolver',
            'data.naming.scheme',
            'data.transfer',
            'executionmanagement.candidatesetgenerator',
            'executionmanagement.dynamicvmdeploy',
            'executionmanagement.executionandplanning',
            'executionmanagement.jobdescription',
            'executionmanagement.jobexecution',
            'executionmanagement.jobmanager',
            'executionmanagement.reservation',
            'information.discovery',
            'information.logging',
            'information.model',
            'information.monitoring',
            'information.provenance',
            'security.accounting',
            'security.attributeauthority',
            'security.authentication',
            'security.authorization',
            'security.credentialstorage',
            'security.delegation',
            'security.identymapping'
        ],
        'int': [
           'MyProxy',
           'SRM',
           'VOBOX',
           'bdii_site',
           'bdii_top',
           'dcap',
           'emi.storm',
           'gsidcap',
           'gsiftp',
           'http',
           'https',
           'lcg-file-catalog',
           'lcg-local-file-catalog',
           'org.glite.ChannelManagement',
           'org.glite.Delegation',
           'org.glite.FileTransfer',
           'org.glite.RTEPublisher',
           'org.glite.ce.ApplicationPublisher',
           'org.glite.ce.CREAM',
           'org.glite.ce.Monitor',
           'org.glite.lb.Server',
           'org.glite.voms',
           'org.glite.voms-admin',
           'org.glite.wms.WMProxy',
           'org.globus.gram',
           'org.nordugrid.gridftpjob',
           'org.nordugrid.ldapglue1',
           'org.nordugrid.ldapglue2',
           'org.nordugrid.ldapng',
           'xroot'
        ],
        'type': [
            'org.dcache.storage',
            'org.glite.wms.WMProxy',
            'org.glite.lb.Server',
            'org.glite.ce.ApplicationPublisher',
            'org.glite.ce.CREAM',
            'org.glite.ce.Monitor',
            'org.glite.rgma.Consumer',
            'org.glite.rgma.Registry',
            'org.glite.rgma.Schema',
            'org.glite.rgma.Browser',
            'org.glite.rgma.PrimaryProducer',
            'org.glite.rgma.SecondaryProducer',
            'org.glite.rgma.OnDemandProducer',
            'org.glite.RTEPublisher',
            'org.glite.voms',
            'org.glite.voms-admin',
            'org.glite.AMGA',
            'org.glite.ChannelManagement',
            'org.glite.FileTransfer',
            'org.glite.FileTransferStats',
            'org.glite.ChannelAgent',
            'org.glite.Delegation',
            'org.glite.KeyStore',
            'org.nordugrid.execution.arex',
            'SRM',
            'gsiftp',
            'org.edg.gatekeeper',
            'it.infn.GridICE',
            'MyProxy',
            'GUMS',
            'gridmap-file',
            'GridCat',
            'edu.caltech.cacr.monalisa',
            'OpenSSH',
            'xroot',
            'BDII',
            'bdii_site',
            'bdii_top',
            'VOBOX',
            'msg.broker.rest',
            'msg.broker.stomp',
            'msg.broker.stomp-ssl',
            'msg.broker.openwire',
            'msg.broker.openwire-ssl',
            'org.lcg.Frontier',
            'org.lcg.Squid',
            'Nagios',
            'National-NAGIOS',
            'Project-NAGIOS',
            'Regional-NAGIOS',
            'data-location-interface',
            'local-data-location-interface',
            'lcg-file-catalog',
            'lcg-local-file-catalog',
            'pbs.torque.server',
            'pbs.torque.maui',
            'other'
        ],
    }
    
    if 'list' in CONF and 'vo' in CONF['list']:
        invalid_list = []
        for i in result:
            # policyRule must begin with VO:
            if not re.match('^VO:', i) and i != 'ALL' and i not in invalid_list:
                invalid_list.append(i)
        for i in invalid_list:
            result.remove(i)
        if invalid_list:
            result.append('INVALID')
    elif 'list' not in CONF:
        for i in result:
            # policyRule must begin with VO:
            if 'vo' in result[i]:
                if result[i]['vo'] != None:
                    for k in range(len(result[i]['vo'])):
                        if not re.match('^VO:', result[i]['vo'][k]) and result[i]['vo'][k] != 'ALL':
                            result[i]['vo'][k] = 'INVALID'

    if 'clean' in CONF or 'strict' in CONF:
        if 'list' in CONF:
            invalid_list = []
            for i in result:
                # all elements which contain spaces or are empty are invalid.
                if not i or i == 'None' or i.find(' ') >= 0:
                    invalid_list.append(i)
                # version should contain a number
                if CONF['list'] in ['intv', 'impv']:
                    if not re.match('^.*[0-9]+', i) and i not in invalid_list:
                        invalid_list.append(i)
                # quality level is restricted
                if CONF['list'] == 'ql':
                    if i not in values['ql'] and i not in invalid_list:
                        invalid_list.append(i)
                # url must be correctly formated
                if CONF['list'] == 'url':
                    if not re.match('(?:(?:([a-z0-9+.-]+)://)|(www\.))+(([a-zA-Z0-9\._-]+\.[a-zA-Z]{2,6})|([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}))(/[a-zA-Z0-9\&amp;%_\./-~-]*)?', i) and i not in invalid_list:
                        invalid_list.append(i)
                if 'strict' in CONF:
                    # "j" is restricted
                    for j in ['cap', 'type', 'int']:
                        if CONF['list'] == j:
                            if i not in values[j] and i not in invalid_list:
                                invalid_list.append(i)
            for i in invalid_list:
                result.remove(i)
            if invalid_list:
                result.append('INVALID')            
        else:
            for i in result:
                # all elements which contain spaces or are empty
                for j in result[i]:
                    if ATTRIBUTES[j][3]:
                        if result[i][j]:
                            for k in range(len(result[i][j])):
                                if not result[i][j][k] or result[i][j][k].find(' ') >= 0:
                                    result[i][j][k] = 'INVALID'
                            continue
                        else:
                            result[i][j] = ['INVALID']
                    elif not result[i][j] or result[i][j].find(' ') >= 0:
                        result[i][j] = 'INVALID'
                # version should contain a number
                for j in ['intv', 'impv']:
                    if j in result[i]:
                        if not re.match('^.*[0-9]+', result[i][j]):
                            result[i][j] = 'INVALID'
                # quality level is restricted
                if 'ql' in result[i]:
                    if result[i]['ql'] not in values['ql']:
                        result[i]['ql'] = 'INVALID'
                # url must be correctly formated
                if 'url' in result[i]:
                    if not re.match('(?:(?:([a-z0-9+.-]+)://)|(www\.))+(([a-zA-Z0-9\._-]+\.[a-zA-Z]{2,6})|([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}))(/[a-zA-Z0-9\&amp;%_\./-~-]*)?', result[i]['url']):
                        result[i]['url'] = 'INVALID'
                if 'strict' in CONF:
                    if 'cap' in result[i]:
                        for k in range(len(result[i]['cap'])):
                            if result[i]['cap'][k] not in values['cap'] and result[i]['cap'][k] != 'INVALID':
                                result[i]['cap'][k] = 'INVALID'
                    for j in ['type', 'int']:
                        # "j" is restricted
                        if j in result[i]:
                            if result[i][j] not in values[j] and result[i][j] != 'INVALID':
                                result[i][j] = 'INVALID'

    #Renaming of attributes
    if 'list' not in CONF:
        result2 = {}
        for i in result:
            result2[i] = {}
            for j in result[i]:
                result2[i][ATTRIBUTES[j][1]+ATTRIBUTES[j][2]] = result[i][j]
        result = result2
    return result
"""

def serialize_output(result):
    """Return the output with the wished format""" 

    if 'list' in CONF['flag']:
        output = '\n'.join(result)
    elif 'csv' in CONF['flag']:
        csv_list = []
        titles = []
        for att in CONF['attribute']:
           titles.append(att) 
        csv_list.append(','.join(titles))
        for id in result:
            tmp_list = []
            for att in CONF['attribute']:
                if result[id][att] == None:
                    tmp_list.append('None')
                elif len(result[id][att]) > 1:
                    tmp_list.append('"'+','.join(result[id][att])+'"')
                else:
                    tmp_list.append(result[id][att][0])
            csv_list.append(','.join(tmp_list))
        output = '\n'.join(csv_list)
    elif 'json' in CONF['flag']:
        output = json.dumps(result)
    else:
        output_list = []
        for id in result:
            for att in result[id]:
                if not result[id][att]:
                    output_list.append(att+': None')
                else:
                    output_list.append(att+': '+', '.join(result[id][att]))
            output_list.append('')
        output = '\n'.join(output_list)
    return output

if __name__ == "__main__":
    main(sys.argv[1:])
