纯真ip数据库读取(python)

python无名氏

1. #!/usr/bin/python3
   
2. # encoding: utf-8
   
3. 
   
4. import socket
   
5. import codecs
   
6. import mmap
   
7. from struct import pack, unpack
   
8. 
   
9. class QQWry(object):
   
10.     def __init__(self, path):
    
11.         self.path = path
    
12.         self.db = None
    
13.         self.open_db()
    
14.         self.idx_start, self.idx_end = self._read_idx()
    
15.         self.total = (self.idx_end - self.idx_start) // 7 + 1
    
16.     def decode_str(self,old):
    
17.         try:
    
18.             return old.decode('gbk').encode('utf-8')
    
19.         except:
    
20.             if old[-1] == b'\x96':
    
21.                 try:
    
22.                     return old[:-1].decode('gbk').encode('utf-8') + b'?'
    
23.                 except:
    
24.                     pass
    
25. 
    
26.             return 'Invalid'.encode('utf-8')
    
27.     def open_db(self):
    
28.         if not self.db:
    
29.             self.db = open(self.path, 'rb')
    
30.             self.db = mmap.mmap(self.db.fileno(), 0, access=1)
    
31.         return self.db
    
32. 
    
33.     def _read_idx(self):
    
34.         self.db.seek(0)
    
35.         start = unpack('I', self.db.read(4))[0]
    
36.         end = unpack('I', self.db.read(4))[0]
    
37.         return start, end
    
38. 
    
39.     def version(self):
    
40.         ip_end_offset = self.read_offset(self.idx_end + 4)
    
41.         a_raw, b_raw = self.read_record(ip_end_offset + 4)
    
42.         return self.decode_str(a_raw + b_raw).decode('utf-8')
    
43. 
    
44.     def read_ip(self, off, seek=True):
    
45.         if seek:
    
46.             self.db.seek(off)
    
47.         buf = self.db.read(4)
    
48.         return unpack('I', buf)[0]
    
49. 
    
50.     def read_offset(self, off, seek=True):
    
51.         if seek:
    
52.             self.db.seek(off)
    
53.         buf = self.db.read(3)
    
54.         return unpack('I', buf + b'\0')[0]
    
55. 
    
56.     def read_string(self, offset):
    
57.         if offset == 0:
    
58.             return 'N/A1'.encode('utf-8')
    
59.         flag = self.get_flag(offset)
    
60.         if flag == 0:
    
61.             return 'N/A2'.encode('utf-8')
    
62.         elif flag == 2:
    
63.             offset = self.read_offset(offset + 1)
    
64.             return self.read_string(offset)
    
65.         self.db.seek(offset)
    
66.         raw_string = b''
    
67.         while True:
    
68.             x = self.db.read(1)
    
69.             if x == b'\0':
    
70.                 break
    
71.             raw_string += x
    
72.         return raw_string
    
73. 
    
74.     def get_flag(self, offset):
    
75.         self.db.seek(offset)
    
76.         c = self.db.read(1)
    
77.         if not c:
    
78.             return 0
    
79.         return c[0]
    
80. 
    
81.     def read_record(self, offset):
    
82.         self.db.seek(offset)
    
83.         flag = self.db.read(1)[0]
    
84.         if flag == 1:
    
85.             buf = self.db.read(3)
    
86.             a_offset = unpack('I', buf + b'\0')[0]
    
87.             a_raw = self.read_string(a_offset)
    
88.             a_flag = self.get_flag(a_offset)
    
89.             if a_flag == 2:
    
90.                 b_raw = self.read_string(a_offset + 4)
    
91.             else:
    
92.                 b_raw = self.read_string(a_offset + len(a_raw) + 1)
    
93.         elif flag == 2:
    
94.             buf = self.db.read(3)
    
95.             a_offset = unpack('I', buf + b'\0')[0]
    
96.             a_raw = self.read_string(a_offset)
    
97.             b_raw = self.read_string(offset + 4)
    
98.         else:
    
99.             a_raw = self.read_string(offset)
    
100.             b_raw = self.read_string(offset + len(a_raw) + 1)
     
101.         return a_raw, b_raw
     
102. 
     
103.     def output(self, output_file='ip.txt'):
     
104.         fp = codecs.open(output_file, 'w', 'utf8')
     
105.         idx = self.idx_start
     
106.         while idx <= self.idx_end:
     
107.             ip_int = self.read_ip(idx)
     
108.             ip_start = socket.inet_ntoa(pack('!I', ip_int))
     
109.             ip_end_offset = self.read_offset(idx + 4)
     
110.             ip_int = self.read_ip(ip_end_offset)
     
111.             ip_end = socket.inet_ntoa(pack('!I', ip_int))
     
112.             a_raw, b_raw = self.read_record(ip_end_offset + 4)
     
113.             a_info = self.decode_str(a_raw)
     
114.             b_info = self.decode_str(b_raw)
     
115.             fp.write('%15s\t%15s\t%s,%s\n' % (
     
116.                 ip_start, ip_end,
     
117.                 a_info.decode('utf8'), b_info.decode('utf8')))
     
118.             idx += 7
     
119.         fp.close()
     
120. 
     
121.     def find(self, ip, l, r):
     
122.         if r - l <= 1:
     
123.             return l
     
124.         m = (l + r) // 2
     
125.         offset = self.idx_start + m * 7
     
126.         new_ip = self.read_ip(offset)
     
127.         if ip < new_ip:
     
128.             return self.find(ip, l, m)
     
129.         else:
     
130.             return self.find(ip, m, r)
     
131. 
     
132.     def query(self, ip):
     
133.         ip = unpack('!I', socket.inet_aton(ip))[0]
     
134.         i = self.find(ip, 0, self.total - 1)
     
135.         o = self.idx_start + i * 7
     
136.         o2 = self.read_offset(o + 4)
     
137.         (c, a) = self.read_record(o2 + 4)
     
138.         return (self.decode_str(c).decode('utf-8'), self.decode_str(a).decode('utf-8'))
     
139. 
     
140.     def __del__(self):
     
141.         if self.db:
     
142.             self.db.close()
     
143. 
     
144. def main():
     
145.     dbpath = "qqwry.dat"
     
146.     ip = "114.114.114.114"
     
147.     qqwry = QQWry(dbpath)
     
148.     c, a = qqwry.query(ip)
     
149.     print(qqwry.version())
     
150.     print('%s %s--%s' % (ip, c, a))
     
151. 
     
152. if __name__ == '__main__':
     
153.     main()
     

复制代码

从别人那看到的，原来是python2的，改成python3了
数据库太大了，放不上来，自己找吧