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

# pylint: disable=consider-using-sys-exit

"""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:

from __future__ import absolute_import, division, print_function

import email
import six
import sys

import Fdl
from GenfdlLib import (
   FdlRegistry,
   getFdlPaths,
)
import GenprefdlLib

def exit( message ): # pylint: disable=redefined-builtin
   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 # pylint: disable=import-outside-toplevel
   if fdlDateTime:
      # pylint: disable-next=consider-using-f-string
      return "%s-%s-%s" % (fdlDateTime[0:4],fdlDateTime[4:6],fdlDateTime[6:8])
   else:
      return datetime.date.today().strftime( "%Y-%m-%d" )

def genfdl( prefdlBytes, mac=None, force=False, fileName=None, fdldir=None,
            raw=False, rawPrefdl=None, compress=False, supervisorUplink=False,
            useSwagFdl=False, useSwagLinecardFdl=False ):
   msg = ( email.message_from_string( prefdlBytes ) if six.PY2
           else email.message_from_bytes( prefdlBytes ) ) # pylint: disable=no-member
   if mac:
      msg['MacAddrBase'] = mac
   if 'CRC' in msg and 'nvalid' in msg and not force:
      exit( "Input contains a bad crc but --force not specified" )

   prefdlFields = GenprefdlLib.decodePrefdl( prefdlBytes )
   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'
   # For SWAG prototyping, we are maintaining a separate FDL which is "SWAG-ified".
   # We will select which FDL to use by appending "Swag" to the end of the SID,
   # similar to the mechanism used for Tundra supervisor uplinks.
   # BUG968909: We should move away from separate FDL files for SWAG mode.
   if useSwagFdl:
      sid += 'Swag'
   elif useSwagLinecardFdl:
      sid += 'SwagLinecard'
   msg[ 'Sid' ] = sid
   hwApi = prefdlFields.hwApi
   cpuPca = prefdlFields.cpuPca
   cpuSid = prefdlFields.cpuSid
   cpuHwApi = prefdlFields.cpuHwApi
   configDict = prefdlFields.configDict

   if fileName:
      fdlName = fileName.split( ',' )
   else:
      registry = FdlRegistry()
      fdlName = registry.lookupFdl( sid, baseAsy, hwApi,
                                    cpuSid=cpuSid, cpuHwApi=cpuHwApi,
                                    configDict=configDict )
      if not fdlName:
         exit(
            # pylint: disable-next=consider-using-f-string
            "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( b"" )
   fdlKeys = set()
   fdlKeysLower = set()
   fdlFactoryData = ""
   fruFactoryIds = None
   fdlVariables = None
   swFeatures = None
   for fdlPath in fdlPaths:
      with open( fdlPath, "rb" ) as fdlFile:
         fdlComponent = Fdl.Rfc822Fdl( fdlFile.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:
         fruFactoryIds = fdlComponent['FruFactoryIds']
      if 'FdlVariables' in fdlComponent:
         fdlVariables = fdlComponent['FdlVariables']
      if 'SwFeatures' in fdlComponent:
         swFeatures = fdlComponent[ 'SwFeatures' ]

   msgKeys = { i.lower() for i in msg.keys() }

   keysWeProvide = set( [ 'mfgdate', 'cpumfgdate' ] )
   keysToIgnore = set( [ 'frufactoryids', 'deviations', 'fdlvariables',
                         'hwapi', 'cpudeviations', 'swfeatures' ] )
   # PCA is not a required field in DataFormat 0003, ignore it if not provided
   if pca == "":
      keysToIgnore.add( "pca" )
   elif 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:
      # pylint: disable-next=consider-using-f-string
      exit("The selected FDL file (%s) requires the"
           " following keys to be provided: %s\n"
           " (I have %s)\n"%
           (",".join(fdlPaths),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 )
   cpuMfgTime = msg.get( "CpuMfgTime2" )
   if not cpuMfgTime:
      cpuMfgTime = msg.get( "CpuMfgTime" )
   fdlCpuDate = dateString( cpuMfgTime )

   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 'CpuMfgDate' in fdlKeys:
      fdl[ 'CpuMfgDate' ] = fdlCpuDate
   # 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 )
   fdlBytes = fdl.fdlBytes()

   # Generate a raw fdl if the user asks
   if raw:
      import Tac # pylint: disable=import-outside-toplevel
      dataFormat = 2 if compress else 1
      gen = Tac.newInstance( "Fdl::Generator", fdlBytes, dataFormat,
                             six.ensure_binary( rawPrefdl or b"" ) )
      if gen.fdlError:
         # pylint: disable-next=consider-using-f-string
         exit( "Error compressing fdl: %s" % gen.fdlError )
      fdlBytes = gen.rawFdl.byteStr
   return fdlBytes

def main():
   import optparse # pylint: disable=import-outside-toplevel,deprecated-module
   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" )
   # BUG968909: We should move away from separate FDL files for SWAG mode.
   parser.add_option( "--useSwagFdl", action="store_true",
                      help="Generate the SWAG version of the system's FDL" )
   parser.add_option( "--useSwagLinecardFdl", action="store_true",
                      help="Generate the SWAG linecard version of the system's FDL" )
   ( options, args ) = parser.parse_args()

   stdinBuffer = ( sys.stdin if six.PY2
      else sys.stdin.buffer ) # pylint: disable=no-member

   if args:
      infile = open( args[0], "rb" ) # pylint: disable=consider-using-with
   else:
      infile = stdinBuffer

   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." )

   stdoutBuffer = ( sys.stdout if six.PY2
      else sys.stdout.buffer ) # pylint: disable=no-member

   if options.output:
      outfile = open( options.output, "wb" ) # pylint: disable=consider-using-with
   else:
      outfile = stdoutBuffer
   prefdlBytes = infile.read()
   if infile is not stdinBuffer:
      infile.close()
   fdlBytes = genfdl( prefdlBytes, mac=options.mac, force=options.force,
         fileName=options.file, fdldir=options.fdldir, raw=options.raw,
         rawPrefdl=options.prefdl, compress=options.compress,
         supervisorUplink=options.supervisorUplink, useSwagFdl=options.useSwagFdl,
         useSwagLinecardFdl=options.useSwagLinecardFdl )
   outfile.write( fdlBytes )
   if outfile is not stdoutBuffer:
      outfile.close()
if __name__ == "__main__":
   main()
