#!/usr/bin/env python3
# Copyright (c) 2019 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

# Please read comment in __main__ for an explanation of this script.

import PyClient, Tac
import os, sys, re, time

_debug = 0
_logging = 1
_logBuf = []
_inserted = "INSERTED"
_removed = "REMOVED"
_recordFilePath = "/var/log/event/xcvrDefault40g-record"

# Debug functions
def debugPrint( s ):
   if _debug:
      print( s )

def log( s ):
   debugPrint( s )
   if _logging:
      msg = str( time.time() ) + '\t' + s
      _logBuf.append( msg )

def flushToLog():
   with open( '/var/log/event/xcvrDefault40g-debug', "a+" ) as f:
      for line in _logBuf:
         f.write( line + '\n' )

def doMaybeSetupLogDir():
   # Create the log/event dir if it does not exist already
   path = '/var/log/event'
   if not os.path.exists( path ):
      # If the directory is made between the if-check and this makedirs
      # call, then we're screwed. The event handler should have already
      # made this path though.
      os.makedirs( path )


# Parses a syslog for the port to configure and how to configure it.
# Two syslog messages we would like to match on:
# 1. Fru: %FRU-6-TRANSCEIVER_REMOVED
# 2. Fru: %FRU-6-TRANSCEIVER_INSERTED
def parseLog( syslogMsg ):
   # Matches transceiver fru syslog
   actionMatch = re.search( r'Fru: %FRU-6-TRANSCEIVER_(INSERTED|REMOVED)',
                            syslogMsg )
   action = None
   if actionMatch is not None:
      action = actionMatch.group( 1 )

   # Matches on Ethernet49 or Ethernet49/1 for example
   portMatch = re.search( r'ethernet(\d+)(?:/\d+)?', syslogMsg.lower() )
   portNum = None
   if portMatch is not None:
      portNum = portMatch.group( 1 )

   if action not in [ _inserted, _removed ] or portNum is None:
      log( f'Issue parsing log -> port number {portNum}, action: {action}' )
      return None, None

   return action, portNum

# Uses Tac.run(), a wrapper around python's subprocess module, to invoke
# the CLI and configure the port to 40g.
def configure40g( primaryIntf ):
   command = """FastCli -p 15 -c $'enable
                config
                interface ethernet %s
                terminal dont-ask
                speed forced 40gfull\n'"""
   command = command % primaryIntf
   log( "Executing command sequence %s" % command )
   Tac.run( command, shell=True )
   log( "Successfully configured 40g" )

# Uses Tac.run(), a wrapper around python's subprocess module, to invoke
# the CLI and remove any speed configuration on the given interface.
def remove40g( primaryIntf ):
   command = """FastCli -p 15 -c $'enable
                config
                interface ethernet %s
                terminal dont-ask
                no speed\n'"""
   command = command % primaryIntf
   log( "Executing command sequence %s" % command )
   Tac.run( command, shell=True )
   log( "Successfully removed speed configuration" )

# Given a port number, configures it to 40G via a CLI session
# but only if the port satisfies some requirements:
# - must be 40g but not 100g capable
# - must be a QSFP
def doMaybeConfigure40g( port ):
   # Use pyclient to get access to the qsfpCtrlSm for 'port'
   intfName = 'Ethernet%s' % port
   pycl = PyClient.PyClient( "ar", "XcvrAgent" )
   ctrlAgent = pycl.agentRoot().entity[ 'xcvrCtrlAgent' ]
   qsfpCtrlDir = ctrlAgent.xcvrControllerAgentSm.qsfpPlusCtrlSm
   qsfpCtrlSm = qsfpCtrlDir.get( intfName )

   if qsfpCtrlSm:
      # Unfortunately we don't know the whether or not our interface
      # ends with a /1 or not. Just try both.
      primaryIntfName = intfName + '/1'
      intfNum = port + '/1'

      capabilities = qsfpCtrlSm.capabilitiesDir.get( primaryIntfName )
      if capabilities is None:
         log( "Primary intf %s did not exist in capabilitiesDir" % primaryIntfName )
         # It wasn't a /1, try it without the /1
         intfNum = port
         capabilities = qsfpCtrlSm.capabilitiesDir.get( intfName )
         if capabilities is None:
            # Somehow the interface did not exist in the capabilities dir
            log( "Intf %s did not exist in capabilitiesDir either" % intfName )
            return

      # Determine if the module is 40g and 100g capable
      is40g = capabilities.linkModeCapabilities.mode40GbpsFull
      is100g = capabilities.linkModeCapabilities.mode100GbpsFull
      log( "Transceiver %s capabilities 40G: %s, 100G: %s" % \
           ( intfName, is40g, is100g ) )

      # The transceivers which are 100g capable will already come up at 100g.
      if is40g and not is100g:
         # Transceiver is capable of 40g but not 100g, so configure 40g via
         # the interface config CLI
         configure40g( intfNum )
         # We just configured the port. Write the interface name to the record file.
         with open( _recordFilePath, "a+" ) as f:
            f.write( "Ethernet%s\n" % intfNum )
      else:
         log( "%s does not meet capability requirements, skipping config" )

   else:
      log( "%s is not a QSFP. No operation." % intfName )

