#! /usr/bin/env python
##**************************************************************
##
## Copyright (C) 1990-2010, Condor Team, Computer Sciences Department,
## University of Wisconsin-Madison, WI.
## 
## Licensed under the Apache License, Version 2.0 (the "License"); you
## may not use this file except in compliance with the License.  You may
## obtain a copy of the License at
## 
##    http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
##
##**************************************************************

import re
import os
import sys
import shutil
import signal
import pwd
from subprocess import *
import time
from optparse import *

# Logger class
class Logger( object ) :
    def __init__( self ) :
        self.__path = None
        self.__file = None
        self.__verbose = 0
        self.__lines = 0
        self.Setup( "/tmp/MasterWrapper.log" )

    def Setup( self, path ) :
        if path == self.__path  :
            return
        old = self.__path
        self.__path = path
        if self.__file is not None :
            if self.__lines :
                print >>self.__file, "Now logging to '"+path+"'"
            self.__file.close( )
            self.__file = None
        try :
            self.__file = open( self.__path, "a" )
            if old is not None  and  self.__lines :
                print >>self.__file, "Previous log was '"+old+"'"
            self.__lines = 0
        except Exception, e :
            print >>sys.stderr, "Can't open log file '"+self.__path+"'", e
            sys.exit( 1 )

    def SetVerbose( self, v ) :
        self.__verbose = v
        if logger is not None :
            logger.Log( 2, "Verbose level now %d" % (v) )

    def ShouldLog( self, level ) :
        return self.__verbose >= level

    def Log( self, level, text, exitval=None ) :
        if level > self.__verbose :
            return
        assert self.__file is not None
        print >>self.__file, text
        self.__lines += 1
        if exitval is not None :
            sys.exit( exitval )


# Condor configuration parameter class
class CondorConfig( object ) :
    def __init__( self, param, append = None, level = 1 ) :
        assert param is not None
        self.__param  = param
        self.__level  = level
        self.__append = append

    def Level( self ) :
        return self.__level

    def ReadConfig( self, ccv, varname ) :
        cmd = [ ccv, self.__param ]
        logger.Log( 3, varname+": "+str(cmd) )
        try :
            p = Popen(cmd, stdout=PIPE, stderr=PIPE)
            (output, err ) = p.communicate()
            if len(err)  and not err.startswith( "Not defined:" ):
                raise IOError(err)
        except Exception, e :
            logger.Log( -1, "Error running '%s': %s" % (cmd, e), 1 )
        if output is not None  and  output.strip() != ""  :
            output = output.strip()
            if self.__append is not None :
                output += self.__append
            logger.Log( 3, " => '"+output+"'" )
            return output
        else :
            logger.Log( 3, " => None" )
            return None


# Manage a single setting
class Setting( object ) :
    def __init__( self, var, description, default=None,
                  export=None, required=False,
                  sysconfig=None, condor_config=None ) :
        self._var = var
        self._description = description
        self._default = default
        self._export = export
        self._value = None
        self._required = required
        self._level = -1
        self.Set( default, "default", 0 )
        if sysconfig == None :
            sysconfig = var
        self._sysconfig = sysconfig
        self._condor_config = condor_config

    def IsEmpty( self ) :
        return self._value is None

    def Get( self ) :
        if self._value is None :
            raise AttributeError
        return str(self._value)

    def Var( self ) :
        return self._var

    def Description( self ) :
        return self._description

    def Export( self ) :
        if self._export :
            assert self._value is not None
            os.environ[self._export] = str(self._value)
            logger.Log( 1,
                        "Exported %s: '%s'='%s'" % \
                        (self._var, self._export, self._value) )

    def Set( self, value, via, level = 9 ) :
        logger.Log( 2,
                    "Set: '%s'='%s' via %s level %d" % \
                    (self._var, str(value), via, level)  )
        if value is None :
            return
        for c in ( "\"", "'" ) :
            if value.startswith( c ) and value.startswith( c ) :
                value = value[1:-1]
                break
        if value != "" and level > self._level :
            self._value = value
            self._level = level

    def IsSysconfig( self, name ) :
        return self._sysconfig == name

    def ReadCondorConfig( self, ccv ) :
        if self._condor_config is None :
            logger.Log( 3, "Skipping %s for %s (no param)"%\
                        (ccv, self._var) )
            return
        if self._level > self._condor_config.Level() :
            logger.Log( 3,
                        "Skipping %s for %s (%d > %d)"%\
                        (ccv, self._var,
                         self._level, self._condor_config.Level())
                        )
            return
        s = self._condor_config.ReadConfig( ccv, self._var )
        if s is not None :
            self.Set( s, "condor_config_val", self._condor_config.Level() )

    def Verify( self ) :
        if self._required and self._value is None :
            print >>sys.stderr, "Required '"+self._var+"' not found"
            sys.exit( 1 )


