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

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

# Utility to insert instructions into ProcMgr's database of cmdline modifications to
# be done to an agent prior to starting it. Normally this database is empty, as
# such modifications are supposed to be temporary for the duration of debugging
# or until the switch is restarted
# The type of supported modifications:
# - insert valgrind
# - enable tcmalloc's heapcheck
# - set environment variables
# - insert a wrapper (similar to valgrind, except wrapper has its own argvs and
#   env vars can be set.
# - whatever can be done with a regular expression.
# So as arguments, we have
# - add/del/show (wheather to add or remove the modification, or show them)
# - the agent name that needs to have its command line modified
# - the type of modification (valgrind/heapcheck/wrapper/environ/re)
# Depending on the modification type, there may be exta arguments: for heapcheck the
# level of checking, and for REs the pattern and the substitution.

# All of those modification types can be achieved with the RE option, but that's
# tedious, thus the added sugar. But to illustrate the RE options, here are some
# examples:
#   ProcMgrModify add <agent-name> re \
#      '(.*)' \
#      'valgrind \1'
#   ProcMgrModify add <agent-name> re \
#      '(.*)' \
#      '\1 __env__ HEAPCHECK=normal LD_PRELOAD=/usr/lib/libtcmalloc.so'

import os
import sys
import time
import pprint
import subprocess

def printHelp():
   print( """
This is used to tell ProcMgr to modify an agent's command line next time it restarts
There is a flexible but complex regular expression option, but some well-known use
cases have their own options to make the internals more transparent. Note that
only one option is possible per agent at a time.

  %s add <agent-name> valgrind
  %s add <agent-name> heapcheck [<level:minimal|normal*|strict|draconian>]
  %s add <agent-name> environ "VAR1=<val1> VAR2=<val2> ..."
  %s add <agent-name> wrap <binary> ["<arg1> <arg2> ..."] ["VAR1=<val1> ..."]
  %s add <agent-name> re <pattern> <substitution>
  %s del <agent-name>
  %s show
""" % ( sys.argv[0], sys.argv[0], sys.argv[0], sys.argv[0],
        sys.argv[0], sys.argv[0], sys.argv[0] ) )
   sys.exit(-1)

# py formated file, read by ProcMgr on service reload.
databaseFile = "/var/run/ProcMgrModifyCmdline"

def getConfig():
   if os.path.isfile( databaseFile ):
      with open( databaseFile ) as f:
         c = eval( f.read() )  # pylint: disable-msg=W0123
      return c
   return {}

def addConfig( agent, config ):
   c = getConfig()
   c[ agent ] = config
   with open( databaseFile, "w+" ) as f:
      f.write( "%s\n" % pprint.pformat( c ) )

def delConfig( agent ):
   c = getConfig()
   if c.get( agent ):
      del c[ agent ]
      with open( databaseFile, "w+" ) as f:
         f.write( "%s\n" % pprint.pformat( c ) )
      if not c:
         os.unlink( databaseFile )
   else:
      print( "Error: agent '%s' not found" % agent )

def quoteIfNeeded( text ):
   if not text:
      return '""'
   if " " in text:
      return '"%s"' % text
   return text

def showConfig():
   c = getConfig()
   for proc, instrs in sorted( c.items() ):
      print( "%-20s" % proc, end=" " )
      print( "%-10s" % instrs[ 0 ], end=" " )
      print( " ".join( [ quoteIfNeeded( instr ) for instr in instrs[ 1: ] ] ) )

def doValgrind( agent ):
   addConfig( agent, ( "valgrind", ) )

def doHeapCheck( agent, level ):
   addConfig( agent, ( "heapcheck", level ) )

def doEnviron( agent, envs ):
   addConfig( agent, ( "environ", envs ) )

def doRe( agent, pattern, subst ):
   addConfig( agent, ( "re", pattern, subst ) )

def doWrap( agent, args ):
   wrapper = args[ 0 ]
   if len( args ) == 3:
      wargs = args[ 1 ]
      envs = args[ 2 ]
   elif len( args ) == 2:
      wargs = args[ 1 ]
      envs = ""
   else:
      wargs = ""
      envs = ""
   addConfig( agent, ( "wrap", wrapper, wargs, envs ) )

def maybePrintDisregardedArgs( args, maxArgs ):
   if len( args ) > maxArgs:
      print( "Note: disregarding extra arguments '%s'" % " ".join( 
         args[ maxArgs: ] ) )
      for _ in range( maxArgs, len( args ) ):
         args.pop()

def addModif( args ):
   if len( args ) < 2:
      print( "Error: missing modif type "
             "(valgrind|heapcheck|environ|wrap|re)" )
      printHelp()
   agent = args[ 0 ]  # mandatory arg
   tipe = args[ 1 ]   # mandatory arg
   xargs = args[ 2: ] # extra args
   mi = len( xargs ) -1 # max index

   if tipe == "valgrind":
      maybePrintDisregardedArgs( xargs, maxArgs=0 )
      doValgrind( agent )
      return

   if tipe == "heapcheck":
      level = "normal"
      if mi >= 0:
         if xargs[ 0 ] not in [ "minimal", "normal", "strict", "draconian" ]:
            print( "Error: bad heapcheck level '%s'" % xargs[ 0 ] )
            printHelp()
         level = xargs[ 0 ]
      maybePrintDisregardedArgs( xargs, maxArgs=1 )
      doHeapCheck( agent, level )
      return

   if tipe == "environ":
      maybePrintDisregardedArgs( xargs, maxArgs=1 )
      if mi < 0:
         print( "Error: no env vars specified" )
         printHelp()
      doEnviron( agent, xargs[ 0 ] )
      return

   if tipe == "wrap":
      if mi < 0:
         print( "Error: no wrapper binary provided" )
         printHelp()
      maybePrintDisregardedArgs( xargs, maxArgs=3 )
      doWrap( agent, xargs )
      return

   if tipe == "re":
      if mi < 1:
         print( "Error: need a pattern and a substitution" )
         printHelp()
      maybePrintDisregardedArgs( xargs, maxArgs=2 )
      doRe( agent, xargs[ 0 ], xargs[ 1 ] )
      return

   print( "Error: unexpected arguments '%s'" % tipe )

def maybeRerunAsRoot( args ):
   ''' Runs the same command as sudo if not su. '''
   if os.geteuid() != 0:
      os.execvp( 'sudo', [ 'sudo' ] + args )

def main( args ):
   if len( args ) == 2 and args[ 1 ] == "show":
      showConfig()
      sys.exit(0)
   maybeRerunAsRoot( args )
   if len( args ) < 3: # <type:add|del> <agent> are mandatory
      printHelp()

   if args[ 1 ] == "add":
      addModif( args[2:] )
   elif args[ 1 ] == "del":
      maybePrintDisregardedArgs( args, 3 )
      delConfig( args[ 2 ] )
   else:
      print( "Error: '%s' not expected argument" % args[ 1 ] )
      printHelp()

   # If ProcMgr is running, tell it to check its config files
   try:
      # Since we've already written out the changes, any reload
      # time after now will ensure ProcMgr has the changes loaded.
      startTime = time.time()
      subprocess.run( "ProcMgr reload".split(),
                      stderr=subprocess.DEVNULL,
                      check=True )
      while time.time() - startTime < 10 and \
            os.path.getmtime( "/var/run/ProcMgrLastWarmstart" ) < startTime:
         time.sleep( .1 )
      print( "ProcMgr has loaded changes." )
   except subprocess.CalledProcessError:
      print( "Changes queued for ProcMgr start" )

   sys.exit( 0 )

if __name__ == "__main__":
   main ( sys.argv )
