addr字段改为ip2location 添加ip2region字段

common/database.py

@@ -79,7 +79,8 @@ class Database(object):
                    f'ttl text,'
                    f'cidr text,'
                    f'asn text,'
-                   f'addr text,'
+                   f'ip2region text,'
+                   f'ip2location text,'
                    f'resolver text,'
                    f'module text,'
                    f'source text,'
@@ -106,11 +107,11 @@ class Database(object):
                     f'id, type, alive, resolve, request, new, url, subdomain,'
                     f'port, level, cname, content, public, status, reason,'
                     f'title, banner, header, response, times, ttl, cidr, asn,'
-                    f'addr, resolver, module, source, elapse, find, brute, valid) '
+                    f'ip2region, ip2location, resolver, module, source, elapse, find, brute, valid) '
                     f'values (:id, :type, :alive, :resolve, :request, :new,'
                     f':url, :subdomain, :port, :level, :cname, :content,'
                     f':public, :status, :reason, :title, :banner, :header,'
-                    f':response, :times, :ttl, :cidr, :asn, :addr, :resolver,'
+                    f':response, :times, :ttl, :cidr, :asn, :ip2region, :ip2location, :resolver,'
                     f':module, :source, :elapse, :find, :brute, :valid)', results)
             except Exception as e:
                 logger.log('ERROR', e)
@@ -234,8 +235,9 @@ class Database(object):
         table_name = table_name.replace('.', '_')
         query = f'select id, type, new, alive, request, resolve, url, ' \
                 f'subdomain, level, cname, content, public, port, status, ' \
-                f'reason, title, banner, times, ttl, cidr, asn, addr, resolver, ' \
+                f'reason, title, banner, times, ttl, cidr, asn, ip2region, ' \
-                f'module, source, elapse, find, brute, valid from "{table_name}"'
+                f'ip2location, resolver, module, source, elapse, find, brute, valid ' \
+                f'from "{table_name}"'
         if alive and limit:
             if limit in ['resolve', 'request']:
                 where = f' where {limit} = 1'

common/ipreg.py

