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

import tempfile, Tac
import optparse, sys, re, os

###########################################################################
##### Hardcoded values for different supported platforms ##################
###########################################################################
# Dictionary containing cpu/flash architecture groups and 
# their member platforms. This maps platform names to an architecture
# group and is used when different actions are needed to be taken
# on different architectures.
# We now primarily use the "Aboot" cmdline option, but fall back on this
# for older Aboot versions.
platToArchGroup = {
      'norcal1' : [ ], # flashUtil was not used for norcal1
      'norcal2' : [ 'raven' ],
      'norcal3' : [ 'oak', 'blackbird', 'eaglepeak' ],
      'norcal4' : [ 'crow', 'mendocino' ],
      'norcal7' : [ 'oldfaithful' ],
      }

# Lookup Table for platforms and corresponding region offsets.
# This is a dict with a record for every key that is a platform.
# For every platform there is a dict with flash regions
# for which we describe start offset and length
# like so 'SECTION' : { 'start' : START, 'length' : LENGTH }
flashOffsetLUT = {
   'norcal2' : { 
               'total' :    { 'start' : 0x000000, 'length' : 0x0800000 },
               'prefdl' :   { 'start' : 0x3F0000, 'length' : 0x0010000 },  
               'image' :    { 'start' : 0x400000, 'length' : 0x03F0000 },
             },
   'norcal3' : {
               'total' :    { 'start' : 0x000000, 'length' : 0x0800000 },
               'prefdl' :   { 'start' : 0x010000, 'length' : 0x000f000 },
               'mfgdata' :  { 'start' : 0x01f000, 'length' : 0x0001000 },
               'fdl' :      { 'start' : 0x020000, 'length' : 0x0010000 },
               'image' :    { 'start' : 0x030000, 'length' : 0x03e0000 },
               'fallback' : { 'start' : 0x410000, 'length' : 0x03f0000 },
              },
   'norcal4' : {
               'total' :    { 'start' : 0x000000, 'length' : 0x0800000 },
               'prefdl' :   { 'start' : 0x010000, 'length' : 0x000f000 },
               'mfgdata' :  { 'start' : 0x01f000, 'length' : 0x0001000 },
               'fdl' :      { 'start' : 0x020000, 'length' : 0x0010000 },
               'fallback' : { 'start' : 0x030000, 'length' : 0x03f0000 },
               'image' :    { 'start' : 0x420000, 'length' : 0x03db000 },
              },
   # norcal5 isn't shipped as a flash image and will not use flashUtil 
   'norcal6' : {
               'total' :    { 'start' : 0x000000, 'length' : 0x1000000 },
               'mfgdata' :  { 'start' : 0x000400, 'length' : 0x0000400 },
               'fdl' :      { 'start' : 0x000800, 'length' : 0x0010000 },
               'prefdl' :   { 'start' : 0x010800, 'length' : 0x000f400 },
               'fallback' : { 'start' : 0x01fc00, 'length' : 0x06e0400 },
               'image' :    { 'start' : 0x700000, 'length' : 0x06fffc0 },
               'agesa' :    { 'start' : 0xdfffc0, 'length' : 0x0200040 },
              },
   'norcal7' : {
               'total' :    { 'start' : 0x000000, 'length' : 0x1000000 },
               'mac' :      { 'start' : 0x001000, 'length' : 0x0001000 },
               'prefdl' :   { 'start' : 0x002000, 'length' : 0x000f000 },
               'fdl' :      { 'start' : 0x011000, 'length' : 0x000f000 },
               'mfgdata' :  { 'start' : 0x020000, 'length' : 0x0001000 },
               'desc' :     { 'start' : 0x000000, 'length' : 0x0001000 },
               'pdr' :      { 'start' : 0x001000, 'length' : 0x0020000 },
               'me' :       { 'start' : 0x021000, 'length' : 0x03df000 },
               'fallback' : { 'start' : 0x400000, 'length' : 0x0600000 },
               'normal' :   { 'start' : 0xa00000, 'length' : 0x0600000 },
              }
   }

