1#!/usr/bin/env python 2# 3# Test for pyunbound lookup. 4# BSD licensed. 5# 6import sys 7import time 8 9import unbound 10 11qname = "www.example.com" 12qname2 = "www2.example.com" 13qtype = unbound.RR_TYPE_A 14qclass = unbound.RR_CLASS_IN 15 16 17def create_context(config_file="ub.lookup.conf", asyncflag=False): 18 """ 19 Create an unbound context to use for testing. 20 21 """ 22 ctx = unbound.ub_ctx() 23 status = ctx.config(config_file) 24 if status != 0: 25 print("read config failed with status: {}".format(status)) 26 sys.exit(1) 27 ctx.set_async(asyncflag) 28 return ctx 29 30 31def callback(data, status, result): 32 """ 33 Callback for background workers. 34 35 """ 36 if status == 0: 37 data['rcode'] = result.rcode 38 data['secure'] = result.secure 39 if result.havedata: 40 data['data'] = result.data 41 data['was_ratelimited'] = result.was_ratelimited 42 data['done'] = True 43 44 45def test_resolve(ctx): 46 """ 47 Test resolving a domain with a foreground worker. 48 49 """ 50 status, result = ctx.resolve(qname, qtype, qclass) 51 if status == 0 and result.havedata: 52 print("Resolve: {}".format(result.data.address_list)) 53 else: 54 print("Failed resolve with: {}".format(status)) 55 56 57def test_async_resolve(ctx): 58 """ 59 Test resolving a domain with a background worker. 60 61 """ 62 cb_data = dict(done=False) 63 retval, async_id = ctx.resolve_async(qname, cb_data, callback, qtype, qclass) 64 while retval == 0 and not cb_data['done']: 65 time.sleep(0.1) 66 retval = ctx.process() 67 68 if cb_data.get('data'): 69 print("Async resolve: {}".format(cb_data['data'].address_list)) 70 else: 71 print("Failed async resolve with: {}".format(retval)) 72 73 74def test_ratelimit_bg_on(ctx): 75 """ 76 Test resolving a ratelimited domain with a background worker. 77 78 """ 79 ctx.set_option("ratelimit:", "1") 80 ctx.set_option("ratelimit-factor:", "0") 81 total_runs = 6 82 success_threshold = 4 # 2/3*total_runs 83 successes = 0 84 for i in range(total_runs): 85 cb_data = dict(done=False) 86 cb_data2 = dict(done=False) 87 retval, async_id = ctx.resolve_async(qname, cb_data, callback, qtype, qclass) 88 retval, async_id = ctx.resolve_async(qname2, cb_data2, callback, qtype, qclass) 89 90 while retval == 0 and not (cb_data['done'] and cb_data['done']): 91 time.sleep(0.1) 92 retval = ctx.process() 93 94 if bool(cb_data.get('was_ratelimited')) ^ bool(cb_data2.get('was_ratelimited')): 95 successes += 1 96 if successes >= success_threshold: 97 break 98 time.sleep(1) 99 if successes >= success_threshold: 100 print("Ratelimit-bg-on: pass") 101 else: 102 print("Failed ratelimit-bg-on") 103 104 105test_resolve(create_context()) 106test_async_resolve(create_context(asyncflag=True)) 107test_ratelimit_bg_on(create_context(asyncflag=True)) 108 109sys.exit(0) 110