#!/usr/bin/python
# 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 = [ "%s=%s" % ( var, 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:] ):
   usage = """
   arsudo -- Run command as root preserving key environment variables

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

   pidlims = resource.getrlimit( resource.RLIMIT_NPROC )
   if pidlims[0] < 2000:  #2000 was used to distinguish it from 2048(autobuild.py)
      resource.setrlimit( resource.RLIMIT_NPROC, ( 2000, 4000 ) )
      # This way,if the pid limit is 2000, we know it's set lower than autobuild.py

   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, e:
      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()
