#!/usr/bin/env python3
# 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 six
import sys
import Ark
import EosInit
import SwagBoot
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 TpmGeneric.Defs as TpmDefs
from TpmGeneric.Tpm import TpmGeneric
from Toggles.AristaSAIToggleLib import toggleAristaSAIStartupConfigModeEnabled

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" )
# pylint: disable-next=consider-using-f-string
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"
swagConfigPath = flashFileUrl( SwagBoot.SWAG_CONFIG_FILE_NAME )
swagConfigLoadedPath = f"file:/var/tmp/{ SwagBoot.SWAG_CONFIG_FILE_NAME }.loaded"

options = None

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

   def registerZeroTouchBlockedSku( self, sku ):
      self.zeroTouchSkuBlockList_.append( sku )

   def zeroTouchSkuBlockList( self ):
      if os.environ.get( "SIMULATION_BLOCK_LIST_SKU" ):
         return [ os.environ[ "SIMULATION_BLOCK_LIST_SKU" ] ]
      else:
         return self.zeroTouchSkuBlockList_

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

def genZeroTouchStartupConfig( writeToFile ):
   try:
      output = Tac.run( [ sys.executable, "/usr/bin/zerotouch-startup" ],
                        stdout=Tac.CAPTURE )
      f = open( writeToFile, 'w' ) # pylint: disable=consider-using-with
      f.write( output )
      f.close()
      return True
   except Tac.SystemCommandError:
      return False

def getSku():
   if "SIMULATION_SKU" in os.environ:
      return os.environ[ "SIMULATION_SKU" ]

   platform = Ark.getPlatform()
   if platform == "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"
   elif platform == "ceoslab":
      # Similarly, we want ztp to be possible inside a
      # cEOS-lab container.
      return "cEOSLab"

   prefdl = EosInit.getPrefdl()
   if prefdl:
      preFdlSku = re.search( br"SKU: (\S+)", prefdl )
      if preFdlSku:
         return six.ensure_str( preFdlSku.group( 1 ) )

      preFdlPca = re.search( br"PCA: (\S+)", prefdl )
      if preFdlPca:
         pca = preFdlPca.group( 1 )
         # Earlier versions of sonomas and helenas didn't get 
         # programmed with a sku in their idprom. See Fru/genfdl
         # for details.
         if pca.startswith( b"PCA00011" ):
            return "DCS-7148SX"
         elif pca.startswith( b"PCA00003" ):
            return "DCS-7124S"

   # If we can't figure out the SKU name we still need to return something
   return ""

def skuIsZtpBlocked():
   if not pluginsCtx:
      return True
   ztpBlock = pluginsCtx.zeroTouchSkuBlockList()
   sku = getSku()
   return not sku or any( re.match( bl, sku ) for bl in ztpBlock )

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

def doPcrExtend( configPath=None, data=None, logMsg='' ):
   # The extension of TPM PCRs is a security feature that allows post boot
   # attestation. We precisely measure every step of EOS that can load code. We
   # record hashes that can be extracted and validated by an external controller
   # in a secure manner. This can used as part of a security model to verify that
   # no tampering of any kind has happened on a device.
   #
   # startup-config being a critical piece that can influence the switche's behavior,
   # it's critical to have a record of this execution.
   #
   # This is what the piece of code does below, if we detect a TPM, support the
   # feature on the platform, and it's enabled (via the 'measuredboot' bit in the
   # secure boot toggle)

   assert configPath or data

   try:
      tpm = TpmGeneric()
      if not tpm.isToggleBitSet( TpmDefs.SBToggleBit.MEASUREDBOOT ):
         return

      if configPath:
         tpm.pcrExtendFromFile( TpmDefs.PcrRegisters.EOS_STARTUP_CONFIG,
                                TpmDefs.PcrEventType.EOS_STARTUP_CONFIG,
                                configPath, log=[ logMsg ] )
      else:
         tpm.pcrExtendFromData( TpmDefs.PcrRegisters.EOS_STARTUP_CONFIG,
                                TpmDefs.PcrEventType.EOS_STARTUP_CONFIG,
                                data, log=[ logMsg ] )
   except ( TpmDefs.NoTpmDevice, TpmDefs.NoTpmImpl, TpmDefs.NoSBToggle ):
      pass
   except TpmDefs.Error as e:
      # pylint: disable-next=consider-using-f-string
      t0( "Failed to extend PCR registers (%s): %s" %
          ( configPath or data, str( e ) ) )

