make country code lookup a background task
This commit is contained in:
parent
52f315b395
commit
c65b465560
|
@ -111,14 +111,23 @@ async def try_ip(db, server_ip, checks, opts: Options):
|
||||||
from asyncio import sleep
|
from asyncio import sleep
|
||||||
|
|
||||||
entries = []
|
entries = []
|
||||||
|
deferred = []
|
||||||
|
|
||||||
success = True
|
success = True
|
||||||
|
|
||||||
|
def maybe_put_ip(ip):
|
||||||
|
from asyncio import QueueFull
|
||||||
|
try:
|
||||||
|
opts.ips.put_nowait(ip)
|
||||||
|
except QueueFull:
|
||||||
|
deferred.append(ip)
|
||||||
|
|
||||||
def finisher(done, pending):
|
def finisher(done, pending):
|
||||||
nonlocal success
|
nonlocal success
|
||||||
for task in done:
|
for task in done:
|
||||||
res, ip, check = task.result()
|
res, ip, check = task.result()
|
||||||
entry = process_result(res, ip, check, opts)
|
entry = process_result(res, ip, check, opts)
|
||||||
|
map(maybe_put_ip, entry.addrs)
|
||||||
entries.append(entry)
|
entries.append(entry)
|
||||||
if not entry.success:
|
if not entry.success:
|
||||||
if opts.early_stopping and success: # only cancel once
|
if opts.early_stopping and success: # only cancel once
|
||||||
|
@ -151,12 +160,8 @@ async def try_ip(db, server_ip, checks, opts: Options):
|
||||||
else:
|
else:
|
||||||
await pooler()
|
await pooler()
|
||||||
|
|
||||||
if opts.ipinfo is not None:
|
for ip in deferred:
|
||||||
await opts.ipinfo.find_country(server_ip, db)
|
await opts.ips.put(ip)
|
||||||
for entry in entries:
|
|
||||||
for addr in entry.addrs:
|
|
||||||
await opts.ipinfo.find_country(addr, db)
|
|
||||||
opts.ipinfo.flush()
|
|
||||||
|
|
||||||
if not opts.dry:
|
if not opts.dry:
|
||||||
for entry in entries:
|
for entry in entries:
|
||||||
|
@ -179,27 +184,34 @@ async def try_ip(db, server_ip, checks, opts: Options):
|
||||||
async def sync_database(db, opts: Options):
|
async def sync_database(db, opts: Options):
|
||||||
from .ips import china, blocks
|
from .ips import china, blocks
|
||||||
|
|
||||||
if db is None:
|
|
||||||
return
|
|
||||||
|
|
||||||
# TODO: handle addresses that were removed from respodns.ips.china.
|
# TODO: handle addresses that were removed from respodns.ips.china.
|
||||||
for ips, kw in ((china, "china"), (blocks, "block_target")):
|
for ips, kw in ((china, "china"), (blocks, "block_target")):
|
||||||
for ip in ips:
|
for ip in ips:
|
||||||
kwargs = dict()
|
kwargs = {kw: True}
|
||||||
kwargs[kw] = True
|
if db is None:
|
||||||
if opts.ipinfo is not None:
|
db.modify_address(ip, **kwargs)
|
||||||
kwargs["country_code"] = await opts.ipinfo.find_country(ip)
|
await opts.ips.put(ip)
|
||||||
db.modify_address(ip, **kwargs)
|
|
||||||
if opts.ipinfo is not None:
|
|
||||||
opts.ipinfo.flush()
|
async def locate_ips(db, opts: Options):
|
||||||
|
seen = set()
|
||||||
|
while (ip := await opts.ips.get()) is not None:
|
||||||
|
if opts.ipinfo is not None and ip not in seen:
|
||||||
|
seen.add(ip)
|
||||||
|
code = await opts.ipinfo.find_country(ip)
|
||||||
|
if db is not None:
|
||||||
|
db.modify_address(ip, country_code=code)
|
||||||
|
opts.ipinfo.flush()
|
||||||
|
|
||||||
|
|
||||||
async def main(db, filepaths, checks, opts: Options):
|
async def main(db, filepaths, checks, opts: Options):
|
||||||
from .util import make_pooler
|
from .util import make_pooler
|
||||||
from asyncio import sleep, create_task
|
from asyncio import sleep, create_task, Queue
|
||||||
from sys import stdin, stderr
|
from sys import stdin, stderr
|
||||||
|
|
||||||
|
opts.ips = Queue()
|
||||||
syncing = create_task(sync_database(db, opts))
|
syncing = create_task(sync_database(db, opts))
|
||||||
|
geoip = create_task(locate_ips(db, opts))
|
||||||
|
|
||||||
def finisher(done, pending):
|
def finisher(done, pending):
|
||||||
for task in done:
|
for task in done:
|
||||||
|
@ -229,6 +241,7 @@ async def main(db, filepaths, checks, opts: Options):
|
||||||
stderr.flush()
|
stderr.flush()
|
||||||
if not first:
|
if not first:
|
||||||
await sleep(opts.ip_wait)
|
await sleep(opts.ip_wait)
|
||||||
|
await opts.ips.put(ip)
|
||||||
await pooler(try_ip(db, ip, checks, opts))
|
await pooler(try_ip(db, ip, checks, opts))
|
||||||
|
|
||||||
if opts.blocking_file_io:
|
if opts.blocking_file_io:
|
||||||
|
@ -249,3 +262,5 @@ async def main(db, filepaths, checks, opts: Options):
|
||||||
|
|
||||||
await pooler()
|
await pooler()
|
||||||
await syncing
|
await syncing
|
||||||
|
await opts.ips.put(None) # end of queue
|
||||||
|
await geoip
|
||||||
|
|
|
@ -4,8 +4,10 @@ from dataclasses import dataclass
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class Options:
|
class Options:
|
||||||
|
# TODO: move these out of Options, since they're really not.
|
||||||
execution: object = None
|
execution: object = None
|
||||||
ipinfo: object = None
|
ipinfo: object = None
|
||||||
|
ips: object = None
|
||||||
|
|
||||||
ip_simul: int = 15 # how many IPs to connect to at once
|
ip_simul: int = 15 # how many IPs to connect to at once
|
||||||
domain_simul: int = 3 # how many domains per IP to request at once
|
domain_simul: int = 3 # how many domains per IP to request at once
|
||||||
|
|
Loading…
Reference in New Issue
Block a user