#!/usr/bin/env python
# Copyright (c) 2007-2011 Arastra, Inc.  All rights reserved.
# Arastra, Inc. Confidential and Proprietary.

import optparse, sys, os, fcntl, time
import EosDisk

op = optparse.OptionParser( usage="usage: %prog [OPTIONS] SWI" )
op.add_option( "--device", action="store",
               help="set up DEVICE instead of default flash device" )
op.add_option( "--fs", action="store", default="vfat",
               help="filesystem types (vfat or ext4)" )
op.add_option( "--force", action="store_true",
               help="proceed even if device contains data" )
op.add_option( "--test", action="store_true",
               help="test the device before writing to it" )
op.add_option( "--boot-config", action="store",
               help="use BOOT_CONFIG instead of default" )
op.add_option( "--startup-config", action="store",
               help="use STARTUP_CONFIG" )
op.add_option( "--copy-tree-from", action="store",
               help="Copy files (recursively) from COPY_TREE_FROM to flash root " \
                    "after setup" )
op.add_option( "--recovery-swi", action="store",
               help="Use this swi in recovery partition" )
op.add_option( "--add-to-recovery", action="store",
               help="Copy files (recursively) from ADD_TO_RECOVERY to recovery " \
                    "disk and thereby to flash root" )
op.add_option( "--flash-size", action="store", type='int', default=0, metavar='SIZE',
               help="Create a flash partition of SIZE MB. By default all the " \
                    "space is used (SIZE=0)" )

opts, args = op.parse_args()

if len( args ) != 1:
   op.error( "no SWI specified" )

if not opts.device:
   opts.device = EosDisk.blockDeviceNode( "flash" )
   if not opts.device:
      sys.stderr.write( "Could not find the block device associated with the " \
                        "the internal flash.\nCheck that the internal flash is " \
                        "properly inserted and was detected by linux.\n" )
      sys.exit( 1 )

if opts.fs.lower() not in [ "vfat", "ext4" ]:
   sys.stderr.write( "Unknown file system type [%s]\n" % opts.fs )
   sys.stderr.write( "Try 'vfat' or 'ext4'" )
   sys.exit( 1 )

if not os.path.exists( opts.device ):
   sys.stderr.write( "Could not find device %s\n" % opts.device )
   sys.exit( 1 )

mounted = EosDisk.mountedHd( opts.device )
if mounted:
   sys.stderr.write( "%s is mounted (%s)\n" % (opts.device, mounted) )
   sys.exit( 1 )

if not opts.force:
   if not EosDisk.emptyHd( opts.device ):
      sys.stderr.write( "%s seems to contain data; use --force to skip check\n"
                        % opts.device )
      sys.exit( 1 )

if opts.copy_tree_from and not os.path.exists( opts.copy_tree_from ):
   sys.stderr.write( "Path '%s' does not exist" % opts.copy_tree_from )
   sys.exit( 1 )

if opts.add_to_recovery and not os.path.exists( opts.add_to_recovery ):
   sys.stderr.write( "Path '%s' does not exist" % opts.add_to_recovery )
   sys.exit( 1 )

swi = args[0]
if not os.path.exists( swi ):
   sys.stderr.write( "File '%s' does not exist" % swi )
   sys.exit( 1 )

if opts.recovery_swi and not os.path.exists( opts.recovery_swi ):
   sys.stderr.write( "Path '%s' does not exits" % opts.recovery_swi )
   sys.exit( 1 )

if opts.test:
   print "Testing '%s'" % opts.device
   size = EosDisk.diskSizeBytes( opts.device )
   for v in [ '\x00', '\x55', '\xaa', '\xff' ]:
      print "Writing 0x%x to the flash" % ord(v)
      f = file( opts.device, "w" )
      for i in range( size / 1024 ):
         f.write( v * 1024 )
      f.flush()

      print "Verifying the contents of the flash"
      f = file( opts.device, "r" )
      for i in range( size / 1024 ):
         buf = f.read( 1024 )
         if buf != v * 1024:
            print "Device %s seems to have stuck bits" % opts.device
            print "Contents incorrect in byte range", i * 1024, "to", (i+1) * 1024
            print "Aborting"
            sys.exit( 1 )

# Lock the disk device (Eos initblockdev script automatically tries to
# mount the filesystem after we rewrite the partition table, but waits if
# the device is locked)
print "Locking disk device %s" % opts.device
disklock = file( opts.device )
fcntl.flock( disklock, fcntl.LOCK_EX | fcntl.LOCK_NB )

print "Partitioning and formatting disk"
options = { "vfat" : [ opts.device ] + EosDisk.filesystems[ 'vfat' ] +
                     [ swi, opts.boot_config, opts.startup_config, None, None,
                       opts.recovery_swi, opts.add_to_recovery, opts.flash_size ],
            "ext3" : [ opts.device ] + EosDisk.filesystems[ 'ext3' ] +
                     [ swi, opts.boot_config, opts.startup_config, None, None,
                       opts.recovery_swi, opts.add_to_recovery, opts.flash_size ],
            "ext4" : [ opts.device ] + EosDisk.filesystems[ 'ext4' ] +
                     [ swi, opts.boot_config, opts.startup_config, None, None,
                       opts.recovery_swi, opts.add_to_recovery, opts.flash_size ] }
parts = EosDisk.formatHd( *options[opts.fs.lower()] )

print "Writing recovery file and populating filesystem"
EosDisk.writeRecovFile( opts.device, parts, swi, opts.boot_config,
      opts.startup_config, opts.copy_tree_from, None, None, opts.recovery_swi,
      opts.add_to_recovery )

# Give some time and space to clear the console of any kernel log messages which may
# confuse the test script. See BUG61948
os.system( 'sync' )
time.sleep( 10 )
print "\n\n\n"

