#!/usr/bin/env python
# Copyright (c) 2006-2011 Arastra, Inc.  All rights reserved.
# Arastra, Inc. Confidential and Proprietary.

"""Loads the switch startup configuration file (flash:/startup-config).
This script is run by Sysdb at the time that the Sysdb process starts."""

import os
import re
import Ark
import EosInit
import Tac, Url, Plugins, SimpleConfigFile, Logging
from UrlPlugin.FlashUrl import flashFileUrl
import Tracing
from FileUrlDefs import ( STARTUP_CONFIG_FILE_NAME,
                          SECURE_MONITOR_STARTUP_CONFIG_FILE_NAME )
import SecMonUtil

traceHandle = Tracing.defaultTraceHandle()
t0 = traceHandle.trace0
t1 = traceHandle.trace1

startupConfigPath = flashFileUrl( STARTUP_CONFIG_FILE_NAME )
oneTimeStartupConfigPath = flashFileUrl( STARTUP_CONFIG_FILE_NAME + ".once" )
secMonStartupConfigPath = flashFileUrl( SECURE_MONITOR_STARTUP_CONFIG_FILE_NAME )
startupConfigLoadedPath = "file:/var/tmp/startup-config.loaded"
oneTimeStartupConfigLoadedPath = flashFileUrl( "startup-config.once.loaded" )
secMonStartupConfigLoadedPath = ( "file:/var/tmp/%s.loaded" %
                                  SECURE_MONITOR_STARTUP_CONFIG_FILE_NAME )
zeroTouchStartupConfigLoadedPath = "file:/var/tmp/zerotouch-startup-config.loaded"
zeroTouchConfigPath = flashFileUrl( "zerotouch-config" )
zeroTouchStartupConfigPath = "file:/var/tmp/zerotouch-startup-config"
kickstartConfigPath = flashFileUrl( "kickstart-config" )
kickstartConfigLoadedPath = "file:/var/tmp/kickstart-config.loaded"

options = None

class EosInitPluginsCtx( object ):
   def __init__( self ):
      self.zeroTouchSkuBlackList_ = []

   def registerZeroTouchBlackListedSku( self, _sku ):
      self.zeroTouchSkuBlackList_.append( _sku )

   def zeroTouchSkuBlackList( self ):
      if os.environ.get( "SIMULATION_BLACK_LIST_SKU" ):
         return [os.environ[ "SIMULATION_BLACK_LIST_SKU" ], ]
      else:
         return self.zeroTouchSkuBlackList_

plugins = None
pluginsCtx = None
def doLoadEosInitPlugins():
   global plugins
   global pluginsCtx
   pluginsCtx = EosInitPluginsCtx()
   if plugins is None:
      plugins = Plugins.loadPlugins( 'EosInitPlugin', context=pluginsCtx )

def genZeroTouchStartupConfig( writeToFile ):
   try:
      output = Tac.run( [ "zerotouch-startup" ], stdout=Tac.CAPTURE )
      f = file( writeToFile, 'w' )
      f.write( output )
      f.close()
      return True
   except Tac.SystemCommandError:
      return False

def getSku():
   sku = None
   if "SIMULATION_SKU" in os.environ:
      return os.environ[ "SIMULATION_SKU" ]
   if Ark.getPlatform() == "veos":
      # EOS running in a VM. 'idprom read' won't work,
      # but we want to distinguish between EOS in a VM and
      # simply broken hardware
      return "vEOS"
   prefdl = EosInit.getPrefdl()
   if prefdl:
      skum = re.search( "SKU: (\S+)", prefdl )
      pcam = re.search( "PCA: (\S+)", prefdl )
      if skum:
         sku = skum.group( 1 )
      if not sku:
         pca = pcam.group( 1 )
         if pca:
            # Earlier versions of sonomas and helenas didn't get 
            # programmed with a sku in their idprom. See Fru/genfdl
            # for details.
            if pca.startswith( "PCA00011" ):
               sku = "DCS-7148SX"
            elif pca.startswith( "PCA00003" ):
               sku = "DCS-7124S"
   if not sku:
      return ""
   else:
      return sku

def urlIsValid( url ):
   valid = ( url.exists() and os.stat( url.localFilename() ).st_size != 0 )
   t1( "Checking: URL", url.localFilename(), "valid", valid )
   return valid

