boot-eos

#!/usr/bin/env python3

import argparse
import hashlib
import logging
import os
import re
import shutil
import subprocess
import urllib
import zipfile


sonicMount = "/host"

def verifySWI(swiPath):
   zf = zipfile.ZipFile(swiPath, 'r')
   reVer = re.compile(r'.*SWI_VERSION=.*(\d+)\.(\d+)\.(\d+).*')
   for str in zf.read("version").splitlines():
      m = reVer.match(str)
      if m:
         if int(m.group(1)) >= 4 and int(m.group(2)) >= 19 and int(m.group(3)) >= 0:
            return ""
         else:
            return "Unsupported version of EOS"
   return "Cannot find version of EOS in SWI file"

unsupportedPlatforms = ["x86_64-arista_7050_qx32", "x86_64-arista_7050_qx32s"]

def verifyPlatform():
   # Check that this script run on an ext4 flash
   platform = subprocess.check_output(["sonic-cfggen",
      "-v", "DEVICE_METADATA.localhost.platform", "-d"]).rstrip()
   if not platform:
      platform = subprocess.check_output(["sonic-cfggen",
         "-v", "platform", "-m", "/etc/sonic/minigraph.xml"]).rstrip()

   return (platform not in unsupportedPlatforms)

def checkSpace(swiPath):
   # for EOS we need ~(2 * <swi size>). swi already on fs, so we need additional
   # <swi size> bytes.
   mb = (1 << 20)
   swiMB = os.stat(swiPath).st_size / mb / 100 * 100 + 100
   statvfs = os.statvfs(sonicMount)
   freeMB = statvfs.f_frsize * statvfs.f_bavail / mb
   return swiMB - freeMB

def getEOSFiles(ignoreFile):
   res = []
   eosFile = os.path.join( sonicMount, "EOS.swi" )
   if eosFile != ignoreFile and os.path.exists(eosFile):
      res.append(eosFile)
   eosFile = os.path.join( sonicMount, ".boot-image.swi" )
   if eosFile != ignoreFile and os.path.exists(eosFile):
      res.append(eosFile)
   return res

def deleteFiles(files):
   for file in files:
      os.remove(file)

def queryYesNo(question, default="yes"):
   """Ask a yes/no question via raw_input() and return their answer.

   "question" is a string that is presented to the user.
   "default" is the presumed answer if the user just hits <Enter>.
       It must be "yes" (the default), "no" or None (meaning
       an answer is required of the user).

   The "answer" return value is True for "yes" or False for "no".
   """
   valid = {"yes": True, "y": True, "ye": True,
             "no": False, "n": False}
   if default is None:
      prompt = " [y/n] "
   elif default == "yes":
      prompt = " [Y/n] "
   elif default == "no":
      prompt = " [y/N] "
   else:
      raise ValueError("invalid default answer: '%s'" % default)

   while True:
      print(question + prompt)
      choice = raw_input().lower()
      if default is not None and choice == '':
         return valid[default]
      elif choice in valid:
         return valid[choice]
      else:
         print("Please respond with 'yes' or 'no' "
                           "(or 'y' or 'n').\n")

def main():
   parser = argparse.ArgumentParser(
      description='Arista tool to boot EOS',
      formatter_class=argparse.ArgumentDefaultsHelpFormatter
   )

   parser._action_groups.pop()
   required = parser.add_argument_group('required arguments')
   optional = parser.add_argument_group('optional arguments')
   required.add_argument('--swi', "-s", help="EOS.swi file", required=True)
   optional.add_argument("--md5", "-m", help="MD5 of EOS.swi file")
   startupConfig = os.path.join(sonicMount, "startup-config")
   startupHelp = "EOS startup-config url. Or you can just put your " \
                "startup-config file to %s" % startupConfig
   if os.path.exists(startupConfig):
      optional.add_argument("--config", "-c", help=startupHelp)
   else:
      required.add_argument("--config", "-c", help=startupHelp, required=True)
   optional.add_argument("--remove", "-r", help="remove old EOS files if there "
      "is not enough space for installation of new version of EOS",
      action="store_true")
   optional.add_argument("--no-reboot", "-n", help="config only (without reboot)",
                     action="store_true")

   args = parser.parse_args()

   # 1) check that we have root privileges
   if os.geteuid() != 0:
      exit("You need to have root privileges to run this script. Please try "
           "again using 'sudo'. Exiting.")

   # 2) download EOS swi
   if args.swi.startswith(('http://', 'https://', 'ftp://')):
      swi_path = "/tmp/EOS.swi"
      try:
         urllib.urlretrieve(args.swi, swi_path)
      except IOError as e:
         exit("Cannot download SWI '{}'. Error: '{}'. Exiting.".format(
            args.swi, e))
   else:
      swi_path = args.swi
      if not os.path.exists(swi_path):
         exit("SWI path %s does not exists. Exiting." % swi_path)

   # 3) check MD5 of EOS swi
   if args.md5:
      hash_md5 = hashlib.md5()
      with open(swi_path, "rb") as f:
         for chunk in iter(lambda: f.read(10 * (1 << 20)), b""):
            hash_md5.update(chunk)
      if hash_md5.hexdigest() != args.md5:
         exit("Invalid MD5 %s. Exiting." % hash_md5.hexdigest())

   # 4) check that platform supports ext4
   if not verifyPlatform():
      exit("Unsupported platform. Exiting.")

   # 5) check that EOS supports installation on ext4
   err = verifySWI(swi_path)
   if err != "":
      exit("%s. Exiting." % err)

   # 6) check that we have enough free space for EOS installation
   need = checkSpace(swi_path)
   if need > 0:
      filesToDelete = getEOSFiles(swi_path)
      if len(filesToDelete):
         if args.remove or queryYesNo("There is not enough space. Do you want to "
                                      "delete old EOS files: "):
            deleteFiles(filesToDelete)
         need = checkSpace(swi_path)
   if need > 0:
      exit("Do not have enough space. We need additional ~%dMB in %s to install "
           "EOS. Exiting." % (need, sonicMount))

   # 7) download startup-config
   if args.config:
      try:
         urllib.urlretrieve(args.config, startupConfig)
      except Exception as e:
         exit("Cannot download startup-config '{}'. Error: '{}'. Exiting.".format(
            args.config, e))

   # 8) move EOS swi file to /host/EOS.swi
   swiName = os.path.basename(swi_path)
   swiPath = os.path.join(sonicMount, swiName)
   logging.info('moving %s to %s', swi_path, swiPath)
   shutil.move(swi_path, swiPath)

   # 9) create backup copy of boot-config and create new boot-config to start EOS
   bootConfig = os.path.join(sonicMount, "boot-config")
   bootConfigBk = os.path.join(sonicMount, "boot-config.sonic")
   if os.path.exists(bootConfig):
      if os.path.exists(bootConfigBk):
         logging.warning('creating backup of current boot-config: %s, but old '
                         'version of boot-config was found',
                         bootConfigBk + ".1")
         os.rename(bootConfig, bootConfigBk + ".1")
      else:
         logging.info('creating backup of sonic boot-config: %s', bootConfigBk)
         os.rename(bootConfig, bootConfigBk)
   with open(bootConfig, 'w+') as f:
      logging.info('creating EOS boot-config')
      f.write("SWI=flash:/%s" % swiName)

   subprocess.call(['sync'])

   # 10) reboot to EOS
   if not args.no_reboot:
      logging.info('rebooting to EOS')
      subprocess.call(['reboot', '-f'])

if __name__ == '__main__':
   main()