Akkariiin/dev

configloader.py

@@ -1,5 +1,5 @@
 #!/usr/bin/python
-# -*- coding: UTF-8 -*-
+# -*- coding: utf-8 -*-
 import importloader
 
 g_config = None

db_transfer.py

@@ -1,5 +1,5 @@
 #!/usr/bin/python
-# -*- coding: UTF-8 -*-
+# -*- coding: utf-8 -*-
 
 import logging
 import time
@@ -9,6 +9,7 @@
 from shadowsocks import common, shell, lru_cache, obfs
 from configloader import load_config, get_config
 import importloader
+import copy
 
 switchrule = None
 db_instance = None
@@ -80,8 +81,10 @@ def push_db_all_user(self):
 	def del_server_out_of_bound_safe(self, last_rows, rows):
 		#停止超流量的服务
 		#启动没超流量的服务
+		keymap = {}
 		try:
 			switchrule = importloader.load('switchrule')
+			keymap = switchrule.getRowMap()
 		except Exception as e:
 			logging.error('load switchrule.py fail')
 		cur_servers = {}
@@ -106,7 +109,10 @@ def del_server_out_of_bound_safe(self, last_rows, rows):
 			read_config_keys = ['method', 'obfs', 'obfs_param', 'protocol', 'protocol_param', 'forbidden_ip', 'forbidden_port', 'speed_limit_per_con', 'speed_limit_per_user']
 			for name in read_config_keys:
 				if name in row and row[name]:
-					cfg[name] = row[name]
+					if name in keymap:
+						cfg[keymap[name]] = row[name]
+					else:
+						cfg[name] = row[name]
 
 			merge_config_keys = ['password'] + read_config_keys
 			for name in cfg.keys():
@@ -392,11 +398,17 @@ def pull_db_all_user(self):
 		return rows
 
 	def pull_db_users(self, conn):
+		keys = copy.copy(self.key_list)
 		try:
 			switchrule = importloader.load('switchrule')
-			keys = switchrule.getKeys(self.key_list)
+			keymap = switchrule.getRowMap()
+			for key in keymap:
+				if keymap[key] in keys:
+					keys.remove(keymap[key])
+				keys.append(key)
+			keys = switchrule.getKeys(keys)
 		except Exception as e:
-			keys = self.key_list
+			logging.error('load switchrule.py fail')
 
 		cur = conn.cursor()
 		cur.execute("SELECT " + ','.join(keys) + " FROM user")
@@ -520,11 +532,17 @@ def update_all_user(self, dt_transfer):
 		return update_transfer
 
 	def pull_db_users(self, conn):
+		keys = copy.copy(self.key_list)
 		try:
 			switchrule = importloader.load('switchrule')
-			keys = switchrule.getKeys(self.key_list)
+			keymap = switchrule.getRowMap()
+			for key in keymap:
+				if keymap[key] in keys:
+					keys.remove(keymap[key])
+				keys.append(key)
+			keys = switchrule.getKeys(keys)
 		except Exception as e:
-			keys = self.key_list
+			logging.error('load switchrule.py fail')
 
 		cur = conn.cursor()
 

importloader.py

@@ -1,5 +1,5 @@
 #!/usr/bin/python
-# -*- coding: UTF-8 -*-
+# -*- coding: utf-8 -*-
 
 def load(name):
 	try:

server_pool.py

@@ -117,14 +117,14 @@ def new_server(self, port, user_config):
 			else:
 				a_config = self.config.copy()
 				a_config.update(user_config)
-				if len(a_config['server_ipv6']) > 2 and a_config['server_ipv6'][0] == "[" and a_config['server_ipv6'][-1] == "]":
+				if len(a_config['server_ipv6']) > 2 and a_config['server_ipv6'][0] == b"[" and a_config['server_ipv6'][-1] == b"]":
 					a_config['server_ipv6'] = a_config['server_ipv6'][1:-1]
