From eb94bd1cc34eb4b228e675011236d090329d5a7f Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Mon, 12 Jan 2015 22:30:03 +0800
Subject: [PATCH] support forbidden iplist

---
 shadowsocks/tcprelay.py | 9 +++++++++
 shadowsocks/utils.py    | 6 +++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py
index 79bd1a5..6afcad4 100644
--- a/shadowsocks/tcprelay.py
+++ b/shadowsocks/tcprelay.py
@@ -123,6 +123,10 @@ class TCPRelayHandler(object):
         self._downstream_status = WAIT_STATUS_INIT
         self._client_address = local_sock.getpeername()[:2]
         self._remote_address = None
+        if 'forbidden_ip' in self._config:
+            self._forbidden_iplist = self._config['forbidden_ip']
+        else:
+            self._forbidden_iplist = None
         if is_local:
             self._chosen_server = self._get_a_server()
         fd_to_handlers[local_sock.fileno()] = self
@@ -331,6 +335,10 @@ class TCPRelayHandler(object):
         if len(addrs) == 0:
             raise Exception("getaddrinfo failed for %s:%d" % (ip,  port))
         af, socktype, proto, canonname, sa = addrs[0]
+        if self._forbidden_iplist:
+            if common.to_str(sa[0]) in self._forbidden_iplist:
+                raise Exception('IP %s is in forbidden list, reject' %
+                                common.to_str(sa[0]))
         remote_sock = socket.socket(af, socktype, proto)
         self._remote_sock = remote_sock
         self._fd_to_handlers[remote_sock.fileno()] = self
@@ -346,6 +354,7 @@ class TCPRelayHandler(object):
         if result:
             ip = result[1]
             if ip:
+
                 try:
                     self._stage = STAGE_CONNECTING
                     remote_addr = ip
diff --git a/shadowsocks/utils.py b/shadowsocks/utils.py
index 0247d0f..a51c965 100644
--- a/shadowsocks/utils.py
+++ b/shadowsocks/utils.py
@@ -100,7 +100,8 @@ def get_config(is_local):
         longopts = ['help', 'fast-open', 'pid-file=', 'log-file=']
     else:
         shortopts = 'hd:s:p:k:m:c:t:vq'
-        longopts = ['help', 'fast-open', 'pid-file=', 'log-file=', 'workers=']
+        longopts = ['help', 'fast-open', 'pid-file=', 'log-file=', 'workers=',
+                    'forbidden-ip=']
     try:
         config_path = find_config()
         optlist, args = getopt.getopt(sys.argv[1:], shortopts, longopts)
@@ -146,6 +147,8 @@ def get_config(is_local):
                 config['fast_open'] = True
             elif key == '--workers':
                 config['workers'] = int(value)
+            elif key == '--forbidden-ip':
+                config['forbidden_ip'] = to_str(value).split(',')
             elif key in ('-h', '--help'):
                 if is_local:
                     print_local_help()
@@ -286,6 +289,7 @@ Proxy options:
   -t TIMEOUT             timeout in seconds, default: 300
   --fast-open            use TCP_FASTOPEN, requires Linux 3.7+
   --workers WORKERS      number of workers, available on Unix/Linux
+  --forbidden-ip IPLIST  comma seperated IP list forbidden to connect
 
 General options:
   -d start/stop/restart  daemon mode