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

# This script reads/writes values from/to the PCIe CAP_EXP structure. The format
# of the structure is defined by the PCIe standard. The location within the device's
# memory is determined by the chip manufacture. The script will walk the list of
# headers until it finds the CAP_EXP then read or write the requested location.

import sys
import argparse
import PciUtil

# Each entry is a tuple of:
# ( read function, format width, write function, max value )
widthInfo = { 'b' : ( PciUtil.read8, 2, PciUtil.write8, 0xff ),
              'w' : ( PciUtil.read16, 4, PciUtil.write16, 0xffff ),
              'l' : ( PciUtil.read32, 8, PciUtil.write32, 0xffffffff ) }

owvSyntax = "Syntax is <offset[width][=value]>. Width is one of .B, .W or .L."
owvError = "Unable to parse '%s'. " + owvSyntax

class OffsetWidthValue( argparse.Action ):
   def __init__( self, option_strings, dest, nargs=None, **kwargs ):
      if nargs is not None:
         raise ValueError( "nargs not allowed" )
      super().__init__( option_strings, dest, **kwargs )

   def __call__( self, parser, namespace, values, option_string=None ):
      # Decode the offset and optional width.
      fields = values.split( '=' )
      regInfo = fields[ 0 ].split( '.' )
      rOffset = PciUtil.strToInt( regInfo[ 0 ] )
      if len( regInfo ) == 1:
         width = 'l'
      elif len( regInfo ) == 2:
         width = regInfo[ 1 ].lower()
      else:
         width = None
      if width not in widthInfo:
         parser.error( owvError % values )

      # Decode and verify the optional new value
      if len( fields ) == 1:
         newVal = None
      elif len( fields ) == 2:
         newVal = PciUtil.strToInt( fields[ 1 ] )
      else:
         parser.error( owvError % values )
      maxVal = widthInfo[ width ][ 3 ]
      if newVal is not None and newVal > maxVal:
         # pylint: disable-next=consider-using-f-string
         parser.error( "New value too large for field width; 0x%x > 0x%x" %
                       ( newVal, maxVal ) )

      # Save the values.
      setattr( namespace, self.dest, ( rOffset, widthInfo[ width ], newVal ) )

argParser = argparse.ArgumentParser( description="Read or write a field (byte, word "
                                     "or long) in a PCI device's PCI Express "
                                     "Capability Structure. The field is specified "
                                     "as an offset in bytes." )
argParser.add_argument( "pciAddress", help='PCI address, e.g. 0000:00:1f.0' )
argParser.add_argument( "offsetValue", action=OffsetWidthValue, help=owvSyntax )

try:
   args = argParser.parse_args()
except ValueError as msg:
   argParser.error( str( msg ) )

pciId = args.pciAddress
regOffset, wInfo, newValue = args.offsetValue
readFunc, fmtWidth, writeFunc, maxValue = wInfo

# Read offset 0x34 (as defined by the PCIe standard) to get the offset to the next
# header.
hdrOffset = PciUtil.read8( pciId, "config", 0x34 )

# Read 32 bits at the header offset. If the low 8 bits are the ID for CAP_EXP,
# then stop. Otherwise get the new header offset from the value that was just
# read. An offset of 0x00 indicates the end of the list.
while True:
   value = PciUtil.read32( pciId, "config", hdrOffset )
   if ( value & 0xff ) == 0x10:
      break
   hdrOffset = ( value >> 8 ) & 0xff
   if hdrOffset == 0:
      print( 'CAP_EXP header not found', file=sys.stderr )
      sys.exit( 1 )

# Found the CAP_EXP header
addr = hdrOffset + regOffset
if newValue is not None:
   writeFunc( pciId, "config", addr, newValue )
else:
   # pylint: disable-next=consider-using-f-string
   print( "0x%0*x" % ( fmtWidth, readFunc( pciId, "config", addr ) ) )
