python drop privileges

drop-privileges.py

import os, pwd, grp


def dropPrivileges(n_uid=None, n_gid=None, new_mask=077, critical=True, target_dir='.'):
    '''
        drop privileges without forking
        e.g.:
            http://www.charleshooper.net/blog/dropping-privileges-in-python-for-tornado-apps/
            http://antonym.org/2005/12/dropping-privileges-in-python.html
        
        if uid and gid are not provided, we try and get the info from stat'ing the local
        directory. this default approach app won't work for most deployments outside of
        your home dir. keep that in mind.
        
        critical: Specifies the return value if the entry point was not root
        
        alternative conversion approach:
            n_uid = pwd.getpwuid(stat_info.st_uid)[0]
            n_gui = grp.getgrgid(stat_info.st_gid)[0]
            
        provide your own logger from the (status, msg) return tuple
    '''
    try:
        
        if os.getuid() != 0:
            if critical:
                msg = 'entry uid was not root but %s'%os.getuid()
                raise Exception(msg)
            else:
                return True, 'no drop needed as not root'
        
        #get a fix on the new uid and  gid
        stat_info = os.stat(target_dir)
        if n_uid is None:
            running_uid = stat_info.st_uid
            if running_uid == 0:
                msg = 'failure to stat uid from current dir, %s'
                vals = (os.path.realpath(__file__), running_uid)
        else:
            running_uid = pwd.getpwnam(n_uid).pw_uid
        
        if n_gid is None:
            running_gid = stat_info.st_gid
        else:
            running_gid = grp.getgrnam(n_gid).gr_gid
        
        # remove group privileges
        os.setgroups([])
        
        #set new uid, gid
        os.setgid(running_gid)
        os.setuid(running_uid)
        
        
        #set umask
        old_umask = os.umask(new_mask)
        
        #just making sure nothing slipped
        if os.getuid() == 0:
            raise Exception()
        
        msg = 'dropped root, %s, to %s, %s, %s'
        vals = (os.getuid(), running_uid, running_gid, new_mask)
        return True, msg%vals
        
    except Exception, e:
        msg = 'dropPrivileges failure of %s, %s, %s: %s'
        vals = (n_uid, n_gid, new_mask, e)
        return False, msg%vals




def testing():
    who_before = subprocess.check_output('whoami', shell=True).strip('\n')
    
    drop, msg = dropPrivileges()
    if not drop:
        print 'Error: %s'%msg
    
    who_after = subprocess.check_output('whoami', shell=True).strip('\n')
    print 'whoami before: %s and after: %s dropPrivileges'%(who_before, who_after)



if __name__=='__main__':
    import subprocess
    testing()