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

import os
import sys
import AgentDirectory
import ExtensionMgrLib
from ExtensionMgr import logs
import PyClient
import Tracing
import Tac

t0 = Tracing.trace0

def printHelpAndExit():
   print( '''Small script we run daemonizied when finalizing extension installs, or 
from the %pre/%preun section of an rpm script in case no swix manifest file is used
(maybe because customer does not like .swix) to report Agents that need restart.
The print api is for debugging only (not visible in capi/non-interactive users).
The rpm name is passed with EXTENSIONMGR_EXTENSION env var.
arguments:
  restart (<rpm>|all) <agent-1> [ { <agent-n> } ]
  print   <text>
  agentsToRestart <agent-1> [ { <agent-n } ]
''' )
   sys.exit( -1 )

sysname = os.environ.get( "SYSNAME", "ar" )
pc = PyClient.PyClient( sysname, "Sysdb" )
root = pc.agentRoot()
status = root.entity[ ExtensionMgrLib.statusPath() ]

def doRestart( rpm, agents ):
   
   ExtensionMgrLib.log( logs.EXTENSION_RESTARTING_AGENTS, " ".join( agents ) )

   #create a pid dict to kill agents
   pidDict = { a[ "name" ] : a[ "pid" ] for a in AgentDirectory.agentList() }  
   
   for agent in agents:
      try:
         if agent in pidDict: 
            Tac.run( [ 'kill', '-9', str( pidDict[ agent ] ) ], stdout=Tac.CAPTURE )
         else:    
            Tac.run( [ 'killall', '-9', agent ], stdout=Tac.CAPTURE )
      except Tac.SystemCommandError:
         pass

   # For the agents we killed successfully, remove them from the agentsToRestart
   for info in status.info.values():
      if rpm != "" and info.filename != rpm: # pylint: disable=consider-using-in
         continue
      info.agentsToRestart.clear()
  
def getStatusInfo( rpm ):
   for info in status.info.values():
      if info.filename != rpm:
         continue
      return info

def printText( text ):
   ttyname = os.environ.get( "REALTTY" )
   if ttyname and ttyname.startswith( "/dev/" ):
      ttyname = ttyname.split( " " )[ 0 ]
      parts = ttyname.split( "/" )
      ttynum = int( parts[ -1 ] )
      ttypath = "/".join( parts[ 0:-1 ] )
      # pylint: disable-next=consider-using-f-string
      ttyname = "%s/%d" % ( ttypath, ttynum + 1 )
      f = open( ttyname, "a" ) # pylint: disable=consider-using-with
      f.write( text )
      f.write( "\n" )
      f.close()

def setAgentsToRestart( rpm, agents ):
   info = getStatusInfo( rpm )
   if info:
      info.agentsToRestart.clear()
      for i, a in enumerate( agents ):
         info.agentsToRestart[ i ] = a
      info.affectedAgents.clear()
      for i, a in enumerate( agents ):
         info.affectedAgents[ i ] = a

def main( argv ):
   
   argc = len( argv )
   if argc < 3:
      printHelpAndExit()

   topic = argv[ 1 ]
   if topic == "restart":
      if argc < 4:
         printHelpAndExit()
      rpm = argv[ 2 ]
      args = argv[ 3: ]
      doRestart( rpm, args )
      return

   rpm = os.environ.get( "EXTENSIONMGR_EXTENSION" )
   if not rpm:
      print( "env var EXTENSIONMGR_EXTENSION not set to extension filename" )
      sys.exit( -1 )

   args = argv[ 2: ]
   if topic == "print":
      printText( " ".join( args ) )
      return
   if topic == "agentsToRestart":
      setAgentsToRestart( rpm, args )
      return
   printHelpAndExit()

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

