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

import GenprefdlLib, GenfdlLib
from EosInit import getPrefdl
import SwagBoot

import argparse
import os
import re
import subprocess
import sys
import sysconfig

pkgExcludeDir = '/etc/EosPkgExclude.d'
excludeOverrideFile = '/mnt/flash/EosPkgRm.exclude'
includeOverrideFile = '/mnt/flash/EosPkgRm.include'

def useSwagPackages():
   '''
   This script runs during the EosPkgRm stage during bootup, which is always run
   before EosStage1 as per the dependency tree - however, EosStage1 is where we write
   the /etc/swag* files, so the SWAG API will not work prior to that point. To this
   end, we need to manually check the feature toggle and kernel command-line to
   detect whether to include SWAG packages.
   '''
   return SwagBoot.isSwagIncarnation3()

def getUnusedPackages( prefdl, printFilename=False ):
   if not prefdl:
      return set()

   try:
      with open( excludeOverrideFile ) as f:
         print( "Overriding excluded RPMs with", excludeOverrideFile )
         return { line for line in ( l.strip() for l in f ) if line }
   except OSError:
      pass

   prefdlFields = GenprefdlLib.decodePrefdl( prefdl )
   registry = GenfdlLib.FdlRegistryWithKeys()
   # For SWAG, alongside the separate SWAG FDL for each SWAG-supported platform, we
   # also maintain a separate RPM exclude file. Select that by appending "Swag" to
   # the current platform's SID.
   # In the future, we may decide to handle SWAG RPM requirements in a different
   # manner. This is tracked by BUG963906.
   sidSuffix = "Swag" if useSwagPackages() else ""
   fdls = registry.lookupFdl( sid=prefdlFields.sid + sidSuffix,
                              baseAsy=prefdlFields.baseAsy,
                              hwApi=prefdlFields.hwApi,
                              cpuSid=prefdlFields.cpuSid,
                              cpuHwApi=prefdlFields.cpuHwApi,
                              configDict={} )

   if not fdls:
      return set()

   if isinstance( fdls, list ):
      fdls = ','.join( fdls )

   if not os.path.isfile( os.path.join( pkgExcludeDir, fdls + '.exclude' ) ):
      for k in registry.fdlKeys():
         if k.get( 'modular' ) and fdls in registry.lookupFdlWithKey( **k ):
            fdls = k[ 'group' ]
            break
      else:
         return set()

   try:
      filename = os.path.join( pkgExcludeDir, fdls + '.exclude' )
      with open( filename ) as f:
         if printFilename:
            print( filename )
         return { line for line in ( l.strip() for l in f ) if line }
   except OSError:
      return set()

def getExcludedPackages( prefdl ):
   # Inspects the EosPkgExclude.d directory for packages
   # that different platforms want to exclude at runtime
   # from the running EOS image and removes the rpms

   cmdlineFilename = '/proc/cmdline'
   if prefdl:
      try:
         # pylint: disable-next=unused-variable
         sid = GenprefdlLib.decodePrefdl( prefdl ).sid
      except TypeError:
         # vEOS has an incomplete prefdl, which is okay
         pass

   with open( cmdlineFilename ) as cmdLineFile:
      m = re.search( r"platform=(\w+)", cmdLineFile.read() )
   pkgFilename = m.group( 1 ) if m else 'default'

   pkgFilePath = os.path.join( pkgExcludeDir, pkgFilename )
   try:
      with open( pkgFilePath ) as pkgFile:
         return { line for line in ( l.strip() for l in pkgFile ) if line }
   except OSError:
      defaultFilePath = os.path.join( pkgExcludeDir, 'default' )
      try:
         with open( defaultFilePath ) as pkgFile:
            return { line for line in ( l.strip() for l in pkgFile ) if line }
      except OSError:
         sys.exit( 0 )

def getIncludedPackages():
   try:
      with open( includeOverrideFile ) as f:
         print( "Retaining RPMs listed in", includeOverrideFile )
         return set( f.read().splitlines() )
   except OSError:
      return set()

def main():
   parser = argparse.ArgumentParser(
      description='Uninstall RPMs that are not used on this platform' )
   parser.add_argument(
      '-l', '--list',
      action='store_true',
      help='List the RPMs to be removed without actually removing them' )
   parser.add_argument(
      '-f', '--file',
      action='store_true',
      help='Show the exclude file matching this platform without removing any RPMs' )
   args = parser.parse_args()

   prefdl = getPrefdl()
   pkgsToRemove = getUnusedPackages( prefdl, printFilename=args.file )

   pkgsToRemove.update( getExcludedPackages( prefdl ) )

   pkgsToRemove.difference_update( getIncludedPackages() )

   if args.list:
      print( '\n'.join( pkgsToRemove ) )

   if args.file or args.list:
      sys.exit( 0 )

   if not pkgsToRemove:
      sys.exit( 0 )

   cmd = ( "LD_PRELOAD=" + sysconfig.get_config_var( "LIBDIR" ) +
           "/libremovexattr_stub.so "
           # Query all RPMs
           "rpm -qa --qf='"
           # if postun
           "%|POSTUN?"
           # print RPM\t<NAME>
           "{RPM\t%{NAME}\n}"
           # else
           ":"
           # print <NAME>\t<FILENAME> for each file in the RPM
           "{[%{=NAME}\t%{FILENAMES}\n]}"
           # endif
           "|' | "
           # Only take unique RPMs or filenames. This is important so
           # that files provided by multiple RPMs are not removed.
           "sort -k2 | uniq -f1 -u | "
           # Filter the RPMs we want to remove
           "grep -E $'(^|\t)(" + "|".join( pkgsToRemove ) + ")($|\t)' | "
           "tee "
           # Uninstall RPMs with postun scripts
           ">(grep $'^RPM\t' | sed $'s/^RPM\t//' | xargs -r rpm -e --nodeps) "
           # Delete files from RPMs without postun scripts
           ">(grep -v $'^RPM\t' | cut -f2 | xargs -r rm -rf) "
           # Discard stdout
           "> /dev/null"
          )
   # pylint: disable-next=consider-using-with
   proc = subprocess.Popen( cmd, shell=True, universal_newlines=True,
                            stderr=subprocess.PIPE, executable="/bin/bash" )
   ( _, stderr ) = proc.communicate()
   # /usr/share/{man,doc} removed in EosUtils/Swi/installrootfs.py,
   # hide corresponding warnings when rpm -e is run here
   for line in stderr.splitlines( True ):
      if 'remove failed: No such file' not in line:
         sys.stderr.write( line )
   if proc.returncode != 0:
      print( "Removal of excluded packages failed" )

if __name__ == "__main__":
   main()
