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

import os, sys

import importlib

try:
   # Python 3
   from importlib import util as import_util

   class RootModuleNotFoundError( Exception ):
      pass

   # importlib.import_mdoule  does not throw ImportError on Python 3, but instead
   # ModuleNotFoundError, so different code is required between versions. Using the
   # implementation below, we can more reliably detect import failure on Python 3
   # Adapted from 
   # docs.python.org/3/library/importlib.html#checking-if-a-module-can-be-imported
   def try_import( name ):
      spec = import_util.find_spec( name )
      if spec is None:
         raise RootModuleNotFoundError
      module = import_util.module_from_spec( spec )
      sys.modules[ name ] = module
      spec.loader.exec_module( module )
      return module

except ImportError:
   # Python 2
   RootModuleNotFoundError = ImportError
   # import_module will throw an ImportError if the imported module incorrectly
   # imports something. AFAICT, we can't detect this in python 2.
   try_import = importlib.import_module

class SwiCommand:
   def __init__( self, name, help, module=None ):
      self.name_ = name
      self.help_ = help or "(undocumented)"
      self.module_ = module or ( "Swi." + name )

   def name( self ):
      return self.name_

   def help( self ):
      return self.help_

   def _module( self ):
      if self.module_ in sys.modules:
         # The module is already imported
         return sys.modules[ self.module_ ]
      return try_import( self.module_ )

   def run( self, args ):
      f = getattr( self._module(), self.name_ + "Handler" )
      f( args=args )

   def printHelp( self ):
      try:
         getattr( self._module(), self.name_ + "Help" )()
      except AttributeError:
         print( self.help() )

swicommands = {c.name(): c for c in (
   SwiCommand( "info",
                help="Display version and signature info for the given swi file" ),
   SwiCommand( "extract",
                help="Extract the contents of the given swi file" ),
   SwiCommand( "extend",
                help="Extend the swi file by adding extension files" ),
   SwiCommand( "create",
                help="Create a swi file with the given name" ),
   SwiCommand( "flavor",
                help="Modify base EOS.swi with a particular flavor" ),
   SwiCommand( "zip",
                help="zip together two swi files into one" ),
   SwiCommand( "installrootfs",
                help="Install packages into a rootfs directory" ),
   SwiCommand( "bless",
                help="Mark an image as officially released" ),
   SwiCommand( "checksum",
                help="Generates checksum data for provided image "
                      "file", ),
   SwiCommand( "checksumverify",
                help="Verifies if image checksums match "
                     "with checksums in provided file" ),
   SwiCommand( "chroot",
                help="chroot into an extracted swi rootfs" ),
   SwiCommand( "umount",
               help="unmount file systems within extracted swi rootfs",
               module="Swi.chroot" ),
   SwiCommand( "update",
               help="update a file/directory in the swi rootfs" ),
   SwiCommand( "rpm",
                help="Perform an rpm operation inside a swi" ),
   SwiCommand( "freshen",
                help="Freshen the RPMs inside a swi with those built in "
                     "this workspace" ),
   SwiCommand( "workspace",
               help="Create (or delete) a workspace from a .swi file" ),
   SwiCommand( "bash",
               help="Extract swi, run Bash in it, repackage upon exit" ),
   SwiCommand( "sign", help="Sign the SWI." ),
   SwiCommand( "licenses", help="Output open-source licenses from a SWI." ),
   SwiCommand( "strip",
               help="Create a stripped EOS.swi from full EOS.swi that only "
                    "supports specified hardware" ),
   SwiCommand( "supports",
               help="List the platforms supported by the swi" ),
   SwiCommand( "flatten",
               help="Flatten a multi squash SWI into a single squash SWI" ),
   SwiCommand( "sxv",
               help="Create sxverity images from a SWI or directory" ),
   # plugin from EosSwimUtils
   SwiCommand( "installswimfs",
               help="Install packages into modularized swim fs" ),
   SwiCommand( "swimadapt",
               help="Read local or remove EOS modularized AKA swim image, "
                    "adapt and store in specified flavor." ),
   )}

if __name__ == "__main__":
   if not sys.argv[1:]:
      sys.stderr.write( "Unknown command.  Try 'swi help' for info.\n" )
      sys.exit( 1 )

   # Figure out where swi lives so the right modules are loaded when the user
   # invokes swi with '/usr/bin/swi' or 'python mytree/src/EosUtils/swi.in'
   swi = os.path.abspath( sys.argv[0] )
   swidir = os.path.dirname( swi )
   if swidir.endswith( "/usr/bin" ):
      libdir = os.environ.get( "PYTHONPATH", "" ).split( ":" )[0]
      if libdir:
         sys.path.insert( 0, os.path.join( swidir, "..", "..", libdir ) )
   else:
      sys.path.insert( 0, swidir )
      swisrcdir = swidir.replace( "/bld/", "/src/" )
      sys.path.insert( 0, swisrcdir )

   cmd, args = sys.argv[1], sys.argv[2:]

   if cmd == "--version":
      sys.stdout.write( "swi version 1.1.19\n" )
   elif cmd == "help" and args and args[0] in swicommands:
      try:
         swicommands[ args[0] ].printHelp()
      except RootModuleNotFoundError:
         sys.stderr.write( "use 'swi help' for help\n" )
         sys.exit( 1 )
   elif cmd in swicommands:
      try:
         swicommands[ cmd ].run( args )
      except RootModuleNotFoundError:
         sys.stderr.write( "use 'swi help' for help\n" )
         sys.exit( 1 )
   elif cmd == "help" and ( not args or args[0] == "commands" ):
      for c in sorted( swicommands ):
         try:
            swicommands[c]._module() # Check if the subcommand is importable
            sys.stdout.write( "\t%-11s %s\n" % (c, swicommands[c].help()) )
         except RootModuleNotFoundError:
            # If not, skip printing the help
            pass
      sys.stdout.flush()
   else:
      sys.stderr.write( "use 'swi help' for help\n" )
      sys.exit( 1 )
