From ae99698b4e2df939b95dbb58b5e748ce85cfa1f0 Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Mon, 2 Feb 2015 15:38:55 +0800
Subject: [PATCH] implement --user

fix @278
---
 shadowsocks/daemon.py | 32 ++++++++++++++++++++++++++++++++
 shadowsocks/local.py  |  1 +
 shadowsocks/server.py |  2 ++
 shadowsocks/utils.py  | 22 +++++++++++++---------
 4 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/shadowsocks/daemon.py b/shadowsocks/daemon.py
index d206ccf..3d2b448 100644
--- a/shadowsocks/daemon.py
+++ b/shadowsocks/daemon.py
@@ -183,3 +183,35 @@ def daemon_stop(pid_file):
         sys.exit(1)
     print('stopped')
     os.unlink(pid_file)
+
+
+def set_user(username):
+    if username is None:
+        return
+
+    import pwd
+    import grp
+
+    try:
+        pwrec = pwd.getpwnam(username)
+    except KeyError:
+        logging.error('user not found: %s' % username)
+        raise
+    user = pwrec[0]
+    uid = pwrec[2]
+    gid = pwrec[3]
+
+    cur_uid = os.getuid()
+    if uid == cur_uid:
+        return
+    if cur_uid != 0:
+        logging.error('can not set user as nonroot user')
+        # will raise later
+
+    # inspired by supervisor
+    if hasattr(os, 'setgroups'):
+        groups = [grprec[2] for grprec in grp.getgrall() if user in grprec[3]]
+        groups.insert(0, gid)
+        os.setgroups(groups)
+    os.setgid(gid)
+    os.setuid(uid)
diff --git a/shadowsocks/local.py b/shadowsocks/local.py
index 994b6d8..87c5215 100755
--- a/shadowsocks/local.py
+++ b/shadowsocks/local.py
@@ -73,6 +73,7 @@ def main():
             sys.exit(1)
         signal.signal(signal.SIGINT, int_handler)
 
+        daemon.set_user(config.get('user', None))
         loop.run()
     except (KeyboardInterrupt, IOError, OSError) as e:
         logging.error(e)
diff --git a/shadowsocks/server.py b/shadowsocks/server.py
index 8eed4ad..345ec0b 100755
--- a/shadowsocks/server.py
+++ b/shadowsocks/server.py
@@ -86,6 +86,8 @@ def main():
             loop = eventloop.EventLoop()
             dns_resolver.add_to_loop(loop)
             list(map(lambda s: s.add_to_loop(loop), tcp_servers + udp_servers))
+
+            daemon.set_user(config.get('user', None))
             loop.run()
         except (KeyboardInterrupt, IOError, OSError) as e:
             logging.error(e)
diff --git a/shadowsocks/utils.py b/shadowsocks/utils.py
index f791a2a..b12f54d 100644
--- a/shadowsocks/utils.py
+++ b/shadowsocks/utils.py
@@ -89,7 +89,11 @@ def check_config(config):
     if config.get('password') in [b'mypassword']:
         logging.error('DON\'T USE DEFAULT PASSWORD! Please change it in your '
                       'config.json!')
-        exit(1)
+        sys.exit(1)
+    if config.get('user', None) is not None:
+        if os.name != 'posix':
+            logging.error('user can be used only on Unix')
+            sys.exit(1)
 
 
 def get_config(is_local):
@@ -97,11 +101,11 @@ def get_config(is_local):
                         format='%(levelname)-s: %(message)s')
     if is_local:
         shortopts = 'hd:s:b:p:k:l:m:c:t:vq'
-        longopts = ['help', 'fast-open', 'pid-file=', 'log-file=']
+        longopts = ['help', 'fast-open', 'pid-file=', 'log-file=', 'user=']
     else:
         shortopts = 'hd:s:p:k:m:c:t:vq'
         longopts = ['help', 'fast-open', 'pid-file=', 'log-file=', 'workers=',
-                    'forbidden-ip=']
+                    'forbidden-ip=', 'user=']
     try:
         config_path = find_config()
         optlist, args = getopt.getopt(sys.argv[1:], shortopts, longopts)
@@ -147,6 +151,8 @@ def get_config(is_local):
                 config['fast_open'] = True
             elif key == '--workers':
                 config['workers'] = int(value)
+            elif key == '--user':
+                config['user'] = to_str(value)
             elif key == '--forbidden-ip':
                 config['forbidden_ip'] = to_str(value).split(',')
             elif key in ('-h', '--help'):
@@ -247,9 +253,7 @@ def print_help(is_local):
 
 
 def print_local_help():
-    print('''usage: sslocal [-h] -s SERVER_ADDR [-p SERVER_PORT]
-               [-b LOCAL_ADDR] [-l LOCAL_PORT] -k PASSWORD [-m METHOD]
-               [-t TIMEOUT] [-c CONFIG] [--fast-open] [-v] -[d] [-q]
+    print('''usage: sslocal [OPTION]...
 A fast tunnel proxy that helps you bypass firewalls.
 
 You can supply configurations via either config file or command line arguments.
@@ -270,6 +274,7 @@ General options:
   -d start/stop/restart  daemon mode
   --pid-file PID_FILE    pid file for daemon mode
   --log-file LOG_FILE    log file for daemon mode
+  --user USER            username to run as
   -v, -vv                verbose mode
   -q, -qq                quiet mode, only show warnings/errors
 
@@ -278,9 +283,7 @@ Online help: <https://github.com/shadowsocks/shadowsocks>
 
 
 def print_server_help():
-    print('''usage: ssserver [-h] [-s SERVER_ADDR] [-p SERVER_PORT] -k PASSWORD
-                -m METHOD [-t TIMEOUT] [-c CONFIG] [--fast-open]
-                [--workers WORKERS] [-v] [-d start] [-q]
+    print('''usage: ssserver [OPTION]...
 A fast tunnel proxy that helps you bypass firewalls.
 
 You can supply configurations via either config file or command line arguments.
@@ -301,6 +304,7 @@ General options:
   -d start/stop/restart  daemon mode
   --pid-file PID_FILE    pid file for daemon mode
   --log-file LOG_FILE    log file for daemon mode
+  --user USER            username to run as
   -v, -vv                verbose mode
   -q, -qq                quiet mode, only show warnings/errors