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

import argparse
import OnieFactoryReset
import subprocess
import os
import sys

def confirmMessage( message ):
   yes = {  'yes', 'y'  }
   # pylint: disable-next=consider-using-f-string
   choice = input( "%s [y/N] " % message ).lower()
   return choice in yes

def executeAction( skipCheck, setBootMode, action ):
   # pylint: disable-next=consider-using-f-string
   if skipCheck or confirmMessage( "Are you sure you want to go back to %s?"
                                   % action ):
      setBootMode()
      # pylint: disable-next=consider-using-f-string
      print( "System will go into %s on the next reboot." % action )
      if skipCheck or confirmMessage( "Would you like to reboot now?" ):
         # pylint: disable-next=consider-using-f-string
         print( "Rebooting into %s..." % action )
         subprocess.call( "/sbin/reboot" )

def main():
   parser = argparse.ArgumentParser()
   group = parser.add_mutually_exclusive_group()
   group.add_argument( '-i', '--install', action="store_true",
                      help="Go into ONIE: install on next reboot." )
   group.add_argument( '-r', '--rescue', action="store_true",
                      help="Go into ONIE: rescue on next reboot." )
   parser.add_argument( '-y', '--yes', action="store_true",
                      help="Skip confirmation and reboot immediately." )

   args = parser.parse_args()
   # Check that user is root before proceeding
   if not os.geteuid() == 0:
      sys.exit( "\nOnly root can run factory-reset!\n" )

   if args.install:
      action, mode = OnieFactoryReset.installOnReboot, "ONIE: Install"
   elif args.rescue:
      action, mode = OnieFactoryReset.rescueOnReboot, "ONIE: Rescue"
   else:
      action, mode = OnieFactoryReset.uninstallOnReboot, "ONIE: Uninstall"

   executeAction ( args.yes, action, mode )

if __name__ == "__main__":
   main()
