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" 12qtype = unbound.RR_TYPE_A 13qclass = unbound.RR_CLASS_IN 14 15def create_context(config_file="ub.lookup.conf", asyncflag=False): 16 """ 17 Create an unbound context to use for testing. 18 19 """ 20 ctx = unbound.ub_ctx() 21 status = ctx.config(config_file) 22 if status != 0: 23 print("read config failed with status: {}".format(status)) 24 sys.exit(1) 25 ctx.set_async(asyncflag) 26 return ctx 27 28 29def callback(data, status, result): 30 """ 31 Callback for background workers. 32 33 """ 34 if status == 0: 35 data['rcode'] = result.rcode 36 data['secure'] = result.secure 37 if result.havedata: 38 data['data'] = result.data 39 data['was_ratelimited'] = result.was_ratelimited 40 data['done'] = True 41 42 43def test_resolve(ctx): 44 """ 45 Test resolving a domain with a foreground worker. 46 47 """ 48 status, result = ctx.resolve(qname, qtype, qclass) 49 if status == 0 and result.havedata: 50 print("Resolve: {}".format(result.data.address_list)) 51 else: 52 print("Failed resolve with: {}".format(status)) 53 54 55def test_async_resolve(ctx): 56 """ 57 Test resolving a domain with a background worker. 58 59 """ 60 cb_data = dict(done=False) 61 retval, async_id = ctx.resolve_async(qname, cb_data, callback, qtype, qclass) 62 while retval == 0 and not cb_data['done']: 63 time.sleep(0.1) 64 retval = ctx.process() 65 66 if cb_data.get('data'): 67 print("Async resolve: {}".format(cb_data['data'].address_list)) 68 else: 69 print("Failed async resolve with: {}".format(retval)) 70 71 72def test_ratelimit_fg_on(ctx): 73 """ 74 Test resolving a ratelimited domain with a foreground worker. 75 76 """ 77 ctx.set_option("ratelimit:", "1") 78 ctx.set_option("ratelimit-factor:", "0") 79 status, result = ctx.resolve(qname, qtype, qclass) 80 if status == 0 and result.was_ratelimited: 81 print("Ratelimit-fg-on: pass") 82 else: 83 print("Failed ratelimit-fg-on with: {}".format(status)) 84 85 86def test_ratelimit_fg_off(ctx): 87 """ 88 Test resolving a non-ratelimited domain with a foreground worker. 89 90 """ 91 status, result = ctx.resolve(qname, qtype, qclass) 92 if status == 0 and result.havedata: 93 print("Ratelimit-fg-off: {}".format(result.data.address_list)) 94 else: 95 print("Failed ratelimit-fg-off with: {}".format(status)) 96 97 98def test_ratelimit_bg_on(ctx): 99 """ 100 Test resolving a ratelimited domain with a background worker. 101 102 """ 103 ctx.set_option("ratelimit:", "1") 104 ctx.set_option("ratelimit-factor:", "0") 105 cb_data = dict(done=False) 106 retval, async_id = ctx.resolve_async(qname, cb_data, callback, qtype, qclass) 107 while retval == 0 and not cb_data['done']: 108 time.sleep(0.1) 109 retval = ctx.process() 110 111 if cb_data.get('was_ratelimited'): 112 print("Ratelimit-bg-on: pass") 113 else: 114 print("Failed ratelimit-bg-on with: {}".format(status)) 115 116 117def test_ratelimit_bg_off(ctx): 118 """ 119 Test resolving a non-ratelimited domain with a background worker. 120 121 """ 122 cb_data = dict(done=False) 123 retval, async_id = ctx.resolve_async(qname, cb_data, callback, qtype, qclass) 124 while retval == 0 and not cb_data['done']: 125 time.sleep(0.1) 126 retval = ctx.process() 127 128 if cb_data.get('data'): 129 print("Ratelimit-bg-off: {}".format(cb_data['data'].address_list)) 130 else: 131 print("Failed ratelimit-bg-off with: {}".format(status)) 132 133 134test_resolve(create_context()) 135test_async_resolve(create_context(asyncflag=True)) 136test_ratelimit_fg_on(create_context()) 137test_ratelimit_fg_off(create_context()) 138test_ratelimit_bg_on(create_context(asyncflag=True)) 139test_ratelimit_bg_off(create_context(asyncflag=True)) 140 141sys.exit(0) 142