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