mirror of
https://github.com/lionsoul2014/ip2region.git
synced 2025-12-08 19:25:22 +00:00
215 lines
6.8 KiB
Python
215 lines
6.8 KiB
Python
# Created by leolin49 on 2022/7/7.
|
|
# Copyright (C) 2022 leolin49. All rights reserved.
|
|
|
|
import logging
|
|
import sys
|
|
import time
|
|
from datetime import datetime
|
|
|
|
import xdb.maker as mk
|
|
import xdb.index as idx
|
|
import xdb.searcher as sc
|
|
import xdb.util as util
|
|
|
|
# format log
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s-%(name)s-%(lineno)s-%(levelname)s - %(message)s')
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def print_help():
|
|
print("ip2region xdb maker")
|
|
print("{} [command] [command options]".format(sys.argv[0]))
|
|
print("Command: ")
|
|
print(" gen generate the binary db file")
|
|
print(" search binary xdb search test")
|
|
print(" bench binary xdb bench test")
|
|
|
|
|
|
def gen_db():
|
|
src_file, dst_file = "", ""
|
|
index_policy = idx.VectorIndexPolicy
|
|
# check input argv
|
|
for i in range(2, len(sys.argv)):
|
|
r = sys.argv[i]
|
|
if len(r) < 5:
|
|
continue
|
|
if not r.startswith("--"):
|
|
continue
|
|
s_idx = r.index("=")
|
|
if s_idx < 0:
|
|
print("missing = for args pair '{}'".format(r))
|
|
return
|
|
if r[2:s_idx] == "src":
|
|
src_file = r[s_idx+1:]
|
|
elif r[2:s_idx] == "dst":
|
|
dst_file = r[s_idx+1:]
|
|
elif r[2:s_idx] == "index":
|
|
index_policy = idx.index_policy_from_string(r[s_idx+1:])
|
|
else:
|
|
print("undefined option `{}`".format(r))
|
|
return
|
|
if src_file == "" or dst_file == "":
|
|
print("{} gen [command options]".format(sys.argv[0]))
|
|
print("options:")
|
|
print(" --src string source ip text file path")
|
|
print(" --dst string destination binary xdb file path")
|
|
return
|
|
|
|
start_time = time.time()
|
|
# make the binary file
|
|
maker = mk.new_maker(index_policy, src_file, dst_file)
|
|
maker.init()
|
|
maker.start()
|
|
maker.end()
|
|
|
|
logging.info("Done, elapsed: {:.0f}m{:.0f}s".format((time.time() - start_time) / 60, (time.time() - start_time) % 60))
|
|
|
|
|
|
def test_search():
|
|
db_file = ""
|
|
for i in range(2, len(sys.argv)):
|
|
r = sys.argv[i]
|
|
if len(r) < 5:
|
|
continue
|
|
if not r.startswith("--"):
|
|
continue
|
|
e_idx = r.index("=")
|
|
if e_idx < 0:
|
|
print("missing = for args pair '{}'".format(r))
|
|
return
|
|
if r[2:e_idx] == "db":
|
|
db_file = r[e_idx+1:]
|
|
else:
|
|
print("undefined option `{}`".format(r))
|
|
return
|
|
if db_file == "":
|
|
print("{} search [command options]".format(sys.argv[0]))
|
|
print("options:")
|
|
print(" --db string ip2region binary xdb file path")
|
|
return
|
|
cb = sc.XdbSearcher.loadContentFromFile(dbfile=db_file)
|
|
searcher = sc.XdbSearcher(contentBuff=cb)
|
|
print("ip2region xdb search test program, commands:\nloadIndex : load the vector index for search "
|
|
"speedup.\nclearIndex: clear the vector index.\nquit : exit the test program")
|
|
while True:
|
|
print("ip2region>> ", end="")
|
|
line = input()
|
|
|
|
# command interception and execution
|
|
if line == "loadIndex":
|
|
searcher.loadVectorIndexFromFile(dbfile=db_file)
|
|
print("vector index cached")
|
|
continue
|
|
elif line == "clearIndex":
|
|
# FIXME need to add 'clearVectorIndex' method in searcher
|
|
print("vector index cleared")
|
|
continue
|
|
elif line == "quit":
|
|
break
|
|
|
|
ip = util.checkip(line)
|
|
if ip == -1:
|
|
print("invalid ip address `{}`".format(line))
|
|
continue
|
|
|
|
s_tm = datetime.now()
|
|
region = searcher.search(ip)
|
|
# TODO calculate io count in `searcher.search` method
|
|
print("\x1b[0;32m[region:{}, took:{:.0f}s]\x1b[0m".format(
|
|
region, (datetime.now().microsecond - s_tm.microsecond) / 1000)
|
|
)
|
|
|
|
|
|
def test_bench():
|
|
db_file, src_file = "", ""
|
|
ignore_error = False
|
|
for i in range(2, len(sys.argv)):
|
|
r = sys.argv[i]
|
|
if len(r) < 5:
|
|
continue
|
|
if not r.startswith("--"):
|
|
continue
|
|
s_idx = r.index("=")
|
|
if s_idx < 0:
|
|
print("missing = for args pair '{}'".format(r))
|
|
return
|
|
if r[2:s_idx] == "db":
|
|
db_file = r[s_idx + 1:]
|
|
elif r[2:s_idx] == "src":
|
|
src_file = r[s_idx + 1:]
|
|
elif r[2:s_idx] == "ignore-error":
|
|
v = r[s_idx + 1:]
|
|
if v == "true" or v == "1":
|
|
ignore_error = True
|
|
elif v == "false" or v == "0":
|
|
ignore_error = False
|
|
else:
|
|
print("invalid value for ignore-error option, could be false/0 or true/1")
|
|
return
|
|
else:
|
|
print("undefined option `{}`".format(r))
|
|
return
|
|
|
|
if db_file == "" or src_file == "":
|
|
print("{} bench [command options]".format(sys.argv[0]))
|
|
print("options:")
|
|
print(" --db string ip2region binary xdb file path")
|
|
print(" --src string source ip text file path")
|
|
print(" --ignore-error bool keep going if bench failed")
|
|
return
|
|
|
|
cb = sc.XdbSearcher.loadContentFromFile(dbfile=db_file)
|
|
searcher = sc.XdbSearcher(contentBuff=cb)
|
|
cnt, err_cnt, s_tm = 0, 0, time.time()
|
|
with open(src_file, 'r', encoding="utf-8") as f:
|
|
lines = f.read().splitlines()
|
|
for line in lines:
|
|
ps = line.split("|", maxsplit=2)
|
|
if len(ps) != 3:
|
|
print("invalid ip segment line `{}`".format(line))
|
|
return
|
|
sip = util.checkip(ps[0])
|
|
if sip == -1:
|
|
print("invalid ip address `{}`".format(line))
|
|
return
|
|
eip = util.checkip(ps[1])
|
|
if eip == -1:
|
|
print("invalid ip address `{}`".format(line))
|
|
return
|
|
print("try to bench segment: `{}`", line)
|
|
mip = util.mid_ip(sip, eip)
|
|
for ip in [sip, util.mid_ip(sip, mip), mip, util.mid_ip(mip, eip), eip]:
|
|
print("|-try to bench ip '{}' ...".format(util.long2ip(ip)), end="")
|
|
region = searcher.search(ip)
|
|
|
|
# check the region info
|
|
cnt += 1
|
|
if region != ps[2]:
|
|
err_cnt += 1
|
|
print(" --[Failed] ({} != {})".format(region, ps[2]))
|
|
if not ignore_error:
|
|
return
|
|
else:
|
|
print(" --[Ok]")
|
|
print("Bench finished, [count: {}, failed: {}, took: {:.3f}s]".format(cnt, err_cnt, time.time() - s_tm))
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) < 2:
|
|
print_help()
|
|
return
|
|
|
|
cmd = sys.argv[1].lower()
|
|
if cmd == "gen":
|
|
gen_db()
|
|
elif cmd == "search":
|
|
test_search()
|
|
elif cmd == "bench":
|
|
test_bench()
|
|
else:
|
|
print_help()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|