pseudoTargetLUT = {
   'norcal7' : {
               'image' : ( 'fallback', 'normal' ),
               },
   }

platformOptionsLUT = {
   # Don't read old contents before programming, it doesn't work with ME operating
   'norcal7' : [ '-A', ]
   }

###########################################################################
##### End hardcoded values section ########################################
###########################################################################

def run( cmd, verboseOut=False, stdoutSetting=Tac.CAPTURE ):
   if verboseOut:
      return Tac.run( cmd, asRoot=True )

   return Tac.run( cmd,
            stdout=stdoutSetting,
            stderr=Tac.DISCARD,
            asRoot=True )

def exitWithPrint( code, err=None, msg=None ):
   if msg:
      sys.stdout.write( msg + '\n' )
   if err:
      sys.stderr.write( err + '\n' )
   sys.exit( code )
   
op = optparse.OptionParser( prog='flashUtil', usage='''
   %prog -- Wrapper Utility around flashrom to perform read
            write action on NorCal SPI flash.

   %prog [-v] [-n] [-p PLATFORM] <-r|w SECTION> <FILENAME>
   
   NorCal1 support is limited and should only be used with the 'total'
   SECTION option or none.
   On NorCal2 and up, the script can be directed to read or
   write any given SECTION.
   During writes <filename> provides the source data and during reads
   data is written to it.
   Caution should be used with the total option as the system can
   be rendered unbootable if a write does not complete.

   Section can be:
     total ( entire flash image )
     prefdl
     fdl
     mac
     image ( aboot )
     fallback

   * not all sections are supported on all platforms.
   ''' )
   
op.add_option( "-r", "--read", action='store', default=None,
      help="Read from Flash" )

op.add_option( "-w", "--write", action='store', default=None,
      help="Write to Flash" )

op.add_option( "-a", "--archGroup", action='store', default=None,
      help="Specify an archGroup (e.g. norcalN)" )

op.add_option( "-n", "--flash-name", dest='name', 
         action='store_true', default=None, help="Only probe for flash chip name" )

op.add_option( "-v", "--verbose", dest='verbose',
         action='store_true', default=None, help="verbose output" )

opts, args = op.parse_args()

# For debugging purposes, to test different flashroms
flashrom = "flashrom"

if opts.name:
   out = run( [ flashrom ] )
   reRes = re.search( r"flash chip \"(.*?)\"", out )
   if reRes:
      exitWithPrint( 0, msg=reRes.group( 1 ) )
   exitWithPrint( 0 )
   
read = opts.read
write = opts.write

# Check right num args
if len( args ) != 1:
   op.print_usage()
   exitWithPrint( 1, err="Invalid argument count." )

# Check that either read or write is specified
if read and not write:
   target = read
elif write and not read:
   flashrom = "flashrom-diag"
   target = write
else:
   op.print_usage()
   exitWithPrint( 1, err="Invalid option." )

fileArg = args[ 0 ]

# Read platform and archGroup from Aboot command line

platform = None
archGroup = None
try:
   cmdline = open( '/proc/cmdline', 'r' ).read()
except IOError:
   pass
else:
   platformRe  = re.search( r"platform=(.*?)($| |\n)", cmdline )
   archGroupRe  = re.search( r"Aboot=Aboot-(\w+)-", cmdline )
   if platformRe:
      platform = platformRe.group( 1 )
   if archGroupRe:
      archGroup = archGroupRe.group( 1 ).lower() 
      # In case some Aboot version strings cause us to get garbage
      if archGroup not in flashOffsetLUT.keys():
         archGroup = None

# Allow user to override
if opts.archGroup:
   archGroup = opts.archGroup

# Older versions of Aboot did not include the archgroup in
# /proc/cmdline, so fall back on using the platform (eventually
# we should get rid of this code)
if not archGroup and platform:
   for key in platToArchGroup.keys():
      if platform in platToArchGroup[ key ]:
         archGroup = key
         break

if not archGroup:
   exitWithPrint( 1,
         err="Cannot determine archGroup, please supply one with -a option." )

pod = flashOffsetLUT[ archGroup.lower() ] # Platform Offset Dictionary
pseudoTargetDict = pseudoTargetLUT.get( archGroup.lower(), {} )