-				a_config['server'] = a_config['server_ipv6']
+				a_config['server'] = common.to_str(a_config['server_ipv6'])
 				a_config['server_port'] = port
 				a_config['max_connect'] = 128
 				a_config['method'] = common.to_str(a_config['method'])
 				try:
-					logging.info("starting server at [%s]:%d" % (common.to_str(a_config['server']), port))
+					logging.info("starting server at [%s]:%d" % (a_config['server'], port))
 
 					tcp_server = tcprelay.TCPRelay(a_config, self.dns_resolver, False, stat_counter=self.stat_counter)
 					tcp_server.add_to_loop(self.loop)
@@ -134,7 +134,7 @@ def new_server(self, port, user_config):
 					udp_server.add_to_loop(self.loop)
 					self.udp_ipv6_servers_pool.update({port: udp_server})
 
-					if common.to_str(a_config['server_ipv6']) == "::":
+					if a_config['server_ipv6'] == "::":
 						ipv6_ok = True
 				except Exception as e:
 					logging.warn("IPV6 %s " % (e,))
@@ -150,7 +150,7 @@ def new_server(self, port, user_config):
 				a_config['max_connect'] = 128
 				a_config['method'] = common.to_str(a_config['method'])
 				try:
-					logging.info("starting server at %s:%d" % (common.to_str(a_config['server']), port))
+					logging.info("starting server at %s:%d" % (a_config['server'], port))
 
 					tcp_server = tcprelay.TCPRelay(a_config, self.dns_resolver, False)
 					tcp_server.add_to_loop(self.loop)

shadowsocks/asyncdns.py

@@ -27,12 +27,12 @@
 if __name__ == '__main__':
     import sys
     import inspect
+
     file_path = os.path.dirname(os.path.realpath(inspect.getfile(inspect.currentframe())))
     sys.path.insert(0, os.path.join(file_path, '../'))
 
 from shadowsocks import common, lru_cache, eventloop, shell
 
-
 CACHE_SWEEP_INTERVAL = 30
 
 VALID_HOSTNAME = re.compile(br"(?!-)[A-Z\d_-]{1,63}(?<!-)$", re.IGNORECASE)
@@ -77,6 +77,7 @@
 QTYPE_NS = 2
 QCLASS_IN = 1
 
+
 def detect_ipv6_supprot():
     if 'has_ipv6' in dir(socket):
         try:
@@ -89,8 +90,10 @@ def detect_ipv6_supprot():
     print('IPv6 not support')
     return False
 
+
 IPV6_CONNECTION_SUPPORT = detect_ipv6_supprot()
 
+
 def build_address(address):
     address = address.strip(b'.')
     labels = address.split(b'.')
@@ -175,7 +178,7 @@ def parse_record(data, offset, question=False):
         )
         ip = parse_ip(record_type, data, record_rdlength, offset + nlen + 10)
         return nlen + 10 + record_rdlength, \
