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

# pylint: disable=redefined-outer-name
# pylint: disable=consider-using-f-string

import os
import sys
import optparse # pylint: disable=deprecated-module
import ctypes
import PyClient, Tac, AgentDirectory, LauncherLib

def _getAgentCfg( sysname, agentName ):
   pcl = PyClient.PyClient( sysname, 'Sysdb' )
   agentCfgDir = LauncherLib.agentConfigDir( pcl.agentRoot() )
   if agentName in agentCfgDir.agent:
      return agentCfgDir.agent[ agentName ]
   return None

def _corelimitBytesInRunningAgentIs( sysname, agentName, corelimitBytes ):
   if not AgentDirectory.agent( sysname, agentName ):
      return False
   # enable coredump in the running agent 
   pcl = PyClient.PyClient( sysname, agentName )
   pcl.execute( 'import resource; '
                'resource.setrlimit( resource.RLIMIT_CORE, ( %d, %d ));' %
                ( corelimitBytes, corelimitBytes ) )
   return True

def corelimitBytesIs( sysname, agentName, corelimitBytes, persist=False ):
   if not persist:
      if _corelimitBytesInRunningAgentIs( sysname, agentName, corelimitBytes ):
         print( "Set corelimit bytes to %s in the running %s agent" % \
               ( str(corelimitBytes) if corelimitBytes != defaultCorelimit else
                 "unlimited", agentName ) )
      else:
         print( "Agent %s is not running, cannot set corelimit." % agentName )
      return
   if AgentDirectory.agent( sysname, 'Sysdb' ):
      # enable coredump in the running Sysdb 
      agentCfg = _getAgentCfg( sysname, agentName )
      if agentCfg:
         agentCfg.stable = False
         agentCfg.corelimitBytes = corelimitBytes
         agentCfg.stable = True
         print( "Set corelimit bytes to %s for agent %s in Sysdb" % \
               ( str(corelimitBytes) if corelimitBytes != defaultCorelimit else
                 "unlimited", agentName ) )
         if AgentDirectory.agent( sysname, 'Launcher' ):
            print( "Agent will be restarted automatically by Launcher" )
   # enable coredump in Sysdb plugin, if it is present
   pluginFound = False
   for dir in sys.path: # pylint: disable=redefined-builtin
      dir += "/SysdbPlugin"
      if os.path.exists( dir ):
         # find agent's Sysdb plugin file
         grepCmd = """grep -l -E "(agentName *= *['\\"]%s['\\"]|"""\
                   """agentCfg *= *{ *['\\"]name['\\"] *: *['\\"]%s['\\"])" """\
                   """%s/*py""" % \
                   ( agentName, agentName, dir )
         files = Tac.run( [ "bash", "-c", 
                            grepCmd ], stdout=Tac.CAPTURE, asRoot=True,
                            ignoreReturnCode=True )
         for f in files.split( "\n" ):
            if f:
               # edit agent's Sysdb plugin file
               pluginFound = True
               corelimitPresent = Tac.run( [ "grep", "-l", 
                                             """['"]corelimitBytes['"]: """, f ], 
                                             ignoreReturnCode=True,
                                             stdout=Tac.CAPTURE,
                                             asRoot=True )
               if not corelimitPresent:
                  Tac.run( [ "sed", "-i", "-e", 
                             """/['"]exe['"]/i "corelimitBytes": %d,""" % 
                             corelimitBytes, f ],
                             asRoot=True  )
               else:
                  Tac.run( [ "sed", "-i", "-e", 
                             's/"corelimitBytes": [0-9]*,/"corelimitBytes": %d,/' % 
                             corelimitBytes, f ], asRoot=True  )
               corelimitIsSet = Tac.run( [ "grep", "-l",
                                           """['"]corelimitBytes['"]: %d""" % 
                                           corelimitBytes, f ], 
                                           ignoreReturnCode=True,
                                           stdout=Tac.CAPTURE, asRoot=True )
               if corelimitIsSet: 
                  print( "Set corelimit bytes to %s for agent %s " \
                        "in Sysdb plugin %s" % ( 
                        str(corelimitBytes) if corelimitBytes != defaultCorelimit 
                        else "unlimited", agentName, f ) )
   if pluginFound:
      return

   # agent does not have a sysdb plugin, perhaps it has a static config?
   configFound = False
   files = [ "/etc/ProcMgr.d/run/" + agentName, "/etc/ProcMgr.d/inst/" + agentName ]
   for f in files:
      if os.path.exists( f ):
         configFound = True
         corelimitPresent = Tac.run( [ "grep", "-l", "corelimitBytes *=", f ],
                                     ignoreReturnCode=True,
                                     stdout=Tac.CAPTURE, asRoot=True )
         if not corelimitPresent:
            echoCmd = """echo "corelimitBytes = %d" >> %s""" % ( corelimitBytes, f )
            Tac.run( [ "bash", "-c", echoCmd ], stdout=Tac.CAPTURE, asRoot=True )
         else:
            Tac.run( [ "sed", "-i", "-e",
                       's/corelimitBytes *= *[0-9]*/corelimitBytes = %d/' %
                       corelimitBytes, f ], asRoot=True  )
         corelimitIsSet = Tac.run( [ "grep", "-l", "corelimitBytes = %d" % 
                                     corelimitBytes, f ], stdout=Tac.CAPTURE, 
                                     asRoot=True )
         if corelimitIsSet: 
            print( "Set corelimit bytes to %s for agent %s " \
                  "in its config file %s" % ( str(corelimitBytes)
                  if corelimitBytes != defaultCorelimit else "unlimited", \
                                                 agentName, f ) )
         if _corelimitBytesInRunningAgentIs( sysname, agentName, corelimitBytes ):
            print( "Set corelimit bytes to %s in the running %s agent" % \
                  ( str(corelimitBytes) if corelimitBytes != defaultCorelimit 
                    else "unlimited", agentName ) )
   if configFound:
      return 

   print( "Sysdb plugin or agent config for agent %s are not found" % agentName )

sysname = os.environ.get('SYSNAME', 'ar')
# Let us have a near unlimited default value (though the 2.6.32 kernel 
# does not seem to honour this and dumps entire core when limit > 0). 
# Some of the agent core-files are pretty large of 400 MB or more.
# If the user wants to get a non truncated core-file (truncated because it 
# does not fit in /var/core), the "system coredump compressed" CLI option 
# can be enabled and tried if the compressed core-file fits in.
defaultCorelimit = ctypes.c_uint32(~0).value
corelimitBytes = defaultCorelimit

parser = optparse.OptionParser( usage="Usage: %prog [-dp] [-c corelimit] "
                                      "[-s sysname] agent agent ...",
                              add_help_option=False)
parser.add_option( '-d', '--disable', action="store_true", 
                   help='disable coredump' )
parser.add_option( '-p', '--persist', action="store_true", 
                   help='persist after agent or Sysdb restart' )
parser.add_option( '-s', '--sysname', 
                   help='specify sysname',
                   default = sysname )
parser.add_option( '-c', '--corelimit', action="store", type="int", 
                   help="specify corelimit bytes (default is unlimited)", 
                   default=corelimitBytes )
parser.add_option( '--help', action="help",help="show this help message and exit")

options, agents = parser.parse_args()

if not agents:
   print( "No agent specified." )
   sys.exit()

if options.disable:
   options.corelimit = 0

for agentName in agents:
   corelimitBytesIs( options.sysname, agentName, options.corelimit, options.persist )
