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

"""This script is used primarily during manufacturing.  It generates
an unsigned fdl when given a prefdl that defines the Serial number,
PCA, and mac address.  The mac address can be specified on the command
line as well. The Serial number, Pca, HwRev, MfgDate, and MacAddrBase
fields in the fdl template file will be replaced with the fields from
the user-provided prefdl."""

# Need to get the following RFC822 headers on a helena fdl:
#   macaddr:
#   serialnumber:
#   mfgdate:
import email, sys, os.path, os
import GenprefdlLib, Fdl
from GenfdlLib import FdlRegistry, getFdlPaths

def exit( message ):
   sys.stderr.write( message )
   sys.exit( 1 )

def validMacAddress( address ):
   if '-' in address:
      a = address.split( '-' )
   elif ':' in address:
      a = address.split( ':' )
   else:
      return False
   if len( a ) != 6:
      return False
   try:
      for i in a:
         val = int(i,16)
         if val > 255 or val < 0:
            return False
   except ValueError:
      return False
   return True

def validateMacAddressFormat( address ):
   if not validMacAddress( address ):
      exit( "Invalid MAC address format; must be"
            "xx:xx:xx:xx:xx:xx or xx-xx-xx-xx-xx-xx" )

def dateString( fdlDateTime=None ):
   import datetime
   if fdlDateTime:
      return "%s-%s-%s" % (fdlDateTime[0:4],fdlDateTime[4:6],fdlDateTime[6:8])
   else:
      return datetime.date.today().strftime( "%Y-%m-%d" )

