#!/usr/bin/python
#
#  vac - Vac command line utility
#
#  Andrew McNab, University of Manchester.
#  Copyright (c) 2013-7. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or
#  without modification, are permitted provided that the following
#  conditions are met:
#
#    o Redistributions of source code must retain the above
#      copyright notice, this list of conditions and the following
#      disclaimer. 
#    o Redistributions in binary form must reproduce the above
#      copyright notice, this list of conditions and the following
#      disclaimer in the documentation and/or other materials
#      provided with the distribution. 
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
#  CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
#  INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
#  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
#  BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
#  TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
#  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
#  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
#
#  Contacts: Andrew.McNab@cern.ch  http://www.gridpp.ac.uk/vac/
#

import os
import errno
import sys
import stat
import glob
import json
import time
import socket
import shutil
import hashlib,base64
import optparse

import pprint

import vac

def queryMachines(options, factoryList):

  totalCount    = 0
  runningCount  = 0
  startingCount = 0
  n             = 0
  
  if options.returnJSON:
    sys.stdout.write('[')
   
  responses = vac.shared.sendMachinesRequests(factoryList, 'vac-command')

  for factoryName in sorted(responses):
    for vmName in sorted(responses[factoryName]['machines']):
          
        if options.returnJSON:
          if n > 0:
            sys.stdout.write(',')

          sys.stdout.write(json.dumps(responses[factoryName]['machines'][vmName]))

        else:
                response = responses[factoryName]['machines'][vmName]
   
                stateStr = ('  ' + response['state'] + 10 * ' ')[0:11]

                if response['state'] == vac.shared.VacState.running:
                                       runningCount += 1
                elif response['state'] == vac.shared.VacState.starting:
                                       startingCount += 1
                totalCount += 1

                if not response['machinetype']:
                  hoursStr = 22 * ' '
                  percentStr = '      '

                elif response['state'] == vac.shared.VacState.running:
                  hoursStr = '%5.2f hrs' % ((time.time() - response['started_time']) / 3600.0)

                  if response['heartbeat_time'] > response['started_time']:
                    hoursStr += ' %0.1f%%' % (100.0 * response['cpu_seconds'] / float(response['heartbeat_time'] - response['started_time']))
                    
                  hoursStr +=  22 * ' '

                  if response['cpu_percentage'] is not None:
                    percentStr = '%5.1f%%' % response['cpu_percentage']
                  else:
                    percentStr = '      '

                elif response['state'] == vac.shared.VacState.starting:
                  hoursStr = ' %0.2f mins' % ((time.time() - response['created_time']) / 60.0) + 22 * ' '
                  percentStr = '      '

                elif response['state'] == vac.shared.VacState.shutdown and response['shutdown_message']:
                  hoursStr = '"' + response['shutdown_message'][0:19] + '"' + 22 * ' '
                  percentStr = '      '
                  
                else:
                  hoursStr = 22 * ' '
                  percentStr = '      '

                if len(response['machine'].split('.')[0]) < 15:
                 vmName = (response['machine'].split('.')[0] + 15*' ')[0:15]
                else:
                 vmName = response['machine'].split('.')[0]

                if 'num_cpus' in response:
                  processorsStr = '%2d' % response['num_cpus']
                else:
                  processorsStr = ' 1'
                  
                if 'machine_model' in response:
                  if response['machine_model'] in vac.shared.vmModels:
                    modelStr = 'VM'
                  elif response['machine_model'] in vac.shared.scModels:
                    modelStr = 'SC'
                  elif response['machine_model'] in vac.shared.dcModels:
                    modelStr = 'DC'
                  else:
                    modelStr = '??'
                else:
                  modelStr = '??'

                print (
                       vmName + ' ' +
                       ((response['machinetype'] + 14 * ' ') if response['machinetype'] else 10 * '-' + '    ')[0:14] + ' ' + 
                       stateStr + '  ' + 
                       processorsStr + ' ' +
                       modelStr + ' ' +
                       hoursStr[0:21] + 
                       percentStr
                      )
        n += 1
        
  if options.returnJSON:
    print ']'