if 'total' not in pod:
   exitWithPrint( 1, err="'total' section layout must be hardcoded for archGroup %s."
         % ( archGroup ) )

# Check that we have a valid target section
if target in pod:
   targets = ( target, )
elif target in pseudoTargetDict:
   if read:
      exitWithPrint( 1, err="PseudoTarget is not supported for read" )
   targets = pseudoTargetDict[ target ]
   valid = set( targets ).issubset( pod.keys() ) and len( targets ) > 1
   if not valid:
      exitWithPrint( 1, 
                     err="PseudoTarget %s is invalid on archGroup %s." 
                     % ( target, archGroup ) )
else:
   exitWithPrint( 1, err="Target section %s is not supported for archGroup %s." 
         % ( target, archGroup ) )


totalLen = pod[ 'total' ][ 'length' ]

# Get platform-specific options for flashrom
platformOptions = platformOptionsLUT.get( archGroup.lower(), [] )

def writeTarget( region, layoutFile ):
   start = pod[ region ][ 'start' ]
   length = pod[ region ][ 'length' ]
   # Check input file size
   imageSize = os.path.getsize( fileArg )
   if imageSize != totalLen:
      if imageSize > length:
         exitWithPrint( 1, err="Image size doesn't match for target section %s (%d)."
                        % ( region, length ) )
      # first we make the image patch
      # All regions should be aligned at the 1KB boundary, if that is not the 
      # case, we need adjust the block size we used.
      blockSize = 1024
      while blockSize > 1:
         if totalLen % blockSize == 0 and start % blockSize == 0:
            break
         blockSize = blockSize / 2
      
      tmpFile = tempfile.NamedTemporaryFile()
      
      run( [ 'dd', 'if=/dev/zero', 'of=' + tmpFile.name,
             'bs=%d' % blockSize, 'count=%d' % ( totalLen / blockSize ) ],
           opts.verbose )
      run( [ 'dd', 'if=' + fileArg, 'of=' + tmpFile.name,
             'bs=%d' % blockSize, 'seek=%d' % ( start / blockSize ), 
             'conv=notrunc' ], opts.verbose )
      cmd = [ flashrom, '-n', '-l', layoutFile, '-i', target, '-w',
             tmpFile.name ]
      cmd += platformOptions
      run( cmd, opts.verbose )
   else:
      cmd = [ flashrom, '-n', '-l', layoutFile, '-i', target, '-w',
             fileArg ]
      cmd += platformOptions
      run( cmd, opts.verbose )

def writePseudoTarget( regions, layoutFile ):
   if os.path.getsize( fileArg ) != totalLen:
      exitWithPrint( 1, err="PseudoTarget needs a full-size image" )

   cmd = [ flashrom, '-n', '-l', layoutFile, ]
   for reg in regions:
      cmd += [ '-i', reg, ]
   cmd += [ '-w', fileArg ]
   cmd += platformOptions
   run( cmd, opts.verbose )

def readTarget( region ):
   start = pod[ target ][ 'start' ]
   length = pod[ target ][ 'length' ]
   cmd = [ flashrom, '--start=%d' % start, '--numBytes=%d' % length,
          '-r', fileArg ]
   cmd += platformOptions
   # if fileArg is stdout, add the quiet option to flashrom
   stdoutSetting = Tac.CAPTURE
   if fileArg == '-':
      cmd += [ '-q' ]
      stdoutSetting = Tac.INHERIT
   run( cmd, opts.verbose, stdoutSetting )

if read:
   # Operation is a READ
   readTarget( target )
else:
   # Operation is a WRITE
   # Create layout file
   layout = tempfile.NamedTemporaryFile()
   for key in pod:
      layout.write( "%X:%X %s\n"
            % ( pod[ key ][ 'start' ], 
                pod[ key ][ 'length' ] + pod[ key ][ 'start' ] - 1,
                key ) )
   layout.flush()
   # Program the target(s)
   if len( targets ) == 1:
      writeTarget( target, layout.name )
   else:
      writePseudoTarget( targets, layout.name )