def genfdl( prefdlStr, mac=None, force=False, fileName=None, fdldir=None, raw=False,
            rawPrefdl=None, compress=False, supervisorUplink=False ):
   msg = email.message_from_string( prefdlStr )
   if mac:
      msg['MacAddrBase'] = mac
   if msg.has_key('CRC') and 'nvalid' in msg and not force:
      exit( "Input contains a bad crc but --force not specified" )

   prefdlFields = GenprefdlLib.decodePrefdl( prefdlStr )
   pca = prefdlFields.pca
   hwRev = prefdlFields.hwRev
   msg['HwRev'] = hwRev
   basePca = prefdlFields.basePca
   baseAsy = prefdlFields.baseAsy
   sid = prefdlFields.sid
   # The supervisorUplink option addresses the special case where a supervisor has
   # uplink capabilities and is segmented into 2 parts - supervisor and uplink.
   # Each segment has a unique FDL to be generated and has its own idprom. However,
   # to simplify MFG, identical prefdl contents will be programmed into both idproms.
   # This option provides the additional context necessary to generate an alternate
   # uplink FDL from the same supervisor prefdl. It does this by simply appending
   # "Uplink" to the SID, which can be used to uniquely key in the fdlMap.
   #
   # i.e. fdlMap = {
   #         ( "Narwhal",       "01" ) : "NarwhalP1",
   #         ( "NarwhalUplink", "01" ) : "NarwhalUplinkP1"
   #      }
   if supervisorUplink:
      sid += 'Uplink'
   msg[ 'Sid' ] = sid
   hwApi = prefdlFields.hwApi
   cpuPca = prefdlFields.cpuPca
   cpuSid = prefdlFields.cpuSid
   cpuHwApi = prefdlFields.cpuHwApi
   configDict = prefdlFields.configDict

   if fileName:
      fdlPaths = fileName.split( ',' )
      for fdlPath in fdlPaths:
         if not os.path.exists( fdlPath ):
            exit( "File requested with --file does not exist: " + fdlPath )
   else:
      registry = FdlRegistry()
      fdlName = registry.lookupFdl( sid, baseAsy, hwApi,
                                    cpuSid=cpuSid, cpuHwApi=cpuHwApi,
                                    configDict=configDict )
      if not fdlName:
         exit(
            "No fdl found for SID %s hwApi %s (base PCA %s and base ASY %s,"
            " cpu SID %s PCA %s cpuHwApi %s),"
            % ( sid, hwApi, basePca, baseAsy, cpuSid, cpuPca, cpuHwApi ) )
      fdlPaths = getFdlPaths( fdlName, fdldir=fdldir )

   fdl = Fdl.Rfc822Fdl( "" )
   fdlKeys = set()
   fdlKeysLower = set()
   fdlFactoryData = ""
   fruFactoryIds = None
   fdlVariables = None
   swFeatures = None
   for fdlPath in fdlPaths:
      fdlComponent = Fdl.Rfc822Fdl( file( fdlPath ).read() )
      # Get the fdl keys from each of the fdls
      fdlKeysLower = fdlKeysLower.union(fdlComponent.keysLower())
      fdlKeys = fdlKeys.union(fdlComponent.keys())
      fdlFactoryData += fdlComponent.factoryData()
      if 'FruFactoryIds' in fdlComponent.keys():
         fruFactoryIds = fdlComponent['FruFactoryIds']
      if 'FdlVariables' in fdlComponent.keys():
         fdlVariables = fdlComponent['FdlVariables']
      if 'SwFeatures' in fdlComponent.keys():
         swFeatures = fdlComponent[ 'SwFeatures' ]

   msgKeys = set([i.lower() for i in msg.keys()])

   keysWeProvide = set(['mfgdate'])
   keysToIgnore = set( [ 'frufactoryids', 'deviations', 'fdlvariables',
                         'hwapi', 'cpudeviations', 'swfeatures' ] )
   if pca.startswith( "PCA00011" ) or pca.startswith( "PCA00003" ):
      # Earlier versions of sonomas and helenas didn't get programmed
      # with a sku in their idprom. We deal with this below
      keysToIgnore.add( "sku" )
   missingKeys = fdlKeysLower - msgKeys - keysToIgnore - keysWeProvide
   if missingKeys:
      exit("The selected FDL file (%s) requires the"
           " following keys to be provided: %s\n"
           " (I have %s)\n"%
           (fdlPath,missingKeys,msgKeys))

   if 'MacAddrBase' in msgKeys:
      validateMacAddressFormat( msg['MacAddrBase'] )

   # Now copy headers to the fdl
   mfgTime = msg.get( "MfgTime2" )
   if not mfgTime:
      mfgTime = msg.get( "MfgTime" )
   fdlDate = dateString( mfgTime )

   for i in fdlKeys:
      if msg[i.lower()]:
         fdl[i] = msg[i.lower()]
   if fruFactoryIds:
      fdl['FruFactoryIds'] = fruFactoryIds
   if 'MfgDate' in fdlKeys:
      fdl['MfgDate'] = fdlDate
   # If fdlVariables doesn't exist in the message, copy it over from
   # the original fdl
   if not 'fdlvariables' in msgKeys and fdlVariables:
      fdl['FdlVariables'] = fdlVariables
   if not 'swfeatures' in msgKeys and swFeatures:
      fdl[ 'SwFeatures' ] = swFeatures
   sku = msg.get( "SKU" )
   if not sku:
      # Earlier versions of sonomas and helenas didn't get programmed
      # with a sku in their idprom. Add it manually here.
      if pca.startswith( "PCA00011" ):
         sku = "DCS-7148SX"
      elif pca.startswith( "PCA00003" ):
         sku = "DCS-7124S"
   # Deal with the fact that we changed the M8 skus at the last minute
   if sku == "DCM-7508":
      sku = "DCS-7508"
   elif sku in [ "SUP-7500", "SUP1-7500" ]:
      sku = "7500-SUP"
   elif sku == "LC-7548S":
      sku = "7548S-LC"
   elif sku == "FM-7508":
      sku = "7508-FM"
   elif sku and ( sku.startswith( "DCS-75" ) or sku.startswith( "75" ) ):
      sku = fdlComponent[ "Sku" ]
   if sku:
      fdl[ 'Sku' ] = sku

   if sid:
      fdl[ 'Sid' ] = sid
   if hwApi:
      fdl[ 'HwApi' ] = hwApi

   # Add the HwEpoch if it is present in the prefdl but not the fdl
   hwEpoch = msg.get( "HwEpoch" )
   if hwEpoch:
      fdl[ 'HwEpoch' ] = hwEpoch

   fdl[ 'Deviations' ] = ", ".join( [ v for (k,v) in msg.items() \
                                         if k=="Deviation" ] )

   if "CpuDeviation" in msg.items():
      fdl[ 'CpuDeviations' ] = ", ".join( [ v for (k,v) in msg.items() \
                                                    if k=="CpuDeviation" ] )

   # output variable part of fdl (macaddr, serialnumber, kvn, mfgdate)
   # concat fixed part of fdl.
   fdl.factoryDataIs( fdlFactoryData )
   fdlText = fdl.fdlStr()

   # Generate a raw fdl if the user asks
   if raw:
      import Tac
      dataFormat = 2 if compress else 1
      gen = Tac.newInstance( "Fdl::Generator", fdlText, dataFormat,
                             rawPrefdl or "" )
      if gen.fdlError:
         exit( "Error compressing fdl: %s" % gen.fdlError )
      fdlText = gen.rawFdl.str
   return fdlText

