Ticket #2506: geoip_compress.py

File geoip_compress.py, 3.8 KB (added by nickm, 9 years ago)
Line 
1#!/usr/bin/python
2
3import collections
4import struct
5import itertools
6import bisect
7
8def encodeInt(x):
9    bytes = []
10    while True:
11        lowbits = x & 0x7f
12        x >>= 7
13        if x:
14            bytes.append(chr(lowbits))
15        else:
16            bytes.append(chr(lowbits | 0x80))
17            return "".join(bytes)
18
19def encodeU32(x): return struct.pack("!L",x)
20def decodeU32(x): return struct.unpack("!L",x)[0]
21
22class GeoIP:
23    def __init__(self):
24        self.rangesPerCountry = collections.defaultdict(lambda: 0)
25        self.rangeSizes = collections.defaultdict(lambda: 0)
26        self.ranges = [] # Tuples of lo, hi, countryname
27
28        self.countryCodes = {} # cc -> integer
29        self.rangeLenCodes = {} # range length -> integer
30        self.countryByCode = {} # integer -> cc
31        self.rangeLenByCode = {} # integer -> range length
32
33    def _addRange(self, lo, hi, cc):
34        rangeLen = hi - lo + 1
35        self.rangesPerCountry[cc] += 1
36        self.rangeSizes[rangeLen] += 1
37        self.ranges.append( (lo, hi, cc) )
38
39    def readCSV(self, f):
40        lasthi = -1
41        for line in f:
42            v = line.replace('"', "").split(",")
43            lo, hi, cc = v[2:5]
44            lo = int(lo)
45            hi = int(hi)
46
47            if lo != lasthi + 1:
48                self._addRange(lasthi + 1, lo - 1, "??")
49
50            self._addRange(lo, hi, cc)
51            lasthi = hi
52
53        if lasthi != (1<<32)-1:
54            self._addRange(lasthi + 1, (1<<32)-1, "??")
55
56    def _mktab(self, d):
57        items = sorted(d.iteritems(), key=lambda i:i[1], reverse=True)
58        codeByValue = {}
59        valueByCode = []
60        i = 0
61        for val, _ in items:
62            codeByValue[val] = i
63            valueByCode.append(val)
64            i += 1
65        return codeByValue, valueByCode
66
67    def generateCodings(self):
68        self.countryCodes,self.countryByCode = self._mktab(self.rangesPerCountry)
69        self.rangeLenCodes,self.rangeLenByCode = self._mktab(self.rangeSizes)
70
71    def writeCountryCodes(self, f):
72        f.write("@CC@")
73        f.write(encodeU32(len(self.countryByCode)))
74        for cc in self.countryByCode:
75            f.write(cc)
76        f.write("@cc@")
77
78    def writeRunlenList(self, f):
79        f.write("@RL@")
80        f.write(encodeU32(len(self.rangeLenByCode)))
81        for rl in self.rangeLenByCode:
82            f.write(encodeU32(rl))
83        f.write("@rl@")
84
85    def writeRanges(self, f_ranges, f_idx):
86        f_ranges.write("@IP@")
87        f_ranges.write(encodeU32(len(self.ranges)))
88        rangePos = 8
89
90        f_idx.write("@IX@")
91        f_idx.write(encodeU32(len(self.ranges) / 256))
92        idxEntries = 0
93        ip = 0
94
95        for rangenum, (lo, hi, cc) in itertools.izip(
96                    xrange(len(self.ranges)), self.ranges):
97
98            if rangenum and (rangenum % 256) == 0:
99                f_idx.write(encodeU32(ip))
100                f_idx.write(encodeU32(rangePos))
101                idxEntries += 1
102
103            rangelen = hi - lo + 1
104
105            rc = encodeInt(self.rangeLenCodes[rangelen])
106            cc = encodeInt(self.countryCodes[cc])
107
108            ip += rangelen
109            rangePos += len(rc)+len(cc)
110            f_ranges.write(rc)
111            f_ranges.write(cc)
112            assert f_ranges.tell() == rangePos
113
114        f_ranges.write("@ip@")
115        assert idxEntries == len(self.ranges) / 256
116        f_idx.write("@ix@")
117
118def realignFile(f):
119    p = f.tell()
120    if p&3:
121        f.write("\0\0\0\0"[:(4-(p&3))])
122
123def generate():
124    G = GeoIP()
125    G.readCSV(open("GeoIPCountryWhois.csv"))
126    G.generateCodings()
127    meta = open("geoip.idx", 'wb')
128    data = open("geoip.dat", 'wb')
129    G.writeCountryCodes(meta)
130    realignFile(meta)
131    G.writeRunlenList(meta)
132    realignFile(meta)
133    G.writeRanges(data, meta)
134    meta.close()
135    data.close()
136
137if __name__ == '__main__':
138    generate()