#!/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.3'
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': ["H", 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':
    {},
'StorageShare':
    {'ID': 'StorageShareSharingID',
    'Path': 'StorageSharePath',
    'AccessMode': 'StorageShareAccessMode',
    'AccessLatency': 'StorageShareAccessLatency',
    'ServingState': 'StorageShareServingState',
    'RetentionPolicy': 'StorageShareRetentionPolicy',
    'ExpirationMode': 'StorageShareExpirationMode',
    'DefaultLifeTime': 'StorageShareDefaultLifeTime',
    'MaximumLifeTime': 'StorageShareMaximumLifeTime',
    'Tag': 'StorageShareTag'},
}

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]
    -H, --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               Print the version of ginfo
    -h, --help                  Print 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.
        StorageShare:
            ID, Path, AccessMode, AccessLatency, ServingState, RetentionPolicy,
            ExpirationMode, DefaultLifeTime, MaximumLifeTime, Tag.
'''

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)
            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 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:])
