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

# Reboot script for 7050QX-32

from __future__ import print_function
import sys
import mmap, os
import subprocess
from struct import pack, unpack

class MmapResource( object ):
   """Resource implementation for a directly-mapped memory region."""

   def __init__( self, path ):
      try:
         fd = os.open( path, os.O_RDWR )
      except EnvironmentError:
         print( "FAIL can not open scd memory-map resource file" )
         print( "FAIL are you running on the proper platform?" )
         sys.exit( 1 )
      try:
         size = os.fstat( fd ).st_size
      except EnvironmentError:
         print( "FAIL can not fstat scd memory-map resource file" )
         print( "FAIL are you running on the proper platform?" )
         sys.exit( 1 )
      try:
         self.mmap_ = mmap.mmap( fd, size, mmap.MAP_SHARED,
                      mmap.PROT_READ | mmap.PROT_WRITE )
      except EnvironmentError:
         print( "FAIL can not map scd memory-map file" )
         print( "FAIL are you running on the proper platform?" )
         sys.exit( 1 )
      finally:
         try:
            # Note that closing the file descriptor has no effect on the memory map
            os.close( fd )
         except EnvironmentError:
            print( "FAIL failed to close scd memory-map file" )
            sys.exit( 1 )

   def read32( self, addr ):
      return unpack( '<L', self.mmap_[ addr : addr + 4 ] )[ 0 ]

   def write32( self, addr, value ):
      self.mmap_[ addr: addr + 4 ] = pack( '<L', value )

def scdRegTest( scd, offset, val1, count ):
   scd.write32( offset, val1 )
   val2 = scd.read32( offset )
   if val1 != val2:
      print( "FAIL: scd write 0x%08x but read back 0x%08x in iter %d" %
             ( val1, val2, count ) )
      sys.exit( 17 )

def scdScrRegTest( scd ):
   scrOffset = 0x0130
   for i in range( 0, 3 ):
      scdRegTest( scd, scrOffset, 0xdeadbeef, i )
      scdRegTest( scd, scrOffset, 0xa5a5a5a5, i )
      scdRegTest( scd, scrOffset, 0x00000000, i )

def reboot( scd ):
   # reboot the system by writing to register 0x7000
   print( "Rebooting" )
   scd.write32( 0x7000, 0xDEAD )
   print( "REBOOTED" )

def main():
   busName = "/sys/bus/pci/devices/0000:04:00.0/resource0"

   subprocess.call( [ 'modprobe', 'scd' ] )

   scd = MmapResource( busName )

   #
   # verify that we can read/write scd scratch register
   #
   scdScrRegTest( scd )

   # reboot the system
   reboot( scd )

if __name__ == "__main__":
   main()