def loadConfigFromCli( url, loadedConfigPath, secureMonitor=False ):
   # returns errorCount
   errors = 0
   output = ''
   try:
      f = url.open()
   except EnvironmentError, e:
      # No initial startup-config
      if ( loadedConfigPath and os.path.exists( loadedConfigPath ) ):
         os.unlink( loadedConfigPath )
   else:
      filename = url.localFilename()

      t0( 'Loading', filename, ', sysname:', options.sysname )
      configLoaded = True
      # CliShell has a built-in 120 second connection timeout
      try:
         cmd = [ 'CliShell',
                 '--disable-aaa', '--disable-guards',
                 '-p', '15', '--startup-config',
                 '--sysname', options.sysname ]
         if options.timeout:
            cmd += [ '-T', str( options.timeout ) ]
         if secureMonitor:
            cmd.append( SecMonUtil.cliArg() )
            output = Tac.run( cmd, stdout=Tac.CAPTURE, stderr=Tac.CAPTURE,
                              input=f.read() )
         else:
            cmd.append( filename )
            output = Tac.run( cmd, stdout=Tac.CAPTURE, stderr=Tac.CAPTURE )
      except Tac.SystemCommandError, e:
         t0( "CliShell ERROR: return code", e.error, "output", e.output )
         errors = 1
         output = e.output
         if ( "Unable to connect" in output or
              "Cannot connect to" in output ):
            configLoaded = False
            output += '\nCLI could not connect to ConfigAgent'

      m = re.search( 'errors in startup-config: (\d+)', output )
      if m:
         errors = int( m.group( 1 ) )

      if not secureMonitor:
         print output
      elif errors:
         # do not print the actual output but only an error summary
         print "errors in secure-monitor startup-config: %d" % errors

      # Write it to /var/tmp/*startup-config.loaded (if valid)
      if configLoaded and loadedConfigPath:
         # Get the full startup-config that was loaded
         # We'll stash the loaded startup-config to /var/tmp/*startup-config.loaded*
         # for ElectionMgr
         if secureMonitor:
            # Get the raw secure-monitor file (encrypted)
            f = open( url.localFilename() )
         else:
            f.seek( 0, 0 )
         startupConfig = f.read()

         with open( loadedConfigPath, "w" ) as loadedFile:
            loadedFile.write( startupConfig )

   return errors
   