def queryMachinetype(options, machinetypeName, factoryList):

  responses = vac.shared.sendMachinetypesRequests(factoryList, clientName = 'vac-command')

  if options.returnJSON:
    sys.stdout.write('[')

  n = 0
  
  for factoryName in sorted(responses):
    if 'machinetypes' in responses[factoryName] and machinetypeName in responses[factoryName]['machinetypes']:

      if options.returnJSON:
        if n > 0:
          sys.stdout.write(',')
      
        sys.stdout.write(json.dumps(responses[factoryName]['machinetypes'][machinetypeName]))

      else:
        if responses[factoryName]['machinetypes'][machinetypeName]['shutdown_time']:
          shutdownTimeStr = time.strftime('%d %b %H:%M:%S', time.localtime(responses[factoryName]['machinetypes'][machinetypeName]['shutdown_time']))
        else:
          shutdownTimeStr = '-'
          
        if responses[factoryName]['machinetypes'][machinetypeName]['shutdown_message']:
          shutdownMessageStr = responses[factoryName]['machinetypes'][machinetypeName]['shutdown_message']
        else:
          shutdownMessageStr = '-'

        if 'running_machines' in responses[factoryName]['machinetypes'][machinetypeName]:
          print (
                  factoryName.split('.')[0] + ' ' + 
                  '%2d ' % responses[factoryName]['machinetypes'][machinetypeName]['running_machines'] +
                  '%5.1f ' % responses[factoryName]['machinetypes'][machinetypeName]['running_hs06'] +
                  shutdownTimeStr + ' "' +
                  shutdownMessageStr + '"'
                )
        else: # legacy for VacQuery 0.3
          print (
                  factoryName.split('.')[0] + ' ' + 
                  '%2d ' % responses[factoryName]['machinetypes'][machinetypeName]['total_machines'] +
                  '%5.1f ' % responses[factoryName]['machinetypes'][machinetypeName]['total_hs06'] +
                  shutdownTimeStr + ' "' +
                  shutdownMessageStr + '"'
                )
      n += 1

  if options.returnJSON:
    print ']'

def queryFactories(options, factoryList, clientName = 'vac-command'):
   
  responses = vac.shared.sendFactoriesRequests(factoryList)

  n = 0
  
  if options.returnJSON:
    sys.stdout.write('[')

  for factoryName in sorted(responses):

    if options.returnJSON:
      if n > 0:
        sys.stdout.write(',')
    
      sys.stdout.write(json.dumps(responses[factoryName]))
    else:
      print '%s %2d/%d %2d/%d %6.2f %9d %9d %8d/%8d "%s"' % ( 
                     factoryName.split('.')[0],                      
                     responses[factoryName]['running_cpus'], responses[factoryName]['total_cpus'],
                     responses[factoryName]['running_machines'], responses[factoryName]['total_machines'],
                     responses[factoryName]['load_average'],
                     responses[factoryName]['root_disk_avail_kb'],
                     responses[factoryName]['vac_disk_avail_kb'],
                     responses[factoryName]['mem_used_kb'], responses[factoryName]['mem_total_kb'], 
                     responses[factoryName]['vac_version']
                            )

    n += 1

  if options.returnJSON:
    print ']'

def makeSyncRecords(args):

  if len(args) > 1:
    targetYearMonths = args[1:]
  else:
    targetYearMonths = [ time.strftime('%Y%m', time.gmtime(time.time() - 86400)) ]

  for targetYearMonth in targetYearMonths:
    vac.vacutils.makeSyncRecord('/var/lib/vac', targetYearMonth, '/var/lib/vac/tmp')

  return 0


def makeSquidConf(inputFile, outputFile):

  try:
    newConfigStr = open(inputFile,'r').read()
  except Exception as e:
    print 'Failed to open input config (' + str(e) + ')'
  
  try:
    currentConfigStr = open(outputFile,'r').read()
  except:
    currentConfigStr = ''
    
  for factory in sorted(vac.shared.factories):
  
    try:
      ip = socket.gethostbyname_ex(factory)[2][0]
    except:
      continue
  
    newConfigStr += 'acl localnet src ' + ip + '/32\n'
    
    if factory != os.uname()[1]:
      newConfigStr += 'cache_peer ' + ip + ' sibling 3128 3130 default\n'

  if newConfigStr == currentConfigStr:
    # Nothing to change
    return 0
    
  try:
    open(outputFile + '.tmp', 'w').write(newConfigStr)
  except Exception as e:
    print 'Failed writing ' + outputFile + '.tmp' 
    try:
      os.remove(outputFile + '.tmp')
    except:
      return 2
 
    return 1
    
  try:
    os.rename(outputFile + '.tmp', outputFile)
  except Exception as e:
    print 'Failed moving ' + outputFile + '.tmp to ' + outputFile
    try:
      os.remove(outputFile + '.tmp')
    except:
      return 2
 
    return 1
    
  return 0