# Manage all of our settings
class Settings( object ) :
    def __init__( self ) :
        tmp = [
            Setting( "SYSCONFIG", "Sysconfig file to use",
                     default="/etc/sysconfig/condor" ),
            Setting( "LOG", "Condor log directory",
                     sysconfig="CONDOR_LOG", default="/tmp",
                     condor_config=CondorConfig("LOG", level=3) ),
            Setting( "START_USER_NAME", "Name of user to start Condor as",
                     sysconfig="CONDOR_START_USER_NAME",
                     condor_config=CondorConfig("MASTER_USER_NAME") ),
            Setting( "CONDOR_CONFIG_VAL",
                     "Path to the condor_config_val program",
                     required=True ),
            Setting( "CONDOR_CONFIG",
                     "Path to the main Condor configuration file",
                     export="CONDOR_CONFIG", required=True ),
            Setting( "LOCAL_MASTER",
                     "Path to the condor_master binary on local disk",
                     sysconfig="CONDOR_LOCAL_MASTER",
                     condor_config=CondorConfig("LOCAL_MASTER"),
                     required=True ),
            Setting( "SHARED_MASTER",
                     "Path to the condor_master binary on shared disk",
                     sysconfig="CONDOR_SHARED_MASTER",
                     condor_config=CondorConfig("SBIN", "/condor_master"),
                     required=True ),
            Setting( "MASTER_PIDFILE",
                     "Condor Master's PID file",
                     sysconfig="PIDFILE",
                     default="/var/run/condor/condor.pid",
                     condor_config=CondorConfig("RUN", append="/condor.pid") ),
            Setting( "REAL_MASTER_ARGS",
                     "Arguments to pass to the 'real' condor_master binary",
                     sysconfig="CONDOR_REAL_MASTER_ARGS",
                     condor_config=CondorConfig("REAL_MASTER_ARGS") ),
            ]
        self._vars = { }
        for setting in tmp :
            self.AddVar( setting )

    def Get( self, var ) :
        if var not in self._vars :
            raise KeyError
        return self._vars[var].Get( )

    def IsEmpty( self, var ) :
        if var not in self._vars :
            return True
        return self._vars[var].IsEmpty( )

    def AddVar( self, setting ) :
        self._vars[setting.Var()] = setting
        
    def SetVar( self, var, value, via, level ) :
        if var not in self._vars :
            raise KeyError( var )
        self._vars[var].Set( value, via, level )
        
    def SetSysconfigVar( self, var, value, via, level ) :
        for setting in self._vars.values() :
            if setting.IsSysconfig( var ) :
                setting.Set( value, via, level )
                return
        raise KeyError( var )

    def GetVarList( self ) :
        l = { }
        for var in self._vars.keys() :
            l[var] = self._vars[var].Description()
        return l

    def PrintVars( self, level = 2 ) :
        if not logger.ShouldLog( level ) :
            return
        logger.Log( level, "Variables:" )
        for var in sorted(self._vars.keys()) :
            try: v = self._vars[var].Get()
            except: v = None
            logger.Log( level, "  "+var+"="+str(v) )

    def ReadCondorConfig( self ) :
        os.environ["CONDOR_CONFIG"] = self._vars["CONDOR_CONFIG"].Get()
        ccv = self._vars["CONDOR_CONFIG_VAL"]
        for setting in self._vars.values() :
            setting.ReadCondorConfig( ccv.Get() )

    def Verify( self ) :
        for setting in self._vars.values() :
            setting.Verify( )

    def Export( self ) :
        for setting in self._vars.values( ) :
            setting.Export( )


