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

# pylint: disable=consider-using-f-string

import optparse, sys, os, fcntl, time # pylint: disable=deprecated-module
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( f"{opts.device} is mounted ({mounted})\n" )
   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 ( b'\x00', b'\x55', b'\xaa', b'\xff' ):
      print( "Writing 0x%x to the flash" % ord(v) )
      with open( opts.device, "wb" ) as f:
         for i in range( size // 1024 ):
            f.write( v * 1024 )
         f.flush()

      print( "Verifying the contents of the flash" )
      with open( opts.device, "rb" ) as f:
         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 = open( opts.device ) # pylint: disable=consider-using-with
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" )