def runLoadConfig():
   Ark.configureLogManager( "LoadConfig" )
   t0( 'LoadConfig started' )
      
   sysname = options.sysname

   import PyClient
   pc = PyClient.PyClient( sysname, "Sysdb" )
   sysdbStatus = pc.agentRoot()[ "Sysdb" ][ "status" ]
   bridgingConfig = pc.agentRoot()[ "bridging" ][ "input" ][ "config" ][ "cli" ]

   # Set Sysdb/status.startupConfigStatus
   sysdbStatus.startupConfigStatus = "inProgress"

   # load plugins
   doLoadEosInitPlugins()

   startupConfigErrorCount = 0
   try:
      # Change entityManager to None once ASU code has moved
      # BUG147045 - horribly ugly to pass in None rather than an
      # EntityManager, but we know that this isn't used in
      # any of the calls below.
      urlContext = Url.Context( None, disableAaa=True )

      # Parse urls
      startupUrl = Url.parseUrl( startupConfigPath, urlContext )
      oneTimeStartupUrl = Url.parseUrl( oneTimeStartupConfigPath, urlContext )
      zeroTouchStartupConfigUrl = Url.parseUrl( zeroTouchStartupConfigPath, 
                                                urlContext )
      zeroTouchConfigUrl = Url.parseUrl( zeroTouchConfigPath, urlContext )
      kickstartConfigUrl = Url.parseUrl( kickstartConfigPath, urlContext )
      # Read zerotouch-config
      zeroTouchDisable = False
      zeroTouchDisableOnce = False
      try:
         zeroTouchConfig = SimpleConfigFile.SimpleConfigFileDict( 
            zeroTouchConfigUrl.localFilename() )
         disable = zeroTouchConfig.get( 'DISABLE' )
         if disable == 'True':
            zeroTouchDisable = True
         elif disable == 'NextReload':
            zeroTouchDisableOnce = True
            del zeroTouchConfig[ 'DISABLE' ]
      except (IOError, SimpleConfigFile.ParseError):
         pass
      # Remove residual .loaded files if any
      loadedFileList = [ startupConfigLoadedPath,
                         secMonStartupConfigLoadedPath,
                         oneTimeStartupConfigLoadedPath,
                         zeroTouchStartupConfigLoadedPath ]
      for loadedFile in loadedFileList:
         localFilename = Url.parseUrl( loadedFile, urlContext ).localFilename()
         if os.path.exists( localFilename ):
            os.unlink( localFilename )
      # Try to load the local startup-config file. If 'startup-config' 
      # does not exist, try 'startup-config.once' - which is a one time
      # load startup-config created by ztp. If 'startup-config.once' 
      # is not present, we fallback to the built in ztp config, which'll
      # configure the switch to not bridge and turn on ztp.
      url = startupUrl
      loadedStartupConfigUrl = None
      loadedStartupConfigPath = None
      if urlIsValid( startupUrl ):
         t0( 'LoadConfig: will load local startup-config' )
         url = startupUrl
         loadedStartupConfigUrl = Url.parseUrl( startupConfigLoadedPath, 
                                                urlContext )
      elif urlIsValid( oneTimeStartupUrl ):
         t0( 'LoadConfig: will load one time startup-config' )
         url = oneTimeStartupUrl
         loadedStartupConfigUrl = Url.parseUrl( oneTimeStartupConfigLoadedPath, 
                                                urlContext )
      elif not (zeroTouchDisable or zeroTouchDisableOnce):
         ztpBl = None
         if pluginsCtx:
            ztpBl = pluginsCtx.zeroTouchSkuBlackList()
            sku = getSku()
            if sku:
               if not any( [ re.match( bl, sku ) for bl in ztpBl ] ):
                  t0( 'LoadConfig: will load zerotouch startup-config' )
                  if genZeroTouchStartupConfig( 
                     zeroTouchStartupConfigUrl.localFilename() ):
                     url = zeroTouchStartupConfigUrl
                     loadedStartupConfigUrl = Url.parseUrl( 
                           zeroTouchStartupConfigLoadedPath, urlContext )

                     # Set default switchport mode to routed
                     bridgingConfig.defaultSwitchportMode = 'routed'
                     try:
                        xcvrConfig = pc.agentRoot()[ "hardware" ][ "xcvr" ]\
                                     [ "xgc" ]
                        xcvrConfig.xuk = 'xcvrUnlockForZtp'
                     except KeyError:
                        # ignore failure due to autobuild dependency
                        print "WARNING: No xcvr info found"
                  else:
                     t0( 'LoadConfig: skipping zerotouch-startup-config, '
                         'failed to generate startup config' )
               else:
                  t0( 'LoadConfig: skipping zerotouch-startup-config, sku', sku,
                      'found in blacklist' )
            else:
               t0( 'LoadConfig: skipping zerotouch-startup-config, empty sku' )
      elif urlIsValid( kickstartConfigUrl ):
         t0( 'LoadConfig: will load kickstart-config' )
         url = kickstartConfigUrl
         loadedStartupConfigUrl = Url.parseUrl( kickstartConfigLoadedPath, 
                                                urlContext )
         t0( 'kickstartConfigUrl.localFileName: ',
                                       kickstartConfigUrl.localFilename()  )
      else:
         t0( 'LoadConfig: no usable startup-config found' )

      if loadedStartupConfigUrl:
         loadedStartupConfigPath = loadedStartupConfigUrl.localFilename()

      startupConfigErrorCount = loadConfigFromCli( url,
                                                   loadedStartupConfigPath )

      # If this was a one time load startup config, delete
      # it, so that we can go back to ztp mode and download 
      # it afresh from the server next time
      if oneTimeStartupUrl.exists() and url is oneTimeStartupUrl:
         t0( 'LoadConfig: deleting one time load startup-config' )
         os.unlink( oneTimeStartupUrl.localFilename() )

      secMonUrl = Url.parseUrl( secMonStartupConfigPath, urlContext )
      if urlIsValid( secMonUrl ):
         loadedSecMonStartupConfigPath = Url.parseUrl( secMonStartupConfigLoadedPath,
                                                       urlContext ).localFilename()
         startupConfigErrorCount += loadConfigFromCli( secMonUrl,
                                                       loadedSecMonStartupConfigPath,
                                                       secureMonitor=True )

      if startupConfigErrorCount > 0:
         Logging.log(
            EosInit.SYSDB_STARTUP_CONFIG_PARSE_ERROR ) # pylint:disable-msg=E1101
   finally:
      # Write to Sysdb/status that we're done loading the startup config
      sysdbStatus.startupConfigErrorCount = startupConfigErrorCount
      sysdbStatus.startupConfigStatus = "completed"
      t0( 'LoadConfig: completed' )

def main():
   import optparse
   global options
   parser = optparse.OptionParser()
   parser.add_option( "-s", "--sysname", action="store", default="ar",
                      help="System name (default: %default)" )
   parser.add_option( "-t", "--timeout", action="store", default=0,
                      type=int, help="Connection timeout (testing only)" )
   options, args = parser.parse_args()
   if args:
      parser.error( "unexpected arguments" )

   runLoadConfig()
      
if __name__ == '__main__':
   main()