@@ -0,0 +1,246 @@
+# encoding = uft-8
+"""
+" ip2region python searcher client module
+"
+" Author: koma<komazhang@foxmail.com>
+" Date : 2015-11-06
+"""
+import io
+import socket
+import struct
+import sys
+
+
+class IpRegInfo(object):
+    __INDEX_BLOCK_LENGTH = 12
+    __TOTAL_HEADER_LENGTH = 8192
+
+    __f = None
+    __headerSip = []
+    __headerPtr = []
+    __headerLen = 0
+    __indexSPtr = 0
+    __indexLPtr = 0
+    __indexCount = 0
+    __dbBinStr = ''
+
+    def __init__(self, db_file):
+        self.init_database(db_file)
+
+    def memory_search(self, ip):
+        """
+        " memory search method
+        " param: ip
+        """
+        if not ip.isdigit():
+            ip = self.ip2long(ip)
+
+        if self.__dbBinStr == '':
+            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
+
+        l, h, data_ptr = (0, self.__indexCount, 0)
+        while l <= h:
+            m = int((l + h) >> 1)
+            p = self.__indexSPtr + m * self.__INDEX_BLOCK_LENGTH
+            sip = self.get_long(self.__dbBinStr, p)
+
+            if ip < sip:
+                h = m - 1
+            else:
+                eip = self.get_long(self.__dbBinStr, p + 4)
+                if ip > eip:
+                    l = m + 1;
+                else:
+                    data_ptr = self.get_long(self.__dbBinStr, p + 8)
+                    break
+
+        if data_ptr == 0:
+            raise Exception("Data pointer not found")
+
+        return self.return_data(data_ptr)
+
+    def binary_search(self, ip):
+        """
+        " binary search method
+        " param: 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
+
+        l, h, data_ptr = (0, self.__indexCount, 0)
+        while l <= h:
+            m = int((l + h) >> 1)
+            p = m * self.__INDEX_BLOCK_LENGTH
+
+            self.__f.seek(self.__indexSPtr + p)
+            buffer = self.__f.read(self.__INDEX_BLOCK_LENGTH)
+            sip = self.get_long(buffer, 0)
+            if ip < sip:
+                h = m - 1
+            else:
+                eip = self.get_long(buffer, 4)
+                if ip > eip:
+                    l = m + 1
+                else:
+                    data_ptr = self.get_long(buffer, 8)
+                    break
+
+        if data_ptr == 0:
+            raise Exception("Data pointer not found")
+
+        return self.return_data(data_ptr)
+
+    def btree_search(self, ip):
+        """
+        " b-tree search method
+        " param: ip
+        """
+        if not ip.isdigit():
+            ip = self.ip2long(ip)
+
+        if len(self.__headerSip) < 1:
+            header_len = 0
+            # pass the super block
+            self.__f.seek(8)
+            # read the header block
+            b = self.__f.read(self.__TOTAL_HEADER_LENGTH)
+            # parse the header block
+            for i in range(0, len(b), 8):
+                sip = self.get_long(b, i)
+                ptr = self.get_long(b, i + 4)
+                if ptr == 0:
+                    break
+                self.__headerSip.append(sip)
+                self.__headerPtr.append(ptr)
+                header_len += 1
+            self.__headerLen = header_len
+
+        l, h, sptr, eptr = (0, self.__headerLen, 0, 0)
+        while l <= h:
+            m = int((l + h) >> 1)
+
+            if ip == self.__headerSip[m]:
+                if m > 0:
+                    sptr = self.__headerPtr[m - 1]
+                    eptr = self.__headerPtr[m]
+                else:
+                    sptr = self.__headerPtr[m]
+                    eptr = self.__headerPtr[m + 1]
+                break
+
+            if ip < self.__headerSip[m]:
+                if m == 0:
+                    sptr = self.__headerPtr[m]
+                    eptr = self.__headerPtr[m + 1]
+                    break
+                elif ip > self.__headerSip[m - 1]:
+                    sptr = self.__headerPtr[m - 1]
+                    eptr = self.__headerPtr[m]
+                    break
+                h = m - 1
+            else:
+                if m == self.__headerLen - 1:
+                    sptr = self.__headerPtr[m - 1]
+                    eptr = self.__headerPtr[m]
+                    break
+                elif ip <= self.__headerSip[m + 1]:
+                    sptr = self.__headerPtr[m]
+                    eptr = self.__headerPtr[m + 1]
+                    break
+                l = m + 1
+
+        if sptr == 0:
+            raise Exception("Index pointer not found")
+
+        index_len = eptr - sptr
+        self.__f.seek(sptr)
+        index = self.__f.read(index_len + self.__INDEX_BLOCK_LENGTH)
+
+        l, h, data_prt = (0, int(index_len / self.__INDEX_BLOCK_LENGTH), 0)
+        while l <= h:
+            m = int((l + h) >> 1)
+            offset = int(m * self.__INDEX_BLOCK_LENGTH)
+            sip = self.get_long(index, offset)
+
+            if ip < sip:
+                h = m - 1
+            else:
+                eip = self.get_long(index, offset + 4)
+                if ip > eip:
+                    l = m + 1;
+                else:
+                    data_prt = self.get_long(index, offset + 8)
+                    break
+
+        if data_prt == 0:
+            raise Exception("Data pointer not found")
+
+        return self.return_data(data_prt)
+
+    def init_database(self, db_file):
+        """
+        " initialize the database for search
+        " param: dbFile
+        """
+        try:
+            self.__f = io.open(db_file, "rb")
+        except IOError as e:
+            print("[Error]: %s" % e)
+            sys.exit()
+
+    def return_data(self, data_ptr):
+        """
+        " get ip data from db file by data start ptr
+        " param: dsptr
+        """
+        data_len = (data_ptr >> 24) & 0xFF
+        data_ptr = data_ptr & 0x00FFFFFF
+
+        self.__f.seek(data_ptr)
+        data = self.__f.read(data_len)
+
+        return {
+            "city_id": self.get_long(data, 0),
+            "region": data[4:]
+        }
+
+    @staticmethod
+    def ip2long(ip):
+        _ip = socket.inet_aton(ip)
+        return struct.unpack("!L", _ip)[0]
+
+    @staticmethod
+    def is_ip(ip):
+        p = ip.split(".")
+        if len(p) != 4:
+            return False
+        for pp in p:
+            if not pp.isdigit():
+                return False
+            if len(pp) > 3:
+                return False
+            if int(pp) > 255:
+                return False
+        return True
+
+    @staticmethod
+    def get_long(b, offset):
+        if len(b[offset:offset + 4]) == 4:
+            return struct.unpack('I', b[offset:offset + 4])[0]
+        return 0
+
+    def close(self):
+        if self.__f is not None:
+            self.__f.close()
+        self.__dbBinStr = None
+        self.__headerPtr = None
+        self.__headerSip = None

common/module.py

@@ -304,7 +304,8 @@ class Module(object):
                       'ttl': None,
                       'cidr': None,
                       'asn': None,
-                      'addr': None,
+                      'ip2region': None,
+                      'ip2location': None,
                       'resolver': None,
                       'module': self.module,
                       'source': self.source,
@@ -336,7 +337,8 @@ class Module(object):
                 public = record.get('public')
                 cidr = record.get('cidr')
                 asn = record.get('asn')
-                addr = record.get('addr')
+                ip2region = record.get('ip2region')
+                ip2location = record.get('ip2location')
                 if isinstance(cname, list):
                     cname = ','.join(cname)
                     content = ','.join(content)
@@ -366,7 +368,8 @@ class Module(object):
                           'ttl': ttl,
                           'cidr': cidr,
                           'asn': asn,
-                          'addr': addr,
+                          'ip2region': ip2region,
+                          'ip2location': ip2location,
                           'resolver': resolver,
                           'module': self.module,
                           'source': self.source,

common/resolve.py

@@ -7,6 +7,7 @@ from common import utils
 from common.database import Database
 from common.ipasn import IPAsnInfo
 from common.ipgeo import IpGeoInfo
+from common.ipreg import IpRegInfo
 
 
 def filter_subdomain(data):
@@ -75,6 +76,8 @@ def deal_output(output_path):
     records = dict()  # 用来记录所有域名解析数据
     ip_asn = IPAsnInfo()
     ip_geo = IpGeoInfo
+    db_path = setting.data_storage_dir.joinpath('ip2region.db')
+    ip_reg = IpRegInfo(db_path)
     with open(output_path) as fd:
         for line in fd:
             line = line.strip()
@@ -108,7 +111,8 @@ def deal_output(output_path):
             ttls = list()
             cidrs = list()
             asns = list()
-            addrs = list()
+            locs = list()
+            regs = list()
             answers = data.get('answers')
             for answer in answers:
                 if answer.get('type') == 'A':
@@ -123,10 +127,12 @@ def deal_output(output_path):
                     asn_info = ip_asn.find(ip)
                     cidrs.append(asn_info.get('cidr'))
                     asns.append(asn_info.get('asn'))
-                    addr = f'{ip_geo.get_country_long(ip)} ' \
+                    loc = f'{ip_geo.get_country_long(ip)} ' \
-                           f'{ip_geo.get_region(ip)} ' \
+                          f'{ip_geo.get_region(ip)} ' \
-                           f'{ip_geo.get_city(ip)}'
+                          f'{ip_geo.get_city(ip)}'
-                    addrs.append(addr)
+                    locs.append(loc)
+                    reg = ip_reg.memory_search(ip).get('region').decode('utf-8')
+                    regs.append(reg)
                     record['resolve'] = 1
                     record['reason'] = status
                     record['cname'] = ','.join(cname)
@@ -135,7 +141,8 @@ def deal_output(output_path):
                     record['ttl'] = ','.join(ttls)
                     record['cidr'] = ','.join(cidrs)
                     record['asn'] = ','.join(asns)
-                    record['addr'] = ','.join(addrs)
+                    record['ip2location'] = ','.join(locs)
+                    record['ip2region'] = ','.join(regs)
                     records[qname] = record
             if not flag:
                 record['alive'] = 0