1#cython: embedsignature=True 2#cython: language_level=2 3# Copyright (c) 2015-2019 by Farsight Security, Inc. 4# 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16 17include "dnstable.pxi" 18 19import math 20 21RRSET = DNSTABLE_QUERY_TYPE_RRSET 22RDATA_IP = DNSTABLE_QUERY_TYPE_RDATA_IP 23RDATA_RAW = DNSTABLE_QUERY_TYPE_RDATA_RAW 24RDATA_NAME = DNSTABLE_QUERY_TYPE_RDATA_NAME 25 26class DnstableException(Exception): 27 pass 28 29class Timeout(DnstableException): 30 pass 31 32cdef class entry(object): 33 cdef dnstable_entry *_instance 34 cdef dnstable_entry_type etype 35 cdef dict d 36 37 def __cinit__(self): 38 self._instance = NULL 39 40 def __dealloc__(self): 41 dnstable_entry_destroy(&self._instance) 42 43 def __init__(self): 44 pass 45 46 cdef from_c(self, dnstable_entry *ent, iszone=False): 47 cdef dnstable_res res 48 cdef uint16_t u16 49 cdef uint64_t u64 50 cdef const uint8_t *data 51 cdef size_t len_data 52 cdef size_t sz 53 cdef size_t i 54 55 self.d = {} 56 self.etype = dnstable_entry_get_type(ent) 57 58 if not self.etype in (DNSTABLE_ENTRY_TYPE_RRSET, DNSTABLE_ENTRY_TYPE_RDATA): 59 raise DnstableException, 'unhandled entry type' 60 61 # rrname 62 res = dnstable_entry_get_rrname(ent, &data, &len_data) 63 if res == dnstable_res_success: 64 self.d['rrname'] = data[:len_data] 65 66 # rrtype 67 res = dnstable_entry_get_rrtype(ent, &u16) 68 if res == dnstable_res_success: 69 self.d['rrtype'] = u16 70 71 # count 72 res = dnstable_entry_get_count(ent, &u64) 73 if res == dnstable_res_success: 74 self.d['count'] = u64 75 76 # time_first 77 res = dnstable_entry_get_time_first(ent, &u64) 78 if res == dnstable_res_success: 79 if iszone: 80 self.d['zone_time_first'] = u64 81 else: 82 self.d['time_first'] = u64 83 84 # time_last 85 res = dnstable_entry_get_time_last(ent, &u64) 86 if res == dnstable_res_success: 87 if iszone: 88 self.d['zone_time_last'] = u64 89 else: 90 self.d['time_last'] = u64 91 92 # rdata 93 res = dnstable_entry_get_num_rdata(ent, &sz) 94 if res == dnstable_res_success: 95 self.d['rdata'] = [] 96 for i from 0 <= i < sz: 97 res = dnstable_entry_get_rdata(ent, i, &data, &len_data) 98 if res != dnstable_res_success: 99 raise DnstableException, 'dnstable_entry_get_rdata() failed' 100 self.d['rdata'].append(data[:len_data]) 101 102 if self.etype == DNSTABLE_ENTRY_TYPE_RRSET: 103 # bailiwick 104 res = dnstable_entry_get_bailiwick(ent, &data, &len_data) 105 if res == dnstable_res_success: 106 self.d['bailiwick'] = data[:len_data].decode('utf-8') 107 108 if iszone: 109 dnstable_entry_set_iszone(ent, iszone) 110 111 self._instance = ent 112 113 def __repr__(self): 114 return self.to_text() 115 116 def to_dict(self): 117 return self.d 118 119 def to_fmt_dict(self): 120 import copy 121 import wdns 122 d = copy.deepcopy(self.d) 123 if 'rdata' in d: 124 new_rdata_list = [] 125 for rdata in d['rdata']: 126 new_rdata_list.append(repr(wdns.rdata(rdata, wdns.CLASS_IN, self.d['rrtype']))) 127 d['rdata'] = new_rdata_list 128 if 'rrname' in d: 129 d['rrname'] = wdns.domain_to_str(d['rrname'].encode('utf-8')) 130 if 'bailiwick' in d: 131 d['bailiwick'] = wdns.domain_to_str(d['bailiwick'].encode('utf-8')) 132 if 'rrtype' in d: 133 d['rrtype'] = wdns.rrtype_to_str(d['rrtype']) 134 return d 135 136 def to_text(self): 137 cdef char *res 138 res = dnstable_entry_to_text(self._instance) 139 s = res.decode('utf-8') 140 free(res) 141 return s 142 143 def to_json(self, rfc3339_time = False, rdata_always_array = False): 144 cdef char *res 145 cdef dnstable_formatter *f 146 147 if not rfc3339_time and not rdata_always_array: 148 # for typical case, use simply C function 149 res = dnstable_entry_to_json(self._instance) 150 s = res.decode('utf-8') 151 free(res) 152 return s 153 else: 154 # create and destroy the formatter each time 155 f = dnstable_formatter_init() 156 dnstable_formatter_set_output_format(f, dnstable_output_format_json) 157 if rfc3339_time: 158 dnstable_formatter_set_date_format(f, dnstable_date_format_rfc3339) 159 else: 160 dnstable_formatter_set_date_format(f, dnstable_date_format_unix) 161 dnstable_formatter_set_rdata_array(f, rdata_always_array) 162 163 res = dnstable_entry_format(f, self._instance) 164 s = res.decode('utf-8') 165 free(res) 166 dnstable_formatter_destroy(&f) 167 return s 168 169@cython.internal 170cdef class iteritems(object): 171 cdef dnstable_iter *_instance 172 cdef object iszone 173 cdef object q 174 175 def __cinit__(self): 176 self._instance = NULL 177 178 def __init__(self, iszone=False): 179 self.iszone = iszone 180 181 def __dealloc__(self): 182 dnstable_iter_destroy(&self._instance) 183 184 def __iter__(self): 185 return self 186 187 def __next__(self): 188 cdef dnstable_res res 189 cdef dnstable_entry *ent 190 191 if self._instance == NULL: 192 raise StopIteration 193 194 res = dnstable_iter_next(self._instance, &ent) 195 196 if res == dnstable_res_failure: 197 raise StopIteration 198 elif res == dnstable_res_timeout: 199 raise Timeout 200 201 d = entry() 202 d.from_c(ent, self.iszone) 203 return d 204 205cdef class query(object): 206 cdef dnstable_query *_instance 207 cdef readonly int qtype 208 cdef readonly str data 209 cdef readonly str rrtype 210 cdef readonly str bailiwick 211 212 def __cinit__(self): 213 self._instance = NULL 214 215 def __init__(self, qtype, str data, str rrtype=None, str bailiwick=None, time_first_before=None, time_first_after=None, time_last_before=None, time_last_after=None, timeout=None, aggregate=True, uint64_t offset=0): 216 cdef dnstable_res 217 cdef timespec ts 218 cdef uint64_t tm 219 220 self.data = data 221 self.rrtype = rrtype 222 self.bailiwick = bailiwick 223 224 if not qtype in (RRSET, RDATA_IP, RDATA_RAW, RDATA_NAME): 225 raise DnstableException, 'invalid qtype' 226 self._instance = dnstable_query_init(qtype) 227 self.qtype = qtype 228 229 res = dnstable_query_set_data(self._instance, data.encode('UTF-8')) 230 if res != dnstable_res_success: 231 raise DnstableException, 'dnstable_query_set_data() failed: %s' % dnstable_query_get_error(self._instance) 232 233 if rrtype: 234 res = dnstable_query_set_rrtype(self._instance, rrtype.encode('UTF-8')) 235 if res != dnstable_res_success: 236 raise DnstableException, 'dnstable_query_set_rrtype() failed: %s' % dnstable_query_get_error(self._instance) 237 238 if offset != 0: 239 res = dnstable_query_set_offset(self._instance, offset) 240 if res != dnstable_res_success: 241 raise DnstableException, 'dnstable_query_set_offset() failed: %s' % dnstable_query_get_error(self._instance) 242 243 res = dnstable_query_set_aggregated(self._instance, aggregate) 244 if res != dnstable_res_success: 245 raise DnstableException, 'dnstable_query_set_aggregated() failed: %s' % dnstable_query_get_error(self._instance) 246 247 if qtype == RRSET and bailiwick: 248 res = dnstable_query_set_bailiwick(self._instance, bailiwick.encode('UTF-8')) 249 if res != dnstable_res_success: 250 raise DnstableException, 'dnstable_query_set_bailiwick() failed: %s' % dnstable_query_get_error(self._instance) 251 252 if timeout: 253 timeout=float(timeout) 254 if timeout < 0: 255 raise ValueError('timeout ({}) is not a positive number'.format(timeout)) 256 ts.tv_sec = math.trunc(timeout) 257 ts.tv_nsec = math.modf(timeout)[0] * 1e9 258 res = dnstable_query_set_timeout(self._instance, &ts) 259 if res != dnstable_res_success: 260 raise DnstableException, 'dnstable_query_set_timeout() failed: %s' % dnstable_query_get_error(self._instance) 261 262 if time_first_before is not None: 263 try: 264 tm = time_first_before 265 except OverflowError: 266 raise DnstableException, 'dnstable_query_set_timeout() failed: overflow error converting %s' % (time_first_before) 267 res = dnstable_query_set_filter_parameter(self._instance, 268 DNSTABLE_FILTER_PARAMETER_TIME_FIRST_BEFORE, &tm, 8) 269 if res != dnstable_res_success: 270 raise DnstableException, 'dnstable_query_set_filter_parameter(time_first_before) failed' 271 272 if time_first_after is not None: 273 try: 274 tm = time_first_after 275 except OverflowError: 276 raise DnstableException, 'dnstable_query_set_timeout() failed: overflow error converting %s' % (time_first_after) 277 res = dnstable_query_set_filter_parameter(self._instance, 278 DNSTABLE_FILTER_PARAMETER_TIME_FIRST_AFTER, &tm, 8) 279 if res != dnstable_res_success: 280 raise DnstableException, 'dnstable_query_set_filter_parameter(time_first_after) failed' 281 282 if time_last_before is not None: 283 try: 284 tm = time_last_before 285 except OverflowError: 286 raise DnstableException, 'dnstable_query_set_timeout() failed: overflow error converting %s' % (time_last_before) 287 res = dnstable_query_set_filter_parameter(self._instance, 288 DNSTABLE_FILTER_PARAMETER_TIME_LAST_BEFORE, &tm, 8) 289 if res != dnstable_res_success: 290 raise DnstableException, 'dnstable_query_set_filter_parameter(time_last_before) failed' 291 292 if time_last_after is not None: 293 try: 294 tm = time_last_after 295 except OverflowError: 296 raise DnstableException, 'dnstable_query_set_timeout() failed: overflow error converting %s' % (time_last_after) 297 res = dnstable_query_set_filter_parameter(self._instance, 298 DNSTABLE_FILTER_PARAMETER_TIME_LAST_AFTER, &tm, 8) 299 if res != dnstable_res_success: 300 raise DnstableException, 'dnstable_query_set_filter_parameter(time_last_after) failed' 301 302 def __dealloc__(self): 303 dnstable_query_destroy(&self._instance) 304 305 def __repr__(self): 306 if self.qtype == RDATA_IP: 307 s = 'ip ' 308 elif self.qtype == RDATA_NAME: 309 s = 'name ' 310 elif self.qtype == RDATA_RAW: 311 s = 'raw ' 312 else: 313 s = '' 314 315 s += self.data 316 if self.rrtype and self.rrtype.lower() != 'any': 317 s += '/' + self.rrtype.upper() 318 if self.bailiwick: 319 if (not self.rrtype) or (self.rrtype and self.rrtype.lower() == 'any'): 320 s += '/ANY' 321 s += '/' + self.bailiwick 322 return s 323 324cdef class reader(object): 325 cdef dnstable_reader *_instance 326 cdef object iszone 327 328 def __cinit__(self): 329 self._instance = NULL 330 331 def __dealloc__(self): 332 dnstable_reader_destroy(&self._instance) 333 334 def __init__(self, str fname, iszone=False): 335 import os 336 if not os.path.isfile(fname.encode('UTF-8')): 337 raise DnstableException, 'cannot open file %s' % fname 338 self._instance = dnstable_reader_init_setfile(fname.encode('UTF-8')) 339 self.iszone = iszone 340 341 def reload(self): 342 dnstable_reader_reload_setfile(self._instance) 343 344 def query(self, query q): 345 it = iteritems(self.iszone) 346 347 it._instance = dnstable_reader_query(self._instance, q._instance) 348 it.q = q 349 return it 350