#!/usr/bin/env python3
# Copyright (c) 2007 Arastra, Inc.  All rights reserved.
# Arastra, Inc. Confidential and Proprietary.
#
# This script stores the existing contents of the BIOS flash,
# tests read, write and erase functionality, and restores the
# contents.
import os, traceback, Tac

print( "Reading the existing flash contents" )
Tac.run( ["flashrom", "-q", "-r", "/tmp/bios.rom"] )

print( "Making sure we can read the same contents again" )
Tac.run( ["flashrom", "-q", "-v", "/tmp/bios.rom"] )

print(
   "Determining the size of the flash and creating all-zero and all-one contents" )
flashsize = os.stat( "/tmp/bios.rom" ).st_size
# pylint: disable-next=consider-using-with
open( "/tmp/zero.rom", "w" ).write( "\x00\x00\x00\x00" * ( flashsize//4 ) )
# pylint: disable-next=consider-using-with
open( "/tmp/one.rom", "w" ).write( "\xff\xff\xff\xff" * ( flashsize//4 ) )
testPass = True

try:
   # Test that we can at least write all ones and zeros to flash
   print( "Checking that the flash image can be written all zeros" )
   Tac.run( ["flashUtil", "-w", "total", "/tmp/zero.rom"] )
   Tac.run( ["flashUtil", "-r", "total", "/tmp/zeroOut.rom"] )
   Tac.run( [ 'diff', '/tmp/zeroOut.rom', '/tmp/zero.rom' ] )
   
   print( "Verifying that the flash image can be written all ones" )
   Tac.run( ["flashUtil", "-w", "total", "/tmp/one.rom"] )
   Tac.run( ["flashUtil", "-r", "total", "/tmp/oneOut.rom"] )
   Tac.run( [ 'diff', '/tmp/oneOut.rom', '/tmp/one.rom' ] )
   
except: # pylint: disable=bare-except
   # Print a backtrace here so it's clear where the failure occurred
   traceback.print_exc()
   testPass = False

finally:
   # Try to restore the original contents of the flash, whether the above
   # failed or succeeded.

   print( "Writing the original flash contents and verifying them" )
   Tac.run( ["flashUtil", "-w", "total", "/tmp/bios.rom"] )
   Tac.run( ["flashrom", "-q", "-v", "/tmp/bios.rom"] )
   assert testPass, "biosFlashrom diagnostic test failed"
   print( "Test completed successfully" )
