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

import subprocess, sys, signal, time, os, syslog
import optparse # pylint: disable=deprecated-module
import ctypes

def now():
   return time.time()

def rounddown( n, div=1 ):
   return float(int( n * div )) / div

def main():
   parser = optparse.OptionParser( """usage: %prog [OPTIONS] COMMANDSTRING
   Ex:
      %prog --daemonize "arhttpd --port 8989" """ )
   parser.allow_interspersed_args = False
   parser.add_option( "--daemonize", action="store_true",
                      help="Fork myself into the background" )
   parser.add_option( "--log", action="store",
                      help="Redirect all output to the given log file" ) 
   parser.add_option( "--logappend", action="store_true",
                      help="Open log file in append mode.  Useful with logrotate." )
   parser.add_option( "--maxcredits", action="store", default='150',
                      help="Max credits for restarting" )
   parser.add_option( "--cost", action="store", default='30',
                      help="Cost for each restart" )
   parser.add_option( "--logpidsuffix", action="store_true",
                      help="Add immortalize pid suffix '-<pid>' to the name of log" +
                           "file." )
   parser.add_option( "--immortalizepidfile", action="store",
                      help="If set, write the pid of immortalize to this file" )
                           
   (opts, args ) = parser.parse_args()

   cmd = args

   if opts.daemonize:
      # double fork
      if os.fork(): 
         os._exit(0) # pylint: disable-msg=W0212
      os.setsid()
      if os.fork(): 
         os._exit(0) # pylint: disable-msg=W0212
      r = open( "/dev/null" ) # pylint: disable=consider-using-with
      w = open( "/dev/null", "w" ) # pylint: disable=consider-using-with
      os.dup2( r.fileno(), 0 )
      os.dup2( w.fileno(), 1 )
      os.dup2( w.fileno(), 2 )

   stdout = None

   def initLogrotateLib( fd ):
      l = ctypes.CDLL( "libAgentLogrotateLib.so" )
      l.reopenFd( fd )

   if opts.log:
      name = opts.log
      if opts.logpidsuffix:
         name = name + '-' + str( os.getpid() )
      if opts.logappend:
         stdout = open( name, "a" )
      else:
         stdout = open( name, "w" ) # pylint: disable=consider-using-with
      # Writing logs to a file in /var/log/agents. Should handle signal from
      # logrotate
      if "/var/log/agents/" in name:
         initLogrotateLib( stdout.fileno() )

   if opts.immortalizepidfile:
      with open( opts.immortalizepidfile, 'w' ) as f:
         f.write( str( os.getpid() ) )
         f.close()

   tracing = "TRACE" in os.environ
   alwaysRestart = "IMMORTALIZE_ALWAYS_RESTART" in os.environ

   def trace( *args ):
      if tracing:
         sys.stdout.write( " ".join( [str(i) for i in args] ) + "\n" )
         sys.stdout.flush()

   # With maxcredits=150 and cost=30, we will restart no more than
   # every 30 seconds.  After 2.5 minutes of succesful runtime, we
   # get 5 more restart attempts.  A constantly crashing program will restart
   # once every 30 seconds.
   maxcredits = float( opts.maxcredits )
   cost = float( opts.cost )

   credit = maxcredits
   first = True
   while True:
      started = now()
      trace( "running now" )

      # pylint: disable-next=consider-using-with
      p = subprocess.Popen( cmd, stdout=stdout, stderr=stdout, close_fds=False )
      if first:
         first = False
      else:
         syslog.syslog( f"immortalize: restarting {sys.argv[1]} {p.pid}" )

      rc = p.wait()
      if rc < 0:
         if -rc in [signal.SIGKILL, signal.SIGTERM] and not alwaysRestart:
            trace( "Program exited with signal", -rc, "exiting." )
            break

      what = (rc < 0) and "signal" or "code" # pylint: disable=consider-using-ternary
      trace( f"Program exited with {what} {rc}, restarting" )

      # Accumulate credit equal to the number of seconds we ran for
      ranfor = now() - started
      # pylint: disable-next=consider-using-f-string
      trace( "Program ran for %s" % ranfor )
      if credit + ranfor > maxcredits:
         credit = maxcredits
      else:
         credit += ranfor

      # Eat up "cost" credit each time we restart
      credit -= cost

      # If credit < 0 then we sleep until credit is positive again.
      # One credit per second of sleep.
      if credit < 0:
         trace( "Restarted too often, sleeping for", rounddown( -credit, 10 ))
         time.sleep( -credit )
         credit = 0

if __name__ == "__main__":
   main()