# "main" class
class Main( object ) :
    _version = "2.0"
    def __init__( self ) :
        self._settings = Settings( )

        name = "MasterScriptVersion"
        setting = Setting( "VERSION", "Version of this program",
                           export="_CONDOR_"+name.upper() )
        setting.Set( self._version, "Version", 10 )
        self._settings.AddVar( setting )

        setting = Setting( "EXPRS", "Condor master expression",
                           export="_CONDOR_MASTER_EXPRS" )
        self._settings.AddVar( setting )
        try:
            exprs  = os.environ.get( "_CONDOR_MASTER_EXPRS", None )
            exprs += ", "+name
        except:
            exprs = name
            
        setting.Set( exprs, "Version", 10 )

        self._parser = OptionParser(
            usage="usage: %prog [options] [var=value ...]",
            version="%prog "+self._version )

    def Setup( self ) :
        self._parser.set_defaults( sysconfig = None )
        self._parser.add_option( "--sysconfig",
                                 action="store", dest="sysconfig",
                                 help="Specify alternate sysconfig file" )

        self._parser.set_defaults( log_stderr = False )
        self._parser.add_option( "--log-stderr",
                                 action="store_true", dest="log_stderr",
                                 help="Force logging to stderr" )

        self._parser.set_defaults( list_vars = False )
        self._parser.add_option( "--list-vars",
                                 action="store_true", dest="list_vars",
                                 help="List variables" )

        self._parser.set_defaults( execute = True )
        self._parser.add_option( "--execute",
                                 action="store_true", dest="execute",
                                 help="Enable execution <default>" )
        self._parser.add_option( "-n", "--no-execute",
                                 action="store_false", dest="execute",
                                 help="Disable execution (for test/debug)" )
        self._parser.set_defaults( verbose = 0 )
        self._parser.add_option( "-v", "--verbose",
                                 action="count", dest="verbose",
                                 help="Increment verbosity level" )
        self._parser.set_defaults( quiet=False )
        self._parser.add_option( "-q", "--quiet",
                                 action="store_true", dest="quiet",
                                 help="be vewwy quiet (I'm hunting wabbits)" )

    def Parse( self ) :
        (self._options, self._args) = self._parser.parse_args()
        logger.SetVerbose( self._options.verbose )
        if self._options.list_vars :
            l = self._settings.GetVarList()
            for var in sorted(l.keys( )) :
                description = l[var]
                print "%-20s %s" % ( var+":", description )
            sys.exit( 0 )
        if self._options.sysconfig is not None :
            self._settings.SetVar( "SYSCONFIG", self._options.sysconfig,
                                   "command line", 2 )

    def ProcessArgs( self ) :
        for arg in self._args :
            a = arg.split( "=" )
            if len(a) != 2 :
                raise OptionValueError( "Invalid argument '"+str(arg)+"'" )
            self._settings.SetVar( a[0], a[1], "command line", 3 )

    def ReadSysconfig( self ) :
        try :
            sysconfig = open( self._settings.Get("SYSCONFIG") )
        except Exception, e :
            print >>sys.stderr, "Can't open sysconfig file", e
            sys.exit( 1 )
        for line in sysconfig :
            line = line.strip( )
            if line.startswith( "#" ) :
                continue
            if len(line) == 0 :
                continue
            a = line.split( "=", 1 )
            if len(a) != 2 :
                continue
            try :
                self._settings.SetSysconfigVar( a[0], a[1],
                                                self._options.sysconfig, 2 )
            except KeyError :
                logger.Log( 3,
                            "Ignoring %s from sysconfig: param not found" % \
                            ( a[0] ) )
                pass
        sysconfig.close( )

    def LogInit( self ) :
        tmp = self._settings.Get( "LOG" ) + "/MasterWrapper.log"
        if self._options.log_stderr :
            tmp = "/dev/stderr"
        assert tmp is not None
        logger.Setup( tmp )

    def SwitchToUser( self ) :
        try:
            user = self._settings.Get( "START_USER_NAME" )
        except:
            return
        try :
            pwent = pwd.getpwnam( user )
        except Exception, e :
            logger.Log( -1, "START_USER "+user+":"+str(e), 1 )
        uid = pwent["pw_uid"]
        logger.Log( 1, "Attempting to switch to user %s (%d)" % (user, uid) )
        if os.getuid != 0  and  os.getuid != uid :
            logger.Log( -1, "START_USER "+user+": not root", 1 )
        if not self._options.excute:
            return
        try:
            os.seteuid( uid )
            os.setuid( uid )
        except Exception, e :
            logger.Log( -1, "START_USER: %s (%d): %s" % (user, uid, str(e)), 1 )

    def CopyMasterBinary( self ) :
        try:
            master = self._settings.Get( "LOCAL_MASTER" )
        except:
            logger.Log( -1, "Local master not specified!", 1 )
        new    = master + ".new"
        old    = master + ".old"

        try:
            shared = self._settings.Get( "SHARED_MASTER" )
        except:
            logger.Log( -1, "Shared master not specified!", 1 )

        logger.Log( 2, "LOCAL: '"+master+"'" )
        logger.Log( 2, "SHARED: '"+shared+"'" )
        try:
            shared_stat = os.stat( shared )
        except Exception, e:
            logger.Log( -1, "Can't stat shared master '%s': %s"%(shared, e), 1 )

        try:
            local_stat = os.stat( master )
        except Exception, e:
            # If this is the first time, stat will return error.
            local_stat = None
            logger.Log( -1, "Can't stat local master '%s': %s"%(master, e), None )

        if local_stat is not None and \
           local_stat.st_mtime == shared_stat.st_mtime and \
           local_stat.st_size  == shared_stat.st_size :
            logger.Log( 2, "Shared and local masters match" )
            return

        if not self._options.execute:
            return

        local  = os.path.dirname( master )
        if not os.path.exists( local ) :
            try:
                os.makedirs( local, 0755 )
            except Exception, e :
                logger.Log( -1, "Can't create local dir '%s': %s"%(local,e), 1 )
        elif not os.path.isdir( local ) :
            logger.Log( -1, "Local dir '%s' isn't a directory!: %s"%(local), 1 )
        

        op = None
        try:
            op = "copy %s to %s" % ( shared, new )
            shutil.copy2( shared, new )
            if os.path.exists( master ) :
                op = "rename %s to %s" % ( master, old )
                os.rename( master, old )
            op = "rename %s to %s" % ( new, master )
            os.rename( new, master )
        except Exception, e :
            logger.Log( -1, "Failed to %s: %s" % (op, str(e)), 1  )

    def ExecMasterBinary( self ) :
        master = self._settings.Get( "LOCAL_MASTER" )
        cmd = [ master ]
        if self._settings.IsEmpty( "REAL_MASTER_ARGS" ) :
            pidfile = self._settings.Get( "MASTER_PIDFILE" )
            cmd += [ "-pidfile", pidfile ]
        else :
            args = self._settings.Get( "REAL_MASTER_ARGS" )
            cmd += args.split( " " )
        print cmd
        logger.Log( 1, "Executing: "+str(cmd) )
        if not self._options.execute:
            return
        try :
            os.execv( cmd[0], cmd )
        except Exception, e :
            logger.Log( -1, "Failed to exec %s: %s" % (str(cmd), e, 1) )


    def Main( self ) :
        self.Setup( )
        self.Parse( )
        self.ProcessArgs( )
        self.LogInit( )
        self.ReadSysconfig( )
        self.LogInit( )
        self._settings.ReadCondorConfig( )
        self._settings.Export( )
        if self._options.verbose >= 1 :
            self._settings.PrintVars( )
        self.SwitchToUser( )
        self.CopyMasterBinary( )
        self.ExecMasterBinary( )

logger = Logger( )
main = Main( )
main.Main( )

### Local Variables: ***
### py-indent-offset:4 ***
### python-indent:4 ***
### python-continuation-offset:4 ***
### tab-width:4  ***
### End: ***
