#!/usr/bin/env python3
# Copyright (c) 2006, 2007, 2008 Arastra, Inc.  All rights reserved.
# Arastra, Inc. Confidential and Proprietary.
#
# This was originally "a4 sudo", but that didn't belong in Artools
# since it is really something used by Artest.

import os, pwd, sys
import resource

def execvp( args ):
   if "P4USER" not in os.environ and "LOGNAME" in os.environ:
      os.environ["P4USER"] = os.environ["LOGNAME"]
   preserve = [ "LD_LIBRARY_PATH", "PATH" ]
   if "ARSUDO_SAVE_ENV" in os.environ:
      preserve += os.environ[ "ARSUDO_SAVE_ENV" ].split( ":" )
   variables = [ "{}={}".format( var, # pylint: disable=consider-using-f-string
         os.environ.get( var, "" ) ) for var in preserve ]
   cmd = [ "/usr/bin/sudo", "/usr/bin/env" ] + variables + args
   os.execv( cmd[0], cmd )

def sudoHandler( args=sys.argv[1:] ): # pylint: disable=dangerous-default-value
   usage = """
   arsudo -- Run command as root preserving key environment variables

   arsudo [-u USER] CMD
   """
   
   if args and args[0] in ( "-h", "--help" ):
      args = []

   # 2000 was used to distinguish it from 2048 (autobuild.py)
   # This way,if the pid limit is 2000, we know it's set lower than autobuild.py
   minLimit = 2000

   pidlims = resource.getrlimit( resource.RLIMIT_NPROC )
   softLimit = pidlims[0]
   if softLimit != resource.RLIM_INFINITY and softLimit < minLimit:
      resource.setrlimit( resource.RLIMIT_NPROC, ( minLimit, 2 * minLimit ) )

   if args and args[0] == "-u":
      args.pop( 0 )
      if args:
         pw = pwd.getpwnam( args.pop( 0 ) )
         os.setgid( pw.pw_gid )
         os.setgroups( [ pw.pw_gid ] )
         os.setuid( pw.pw_uid )
         os.environ["USER"] = pw.pw_name
         # getpass.getuser() check's LOGNAME. If we don't set this,
         # ArtoolsTestLib can't spin up a MySQL instance
         os.environ["LOGNAME"] = pw.pw_name

   if not args:
      print( usage )
      sys.exit()

   try:
      os.execvp( args[0], args )
   except OSError as e:
      # pylint: disable-next=consider-using-f-string
      sys.stderr.write( "%s: %s\n" %( args[0], e ))
      sys.exit( 1 )

# Rerun the arsudo command as root
if os.getuid() != 0:
   execvp( sys.argv ) 

sudoHandler()
