优化ip2region查询

common/ipreg.py

@@ -10,6 +10,8 @@ import socket
 import struct
 import sys
 
+from config import settings
+
 
 class IpRegInfo(object):
     __INDEX_BLOCK_LENGTH = 12
@@ -39,7 +41,8 @@ class IpRegInfo(object):
             self.__dbBinStr = self.__f.read()  # read all the contents in file
             self.__indexSPtr = self.get_long(self.__dbBinStr, 0)
             self.__indexLPtr = self.get_long(self.__dbBinStr, 4)
-            self.__indexCount = int((self.__indexLPtr - self.__indexSPtr) / self.__INDEX_BLOCK_LENGTH) + 1
+            self.__indexCount = int((self.__indexLPtr - self.__indexSPtr) /
+                                    self.__INDEX_BLOCK_LENGTH) + 1
 
         l, h, data_ptr = (0, self.__indexCount, 0)
         while l <= h:
@@ -67,14 +70,16 @@ class IpRegInfo(object):
         " binary search method
         " param: ip
         """
-        if not ip.isdigit(): ip = self.ip2long(ip)
+        if not ip.isdigit():
+            ip = self.ip2long(ip)
 
         if self.__indexCount == 0:
             self.__f.seek(0)
             super_block = self.__f.read(8)
             self.__indexSPtr = self.get_long(super_block, 0)
             self.__indexLPtr = self.get_long(super_block, 4)
-            self.__indexCount = int((self.__indexLPtr - self.__indexSPtr) / self.__INDEX_BLOCK_LENGTH) + 1
+            self.__indexCount = int((self.__indexLPtr - self.__indexSPtr) /
+                                    self.__INDEX_BLOCK_LENGTH) + 1
 
         l, h, data_ptr = (0, self.__indexCount, 0)
         while l <= h:
@@ -208,10 +213,9 @@ class IpRegInfo(object):
         self.__f.seek(data_ptr)
         data = self.__f.read(data_len)
 
-        return {
-            "city_id": self.get_long(data, 0),
-            "region": data[4:]
-        }
+        info = {"city_id": self.get_long(data, 0),
+                "region": data[4:].decode('utf-8')}
+        return info
 
     @staticmethod
     def ip2long(ip):
@@ -244,3 +248,25 @@ class IpRegInfo(object):
         self.__dbBinStr = None
         self.__headerPtr = None
         self.__headerSip = None
+
+
+class IpRegData(IpRegInfo):
+    def __init__(self):
+        path = settings.data_storage_dir.joinpath('ip2region.db')
+        IpRegInfo.__init__(self, path)
+
+    def query(self, ip, algorithm='memory'):
+        algorithms = ['memory', 'binary', 'btree']
+        if algorithm not in algorithms:
+            raise Exception(f"Only three query algorithms are supported: {algorithms}")
+        if algorithm == 'memory':
+            result = self.memory_search(ip)
+        elif algorithm == 'binary':
+            result = self.binary_search(ip)
+        else:
+            result = self.btree_search(ip)
+        addr_list = result.get('region').split('|')
+        isp = addr_list[-1]
+        addr = ''.join(filter(lambda x: x != '0', addr_list[:-1]))
+        info = {'addr': addr, 'isp': isp}
+        return info

common/resolve.py

@@ -6,13 +6,12 @@ from config import settings
 from common import utils
 from common.ipasn import IPAsnInfo
 from common.ipgeo import IpGeoInfo
-from common.ipreg import IpRegInfo
+from common.ipreg import IpRegData
 
 
 ip_asn = IPAsnInfo()
 ip_geo = IpGeoInfo()
-db_path = settings.data_storage_dir.joinpath('ip2region.db')
-ip_reg = IpRegInfo(db_path)
+ip_reg = IpRegData()
 
 
 def filter_subdomain(data):
@@ -107,7 +106,7 @@ def gen_infos(data, qname, info, infos):
                   f'{ip_geo.get_region(ip)} ' \
                   f'{ip_geo.get_city(ip)}'
             locs.append(loc)
-            reg = ip_reg.memory_search(ip).get('region').decode('utf-8')
+            reg = ip_reg.query(ip).get('addr')
             regs.append(reg)
             info['resolve'] = 1
             info['reason'] = 'OK'