-            (name, ip, record_type, record_class, record_ttl)
+               (name, ip, record_type, record_class, record_ttl)
     else:
         record_type, record_class = struct.unpack(
             '!HH', data[offset + nlen:offset + nlen + 4]
@@ -209,7 +212,7 @@ def parse_response(data):
             if not header:
                 return None
             res_id, res_qr, res_tc, res_ra, res_rcode, res_qdcount, \
-                res_ancount, res_nscount, res_arcount = header
+            res_ancount, res_nscount, res_arcount = header
 
             qds = []
             ans = []
@@ -266,14 +269,22 @@ def __str__(self):
 
 
 class DNSResolver(object):
-
-    def __init__(self):
+    def __init__(self, black_hostname_list=None):
         self._loop = None
         self._hosts = {}
         self._hostname_status = {}
         self._hostname_to_cb = {}
         self._cb_to_hostname = {}
         self._cache = lru_cache.LRUCache(timeout=300)
+        # read black_hostname_list from config
+        if type(black_hostname_list) != list:
+            self._black_hostname_list = []
+        else:
+            self._black_hostname_list = list(map(
+                (lambda t: t if type(t) == bytes else t.encode('utf8')),
+                black_hostname_list
+            ))
+        logging.info('black_hostname_list init as : ' + str(self._black_hostname_list))
         self._sock = None
         self._servers = None
         self._parse_resolv()
@@ -377,7 +388,7 @@ def _handle_data(self, data):
             ip = None
             for answer in response.answers:
                 if answer[1] in (QTYPE_A, QTYPE_AAAA) and \
-                        answer[2] == QCLASS_IN:
+                                answer[2] == QCLASS_IN:
                     ip = answer[0]
                     break
             if IPV6_CONNECTION_SUPPORT:
@@ -462,19 +473,22 @@ def resolve(self, hostname, callback):
             ip = self._hosts[hostname]
             callback((hostname, ip), None)
         elif hostname in self._cache:
-            logging.debug('hit cache: %s', hostname)
+            logging.debug('hit cache: %s ==>> %s', hostname, self._cache[hostname])
             ip = self._cache[hostname]
             callback((hostname, ip), None)
+        elif any(hostname.endswith(t) for t in self._black_hostname_list):
+            callback(None, Exception('hostname <%s> is block by the black hostname list' % hostname))
+            return
         else:
             if not is_valid_hostname(hostname):
                 callback(None, Exception('invalid hostname: %s' % hostname))
                 return
             if False:
                 addrs = socket.getaddrinfo(hostname, 0, 0,
-                                       socket.SOCK_DGRAM, socket.SOL_UDP)
+                                           socket.SOCK_DGRAM, socket.SOL_UDP)
                 if addrs:
                     af, socktype, proto, canonname, sa = addrs[0]
-                    logging.debug('DNS resolve %s %s' % (hostname, sa[0]) )
+                    logging.debug('DNS resolve %s %s' % (hostname, sa[0]))
                     self._cache[hostname] = sa[0]
                     callback((hostname, sa[0]), None)
                     return
@@ -506,7 +520,11 @@ def close(self):
 
 
 def test():
-    dns_resolver = DNSResolver()
+    black_hostname_list = [
+        'baidu.com',
+        'yahoo.com',
+    ]
+    dns_resolver = DNSResolver(black_hostname_list=black_hostname_list)
     loop = eventloop.EventLoop()
     dns_resolver.add_to_loop(loop)
 
@@ -521,16 +539,20 @@ def callback(result, error):
             # TODO: what can we assert?
             print(result, error)
             counter += 1
-            if counter == 9:
+            if counter == 12:
                 dns_resolver.close()
                 loop.stop()
+
         a_callback = callback
         return a_callback
 
-    assert(make_callback() != make_callback())
+    assert (make_callback() != make_callback())
 
     dns_resolver.resolve(b'google.com', make_callback())
     dns_resolver.resolve('google.com', make_callback())
+    dns_resolver.resolve('baidu.com', make_callback())
+    dns_resolver.resolve('map.baidu.com', make_callback())
+    dns_resolver.resolve('yahoo.com', make_callback())
     dns_resolver.resolve('example.com', make_callback())
     dns_resolver.resolve('ipv6.google.com', make_callback())
     dns_resolver.resolve('www.facebook.com', make_callback())
@@ -546,10 +568,25 @@ def callback(result, error):
                          'ooooooooooooooooooooooooooooooooooooooooooooooooooo'
                          'ooooooooooooooooooooooooooooooooooooooooooooooooooo'
                          'long.hostname', make_callback())
-
     loop.run()
+    # test black_hostname_list
+    dns_resolver = DNSResolver(black_hostname_list=[])
+    assert type(dns_resolver._black_hostname_list) == list
+    assert len(dns_resolver._black_hostname_list) == 0
+    dns_resolver.close()
+    dns_resolver = DNSResolver(black_hostname_list=123)
+    assert type(dns_resolver._black_hostname_list) == list
+    assert len(dns_resolver._black_hostname_list) == 0
+    dns_resolver.close()
+    dns_resolver = DNSResolver(black_hostname_list=None)
+    assert type(dns_resolver._black_hostname_list) == list
+    assert len(dns_resolver._black_hostname_list) == 0
+    dns_resolver.close()
+    dns_resolver = DNSResolver()
+    assert type(dns_resolver._black_hostname_list) == list
+    assert dns_resolver._black_hostname_list.__len__() == 0
+    dns_resolver.close()
 
 
 if __name__ == '__main__':
     test()
