1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
|
# DNS messages (RFC 1035) in plain Python: the reference 140 counts aup's lines against. A message's bytes read into its
# header and records, names in their wire form uncompressed (the owner's and those inside NS, CNAME, PTR, MX, SOA and
# SRV data), and written back with the owners compressed. Each pointer points before the last, so they end.
NAMED = {2: 0, 5: 0, 12: 0, 15: 2, 33: 6} # NS, CNAME, PTR, MX, SRV: the bytes before the name
def name(d, i):
"""The name at i uncompressed, and where it ends in place."""
out, j, end, below = b'', i, None, i
while d[j]:
if d[j] >= 0xc0:
to = (d[j] & 0x3f) << 8 | d[j + 1]
if to >= below:
raise ValueError('a pointer forward')
end, below, j = end or j + 2, to, to
elif d[j] < 0x40:
out += d[j:j + d[j] + 1]
j += d[j] + 1
else:
raise ValueError('a label type')
if len(out) >= 255:
raise ValueError('a name too long')
return out + b'\0', end or j + 1
def data(d, i, size, type):
"""A record's data, the names in it uncompressed."""
if type in NAMED:
n, end = name(d, i + NAMED[type])
if end != i + size:
raise ValueError('data length')
return d[i:i + NAMED[type]] + n
if type == 6:
m, after = name(d, i)
r, end = name(d, after)
if end + 20 != i + size:
raise ValueError('data length')
return m + r + d[end:end + 20]
return d[i:i + size]
def read(d):
"""The id, the flags and the records (section, name, type, class, ttl, data) of a message."""
if len(d) < 12:
raise ValueError('a header is 12 bytes')
word = lambda i: d[i] << 8 | d[i + 1]
records, at = [], 12
for section in range(4):
for _ in range(word(4 + 2 * section)):
owner, at = name(d, at)
type, klass = word(at), word(at + 2)
at += 4
if section == 0:
records.append((0, owner, type, klass, 0, b''))
continue
ttl, size = word(at) << 16 | word(at + 2), word(at + 4)
if at + 6 + size > len(d):
raise ValueError('data length')
records.append((section, owner, type, klass, ttl, data(d, at + 6, size, type)))
at += 6 + size
if at != len(d):
raise ValueError('bytes after the records')
return word(0), word(2), records
def write(id, flags, records):
"""The message's bytes, each owner from its first label whose rest was written before, as a pointer to it."""
out = bytearray(id.to_bytes(2, 'big') + flags.to_bytes(2, 'big'))
for section in range(4):
out += sum(r[0] == section for r in records).to_bytes(2, 'big')
seen = {}
for section, owner, type, klass, ttl, rdata in records:
k = 0
while owner[k] and owner[k:] not in seen:
if len(out) < 0x4000:
seen[owner[k:]] = len(out)
out += owner[k:k + owner[k] + 1]
k += owner[k] + 1
out += (0xc000 | seen[owner[k:]]).to_bytes(2, 'big') if owner[k] else b'\0'
out += type.to_bytes(2, 'big') + klass.to_bytes(2, 'big')
if section:
out += ttl.to_bytes(4, 'big') + len(rdata).to_bytes(2, 'big') + rdata
return bytes(out)
if __name__ == '__main__':
import sys, time
with open(sys.argv[1], 'rb') as f: # dns_check.py's messages: each a u16 big-endian length, then its bytes
given = f.read()
messages, at = [], 0
while at < len(given):
messages.append(given[at + 2:at + 2 + int.from_bytes(given[at:at + 2], 'big')])
at += 2 + len(messages[-1])
size = sum(map(len, messages))
runs = []
for _ in range(7):
start = time.perf_counter()
for m in messages:
write(*read(m))
runs.append((time.perf_counter() - start) * 1e9 / size)
print('{"messages":%d,"bytes":%d,"python_ns_per_byte":%.1f}' % (len(messages), size, sorted(runs)[3]))
|