def loadConfigFromCli( url, loadedConfigPath, secureMonitor=False ):
   # returns errorCount
   errors = 0
   output = ''
   try:
      f = url.open()
   except OSError as e: # pylint: disable=unused-variable
      # No initial startup-config
      doPcrExtend( data=TpmDefs.DefaultPcrExtend.NO_STARTUP_CONFIG.value,
                   logMsg='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
      doPcrExtend( configPath=filename,
                   # pylint: disable-next=consider-using-f-string
                   logMsg='loaded startup-config: %s' % filename )
      # CliShell has a built-in 120 second connection timeout
      try:
         cmd = [ sys.executable, '/usr/bin/CliShell',
                 '--disable-aaa', '--disable-guards',
                 '-p', '15', '--startup-config',
                 '--sysname', options.sysname ]
         if options.timeout:
            cmd += [ '-T', str( options.timeout ) ]
         if secureMonitor:
            cmd.append( "--secure-monitor-operator" )
            output = Tac.run( cmd, stdout=Tac.CAPTURE, stderr=Tac.CAPTURE,
                              input=f.read() )
         else:
            cmd.append( filename )
            # Note: the following two markers are used for performance testing of
            # loading startup-config (runtime edition of code!). DO NOT REMOVE!
            # MARKER: BEGIN LOADING STARTUP-CONFIG
            output = Tac.run( cmd, stdout=Tac.CAPTURE, stderr=Tac.CAPTURE )
            # MARKER: END LOADING STARTUP-CONFIG
      except Tac.SystemCommandError as 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( r'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
         # pylint: disable-next=consider-using-f-string
         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() ) # pylint: disable=consider-using-with
         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 # pylint: disable=import-outside-toplevel
   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 )
      swagConfigUrl = Url.parseUrl( swagConfigPath, 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 ( OSError, SimpleConfigFile.ParseError ):
         pass
      # Remove residual .loaded files if any
      loadedFileList = [ swagConfigLoadedPath,
                         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
      # AristaSAI cEOS uses kickstart-config only. Loading a startup config
      # Can prevent the device from initializing correctly. This toggle
      # will only ever be set manually on DUT for debugging when required.
      if ( Ark.getPlatform() == 'ceossai' and not
           toggleAristaSAIStartupConfigModeEnabled() ):
         if urlIsValid( kickstartConfigUrl ):
            t0( 'LoadConfig: will load kickstart-config' )
            url = kickstartConfigUrl
            loadedStartupConfigUrl = Url.parseUrl( kickstartConfigLoadedPath,
                                                   urlContext )
            t0( 'kickstartConfigUrl.localFileName: ',
                kickstartConfigUrl.localFilename() )
      elif SwagBoot.shouldLoadSwagConfig() and urlIsValid( swagConfigUrl ):
         t0( 'LoadConfig: will load local swag-config' )
         url = swagConfigUrl
         loadedStartupConfigUrl = Url.parseUrl( swagConfigLoadedPath, 
                                                urlContext )
      elif urlIsValid( startupUrl ):
         t0( 'LoadConfig: will load local startup-config' )
         url = startupUrl
         loadedStartupConfigUrl = Url.parseUrl( startupConfigLoadedPath, 
                                                urlContext )
         if os.path.exists( "/export/swi/startup-config" ):
            Logging.log(
               EosInit.SYSDB_STARTUP_CONFIG_FROM_BOOT_IMAGE,
               "startup-config" )
         elif os.path.exists( "/export/swi/startup-config.xz" ):
            Logging.log(
               EosInit.SYSDB_STARTUP_CONFIG_FROM_BOOT_IMAGE,
               "startup-config.xz" )
      elif urlIsValid( oneTimeStartupUrl ):
         t0( 'LoadConfig: will load one time startup-config' )
         url = oneTimeStartupUrl
         loadedStartupConfigUrl = Url.parseUrl( oneTimeStartupConfigLoadedPath, 
                                                urlContext )
      elif not (zeroTouchDisable or zeroTouchDisableOnce):
         if skuIsZtpBlocked():
            t0( 'LoadConfig: skipping zerotouch-startup-config, sku empty or '
                'found in blocklist' )
         else:
            # CliSession uses Sysdb EntityManager to generate the cleanconfig
            # So, wait for the cleanconfig generation to be complete before we
            # write anything into Sysdb(for example, defaultSwitchportMode
            # below). Otherwise, lines like
            # "switchport default mode routed" will be seen in the
            # cleanconfig causing CVP services to get confused
            def _ccComplete():
               ar = pc.agentRoot()
               ccStatus = ar[ "cli" ][ "session" ][ "cleanConfigStatus" ]
               return ccStatus.cleanConfigComplete

            try:
               Tac.waitFor( _ccComplete,
                            description="clean-config to complete",
                            maxDelay=1 )
            except Tac.Timeout:
               print( "Timed out waiting for clean-config to complete" )
               sys.exit( 1 )

            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' )
      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 )
   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 # pylint: disable=import-outside-toplevel,deprecated-module
   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()