-

shadowsocks/common.py

@@ -121,7 +121,19 @@ def is_ip(address):
     return False
 
 
+def sync_str_bytes(obj, target_example):
+    """sync (obj)'s type to (target_example)'s type"""
+    if type(obj) != type(target_example):
+        if type(target_example) == str:
+            obj = to_str(obj)
+        if type(target_example) == bytes:
+            obj = to_bytes(obj)
+    return obj
+
+
 def match_regex(regex, text):
+    # avoid 'cannot use a string pattern on a bytes-like object'
+    regex = sync_str_bytes(regex, text)
     regex = re.compile(regex)
     for item in regex.findall(text):
         return True
@@ -253,7 +265,7 @@ def __init__(self, addrs):
         list(map(self.add_network, addrs))
 
     def add_network(self, addr):
-        if addr is "":
+        if addr == "":
             return
         block = addr.split('/')
         addr_family = is_ip(block[0])
@@ -265,9 +277,9 @@ def add_network(self, addr):
             ip = (hi << 64) | lo
         else:
             raise Exception("Not a valid CIDR notation: %s" % addr)
-        if len(block) is 1:
+        if len(block) == 1:
             prefix_size = 0
-            while (ip & 1) == 0 and ip is not 0:
+            while (ip & 1) == 0 and ip != 0:
                 ip >>= 1
                 prefix_size += 1
             logging.warn("You did't specify CIDR routing prefix size for %s, "
@@ -381,12 +393,12 @@ def test_inet_conv():
 
 def test_parse_header():
     assert parse_header(b'\x03\x0ewww.google.com\x00\x50') == \
-        (0, b'www.google.com', 80, 18)
+        (0, ADDRTYPE_HOST, b'www.google.com', 80, 18)
     assert parse_header(b'\x01\x08\x08\x08\x08\x00\x35') == \
-        (0, b'8.8.8.8', 53, 7)
+        (0, ADDRTYPE_IPV4, b'8.8.8.8', 53, 7)
     assert parse_header((b'\x04$\x04h\x00@\x05\x08\x05\x00\x00\x00\x00\x00'
                          b'\x00\x10\x11\x00\x50')) == \
-        (0, b'2404:6800:4005:805::1011', 80, 19)
+        (0, ADDRTYPE_IPV6, b'2404:6800:4005:805::1011', 80, 19)
 
 
 def test_pack_header():
@@ -411,7 +423,25 @@ def test_ip_network():
     assert 'www.google.com' not in ip_network
 
 
+def test_sync_str_bytes():
+    assert sync_str_bytes(b'a\.b', b'a\.b') == b'a\.b'
+    assert sync_str_bytes('a\.b', b'a\.b') == b'a\.b'
+    assert sync_str_bytes(b'a\.b', 'a\.b') == 'a\.b'
+    assert sync_str_bytes('a\.b', 'a\.b') == 'a\.b'
+    pass
+
+
+def test_match_regex():
+    assert match_regex(br'a\.b', b'abc,aaa,aaa,b,aaa.b,a.b')
+    assert match_regex(r'a\.b', b'abc,aaa,aaa,b,aaa.b,a.b')
+    assert match_regex(br'a\.b', b'abc,aaa,aaa,b,aaa.b,a.b')
+    assert match_regex(r'a\.b', b'abc,aaa,aaa,b,aaa.b,a.b')
+    assert match_regex(r'\bgoogle\.com\b', b' google.com ')
+    pass
+
 if __name__ == '__main__':
+    test_sync_str_bytes()
+    test_match_regex()
     test_inet_conv()
     test_parse_header()
     test_pack_header()