#
# PROGRAM MAIN !!!
# 

if __name__ == '__main__':

    parser = optparse.OptionParser(usage="usage: %prog [options] command [target]\n\nCommands:\n  machines\n  machinetype\n  factories\n  proxy-init\n  cernvm-signature\n  apel-sync")

    parser.add_option("-s", 
                      "--space",
                      dest="spaceName",
                      help="override Vac space given in configuration file")

    parser.add_option("-t", 
                      "--timeout",
                      dest="udpTimeoutSeconds", 
                      help="set timeout in seconds for UDP queries")

    parser.add_option("-H", 
                      "--include-hs06",
                      action="store_true",
                      dest="includeHS06", 
                      help="include HS06 information")

    parser.add_option("-J", 
                      "--json",
                      action="store_true",
                      dest="returnJSON", 
                      help="return JSON")

    parser.add_option("-l", 
                      "--legacy-proxy",
                      action="store_true",
                      dest="isLegacyProxy", 
                      help="create a legacy Globus proxy")

    parser.add_option("-c",
                      "--cert",
                      dest="certPath", 
                      help="X509 certificate file")

    parser.add_option("-k",
                      "--key",
                      dest="keyPath", 
                      help="X509 certificate's private key file")

    (options, args) = parser.parse_args()

    if len(args) > 0 and args[0]:
    
        readConfError = vac.shared.readConf()
        
        if readConfError:
# Perhaps we should try to ignore these errors in case the command
# line arguments supplement missing info in vac.conf etc? Anyway,
# for now, bad config means we do nothing.
          print 'Reading configuration fails with error: ' + readConfError
          sys.exit(1)
                
        if options.spaceName:
             vac.shared.spaceName = options.spaceName.strip()

        if options.udpTimeoutSeconds:
             vac.shared.udpTimeoutSeconds = float(options.udpTimeoutSeconds.strip())

        if args[0] == 'scan':
            print "It's  vac machines  now!"
            sys.exit(1)

        if args[0] == 'machines' and len(args) > 1:
            queryMachines(options, args[1:])
            sys.exit(0)
        elif args[0] == 'machines':
            queryMachines(options, None)
            sys.exit(0)

        if args[0] == 'factories' and len(args) > 1:
            queryFactories(options, args[1:])
            sys.exit(0)
        elif args[0] == 'factories':
            queryFactories(options, None)
            sys.exit(0)

        if args[0] == 'machinetype' and len(args) == 2:
            queryMachinetype(options, args[1], None)
            sys.exit(0)
        elif args[0] == 'machinetype' and len(args) > 2:
            queryMachinetype(options, args[1], args[2:])
            sys.exit(0)
        elif args[0] == 'machinetype':
            print 'machinetype command requires an argument!'
            sys.exit(1)

        if args[0] == 'proxy-init':
            if not options.certPath:
              print 'proxy-init command requires the --cert/-c option!'
              sys.exit(1)
            elif not options.keyPath:
              print 'proxy-init command requires the --key/-k option!'
              sys.exit(1)
            else:
              print vac.vacutils.makeX509Proxy(certPath = options.certPath, 
                                               keyPath = options.keyPath, 
                                               expirationTime = int(time.time()) + 43200, 
                                               isLegacyProxy = options.isLegacyProxy)
              sys.exit(0)

        if args[0] == 'cernvm-signature':
            if len(args) == 2:
              dict = vac.vacutils.getCernvmImageData(args[1])
              if dict is None:
                print 'Failed to verify signature'
                sys.exit(1)
              else:
                print dict
                sys.exit(0)                
            else:
              print 'cernvm-signature command requires the filename argument!'
              sys.exit(2)

        if args[0] == 'apel-sync':
            sys.exit(makeSyncRecords(args))

        if args[0] == 'squid-conf':
            if len(args) == 3:              
              sys.exit(makeSquidConf(args[1], args[2]))                
            else:
              print 'squid-conf command requires input and output filename arguments!'
              sys.exit(2)

    print 'Must give a recognised command argument to vac!'
    sys.exit(1)
                                                                      