#!/usr/bin/env python
# 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
file( "/tmp/zero.rom", "w" ).write( "\x00\x00\x00\x00" * (flashsize/4) )
file( "/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:
   # 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"
