pvl/args.py
changeset 96 d02b0b67c29c
parent 86 5bd98d185c48
child 187 cbf9371a472d
--- a/pvl/args.py	Sat Jan 12 00:54:36 2013 +0200
+++ b/pvl/args.py	Sat Jan 12 16:14:00 2013 +0200
@@ -5,6 +5,10 @@
 import optparse
 import logging
 
+import pwd, grp, os, sys
+
+import logging; log = logging.getLogger('pvl.args')
+
 def parser (parser) :
     """
         Return an optparse.OptionGroup.
@@ -18,6 +22,9 @@
     general.add_option('--log-file',                                                                    help="Log to file")
     general.add_option('--debug-module',    action='append', metavar='MODULE', 
             help="Enable logging for the given logger/module name")
+    
+    general.add_option('--uid',             help="Change uid")
+    general.add_option('--gid',             help="Change gid")
 
     # defaults
     parser.set_defaults(
@@ -28,7 +35,49 @@
  
     return general
 
-def apply (options, logname=None) :
+def apply_setid (options, rootok=None) :
+    """
+        Drop privileges if running as root.
+
+        XXX: this feature isn't very useful (import-time issues etc), but in certain cases (syslog-ng -> python),
+        it's difficult to avoid this without some extra wrapper tool..?
+    """
+
+    # --uid -> pw
+    if not options.uid :
+        pw = None
+    elif options.uid.isdigit() :
+        pw = pwd.getpwuid(int(options.uid))
+    else :
+        pw = pwd.getpwnam(options.uid)
+
+    # --gid -> gr
+    if not options.gid and not pw :
+        gr = None
+    elif not options.gid :
+        gr = grp.getgrgid(pw.pw_gid)
+    elif options.gid.isdigit() :
+        gr = grp.getgrgid(str(options.gid))
+    else :
+        gr = grp.getgrnam(options.gid)
+    
+    if gr :
+        # XXX: secondary groups? seem to get cleared
+        log.info("setgid: %s: %s", gr.gr_name, gr.gr_gid)
+        os.setgid(gr.gr_gid)
+
+    if pw :
+        log.info("setuid: %s: %s", pw.pw_name, pw.pw_uid)
+        os.setuid(pw.pw_uid)
+    
+    elif os.getuid() == 0 :
+        if rootok :
+            log.info("running as root")
+        else :
+            log.error("refusing to run as root, use --uid 0 to override")
+            sys.exit(2)
+
+def apply (options, logname=None, rootok=True) :
     """
         Apply the optparse options.
     """
@@ -45,8 +94,12 @@
         level       = options.loglevel,
         filename    = options.log_file,
     )
+    
+    if options.uid or options.gid or not rootok :
+        # set uid/gid
+        apply_setid(options, rootok=rootok)
 
     # enable debugging for specific targets
     for logger in options.debug_module :
         logging.getLogger(logger).setLevel(logging.DEBUG)
-
+