def main():
   import optparse
   usage = "%prog [-m <mac address>] [-o <outfile>] [-z] [-f] [--raw] [infile]"
   description = """
   Take a prefdl (in rfc822 format) as input and generate a fdl as output.

   %prog is intended for manufacturing and internal use.  It generates an
   (unsigned but possibly compressed) fdl. The input is a sequence of
   lines of text in rfc822 header format, defining the following fields:

      SerialNumber
      Pca
      KVN
      Sku
      MacAddrBase
      FdlVariables

   The mac address may optionally be specified on the command line with
   the -m option instead of in the input file.  The PCA number in the
   input is used to select a FDL template file.  The Serial number, Kvn,
   Pca, HwRev, MfgDate, and MacAddrBase fields in the fdl template file
   are replaced in the output by the user-provided input fields.  If no
   infile is specified, input comes from stdin.  The generated fdl is
   optionally compressed.
   """

   class HelpFormatter(optparse.IndentedHelpFormatter):
      def format_description(self, description):
         return description or ""

   parser = optparse.OptionParser(usage=usage,description=description,
                                  formatter=HelpFormatter())
   parser.add_option( "-m", "--mac", action="store",
                      help="First mac address for this fdl, in "
                      "xx:xx:xx:xx:xx:xx or xx-xx-xx-xx-xx-xx format" )
   parser.add_option( "-o", "--output", action="store",
                      help="Name of file to output to (default is stdout)" )
   parser.add_option( "-z", "--compress", action="store_true",
                      help="Generate a bz2-compressed fdl" )
   parser.add_option( "", "--prefdl", action="store",
                      help="Raw prefdl to encode into the fdl. Only valid for "
                           "raw fdl generation" )
   parser.add_option( "-r", "--raw", action="store_true",
                      help="Generate a raw FDL, suitable for directly programming in"
                      " a seeprom device, instead of text" )
   parser.add_option( "-f", "--force", action="store_true",
                      help="Generate a FDL even if the prefdl crc is invalid" )
   parser.add_option( "--file", action="store",
                      help="Force a specific FDL template file" )
   parser.add_option( "--fdldir", action="store",
                      help="Force lookup in a specific FDL directory" )
   parser.add_option( "--supervisorUplink", action="store_true",
                      help="Generate a FDL for the supervisor uplink" )
   ( options, args ) = parser.parse_args()

   if args:
      infile = file( args[0] )
   else:
      infile = sys.stdin

   if options.compress and not options.raw:
      parser.error( "compression (-z) cannot be requested without -r/--raw." )

   if options.prefdl and not options.raw:
      parser.error( "prefdl cannot be specified without -r/--raw." )

   if options.output:
      outfile = file( options.output, "w" )
   else:
      outfile = sys.stdout
   prefdlStr = infile.read()
   fdlText = genfdl( prefdlStr, mac=options.mac, force=options.force,
         fileName=options.file, fdldir=options.fdldir, raw=options.raw,
         rawPrefdl=options.prefdl, compress=options.compress,
         supervisorUplink=options.supervisorUplink )
   outfile.write( fdlText )
if __name__ == "__main__":
   main()