# Removes the speed config of the given port but only if this script
# was the one to write the speed config in the first place.
def doMaybeRemove40g( port ):
   # Read in the record file to determine every interface which
   # currently is configured by this script.
   configuredIntfs = readConfigRecord()
   log( "The following intfs are configured by the script: %s" % configuredIntfs )

   # Again we don't know if this intf has a /1 or not, so try both
   intfName = 'Ethernet%s' % port
   primaryIntfName = intfName + '/1'
   intf = None
   intfNum = None
   if intfName in configuredIntfs:
      intf = intfName
      intfNum = port
   elif primaryIntfName in configuredIntfs:
      intf = primaryIntfName
      intfNum = port + '/1'

   if intf is not None:
      # Remove the config and remove the interface from the file.
      # We remove the intf from the file by overwriting the entire file
      # with a new set of interfaces.
      log( "Removing config on %s" % intf )
      remove40g( intfNum )
      newIntfs = [ i for i in configuredIntfs if i != intf ]
      writeConfigRecord( newIntfs )
   else:
      # No-op, just write to the debug log
      log( "Interface %s not in collection of interfaces configured by this"
           " script " % intf )

# Reads in the config record file and returns a list of primary interface
# names which we have configured.
#
# This function can error with file not found exception.
def readConfigRecord():
   log( "Reading config record file" )
   intfs = []
   if not os.path.isfile( _recordFilePath ):
      log( "Record file does not exist. No configuration to remove" )
   else:
      with open( _recordFilePath ) as f:
         for line in f:
            intfs.append( line.strip( ' \n' ) )
   return intfs

# Write a record of all interfaces we have configured to 40g to a file.
# Takes in a list of all interfaces and overwrites the file with the list
# of all interfaces. Each interface is separated by a newline. The file
# might look like this, for example:
#  Ethernet49
#  Ethernet51
#  Ethernet3
def writeConfigRecord( intfs ):
   with open( _recordFilePath, "w" ) as f:
      for intf in intfs:
         f.write( intf + "\n" )

def run():
   global _debug
   # Setup logging
   doMaybeSetupLogDir()

   # Activate debug flag based on passed in argument
   if len( sys.argv ) == 2:
      _debug = 1 if sys.argv[ 1 ].lower() == 'debug' else 0

   # This script has two parts--both are designed to be triggered on a CLI
   # event handler configured by the customer/user.
   #
   # 1. Insertion
   # This script is designed to detect 40G capable modules and configure them
   # to 40G automatically.  The script will only run if the module is 40G capable
   # and not 100G capable.  Since this script invokes a CLI speed change, the
   # forwarding agent should be hit-less on speed change.
   #
   # 2. Removal
   # When a module is removed, the speed change command is deleted.  This will only
   # happen if the customer/user has left the log, /var/log/event/xcvrDefault40g,
   # untouched.  If that log is tampered with in any way, the speed change will not
   # be reverted.
   #
   # Overview:
   # First we parse the syslog to determine what action, removal or insertion, has
   # occured and on what port.  On insertion we spin up a CLI session and send the
   # 'speed' command.  On removal we parse the log starting from the end. If the
   # last event for the removed transceiver was an insertion, then we send the
   # 'no speed' command to remove the configuration.

   # The transceiver removal syslog will be stored in this environment variable.
   syslogMsg = os.environ.get( 'EVENT_LOG_MSG' )
   if syslogMsg is None:
      log( "Blank syslog message (syslogMsg==None). Exiting..." )
      sys.exit( 0 )

   log( "Starting...event log received: %s" % syslogMsg )

   action, port = parseLog( syslogMsg )
   log( f"Action: {action}, port: {port}" )
   if port is None or action is None:
      log( "Could not recognize parsed log, resulting in no-op. Exiting..." )
      sys.exit( 0 )
   else:
      actionStr = "'speed forced 40g'" if action == _inserted else "'no speed'"
      log( f"Preparing to potentially configure {actionStr} on port {port}" )
      if action == _inserted:
         doMaybeConfigure40g( port )
      else:
         doMaybeRemove40g( port )

   log( "Finished with port %s" % port )

   # Writes the log buffer to a file in one go
   flushToLog()
   sys.exit( 0 )

if __name__ == '__main__':
   run()
