aiomisc 17.5.19__tar.gz → 17.5.20__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. {aiomisc-17.5.19 → aiomisc-17.5.20}/PKG-INFO +3 -1
  2. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/process_pool.py +1 -3
  3. aiomisc-17.5.20/aiomisc/service/dns/__init__.py +14 -0
  4. aiomisc-17.5.20/aiomisc/service/dns/records.py +400 -0
  5. aiomisc-17.5.20/aiomisc/service/dns/service.py +131 -0
  6. aiomisc-17.5.20/aiomisc/service/dns/store.py +51 -0
  7. aiomisc-17.5.20/aiomisc/service/dns/tree.py +59 -0
  8. aiomisc-17.5.20/aiomisc/service/dns/zone.py +44 -0
  9. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/grpc_server.py +17 -1
  10. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/thread_pool.py +1 -3
  11. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/version.py +2 -2
  12. {aiomisc-17.5.19 → aiomisc-17.5.20}/pyproject.toml +17 -7
  13. {aiomisc-17.5.19 → aiomisc-17.5.20}/COPYING +0 -0
  14. {aiomisc-17.5.19 → aiomisc-17.5.20}/README.rst +0 -0
  15. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/__init__.py +0 -0
  16. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/_context_vars.py +0 -0
  17. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/aggregate.py +0 -0
  18. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/backoff.py +0 -0
  19. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/circuit_breaker.py +0 -0
  20. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/compat.py +0 -0
  21. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/context.py +0 -0
  22. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/counters.py +0 -0
  23. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/cron.py +0 -0
  24. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/entrypoint.py +0 -0
  25. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/io.py +0 -0
  26. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/iterator_wrapper.py +0 -0
  27. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/log.py +0 -0
  28. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/periodic.py +0 -0
  29. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/plugins/__init__.py +0 -0
  30. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/plugins/__main__.py +0 -0
  31. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/pool.py +0 -0
  32. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/py.typed +0 -0
  33. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/recurring.py +0 -0
  34. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/__init__.py +0 -0
  35. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/aiohttp.py +0 -0
  36. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/asgi.py +0 -0
  37. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/base.py +0 -0
  38. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/carbon.py +0 -0
  39. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/cron.py +0 -0
  40. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/periodic.py +0 -0
  41. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/process.py +0 -0
  42. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/profiler.py +0 -0
  43. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/raven.py +0 -0
  44. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/sdwatchdog.py +0 -0
  45. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/tcp.py +0 -0
  46. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/tls.py +0 -0
  47. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/tracer.py +0 -0
  48. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/udp.py +0 -0
  49. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/service/uvicorn.py +0 -0
  50. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/signal.py +0 -0
  51. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/timeout.py +0 -0
  52. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/utils.py +0 -0
  53. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc/worker_pool.py +0 -0
  54. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_log/__init__.py +0 -0
  55. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_log/enum.py +0 -0
  56. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_log/formatter/__init__.py +0 -0
  57. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_log/formatter/color.py +0 -0
  58. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_log/formatter/journald.py +0 -0
  59. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_log/formatter/json.py +0 -0
  60. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_log/formatter/rich.py +0 -0
  61. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_log/py.typed +0 -0
  62. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_worker/__init__.py +0 -0
  63. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_worker/__main__.py +0 -0
  64. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_worker/forking.py +0 -0
  65. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_worker/process.py +0 -0
  66. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_worker/process_inner.py +0 -0
  67. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_worker/protocol.py +0 -0
  68. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_worker/py.typed +0 -0
  69. {aiomisc-17.5.19 → aiomisc-17.5.20}/aiomisc_worker/worker.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: aiomisc
3
- Version: 17.5.19
3
+ Version: 17.5.20
4
4
  Summary: aiomisc - miscellaneous utils for asyncio
5
5
  Home-page: https://github.com/aiokitchen/aiomisc
6
6
  License: MIT
@@ -44,6 +44,7 @@ Provides-Extra: aiohttp
44
44
  Provides-Extra: asgi
45
45
  Provides-Extra: carbon
46
46
  Provides-Extra: cron
47
+ Provides-Extra: dns
47
48
  Provides-Extra: grpc
48
49
  Provides-Extra: raven
49
50
  Provides-Extra: rich
@@ -55,6 +56,7 @@ Requires-Dist: aiohttp-asgi (>=0.5.2,<0.6.0) ; extra == "asgi"
55
56
  Requires-Dist: asgiref (>=3.7,<4.0) ; extra == "uvicorn"
56
57
  Requires-Dist: colorlog (>=6.0,<7.0)
57
58
  Requires-Dist: croniter (==2.0) ; extra == "cron"
59
+ Requires-Dist: dnslib (>=0.9,<0.10) ; extra == "dns"
58
60
  Requires-Dist: grpcio (>=1.56,<2.0) ; extra == "grpc"
59
61
  Requires-Dist: grpcio-reflection (>=1.56,<2.0) ; extra == "grpc"
60
62
  Requires-Dist: grpcio-tools (>=1.56,<2.0) ; extra == "grpc"
@@ -39,9 +39,7 @@ class ProcessPoolExecutor(ProcessPoolExecutorBase, EventLoopMixin):
39
39
  self._statistic.sum_time += loop.time() - start_time
40
40
 
41
41
  def submit(self, *args: Any, **kwargs: Any) -> Future:
42
- """
43
- Submit blocking function to the pool
44
- """
42
+ """Submit blocking function to the pool"""
45
43
  loop = asyncio.get_running_loop()
46
44
  start_time = loop.time()
47
45
  future = super().submit(*args, **kwargs)
@@ -0,0 +1,14 @@
1
+ from . import records
2
+ from .service import DNSServer, TCPDNSServer, UDPDNSServer
3
+ from .store import DNSStore
4
+ from .zone import DNSZone
5
+
6
+
7
+ __all__ = (
8
+ "DNSServer",
9
+ "DNSStore",
10
+ "DNSZone",
11
+ "TCPDNSServer",
12
+ "UDPDNSServer",
13
+ "records",
14
+ )
@@ -0,0 +1,400 @@
1
+ import enum
2
+ from abc import ABC, abstractmethod
3
+ from dataclasses import dataclass, field
4
+ from typing import Any, Hashable, Iterable, List
5
+
6
+ import dnslib # type: ignore[import-untyped]
7
+
8
+
9
+ class RecordType(enum.IntEnum):
10
+ A = 1
11
+ A6 = 38
12
+ AAAA = 28
13
+ AFSDB = 18
14
+ ANY = 255
15
+ APL = 42
16
+ AXFR = 252
17
+ CAA = 257
18
+ CDNSKEY = 60
19
+ CDS = 59
20
+ CERT = 37
21
+ CNAME = 5
22
+ CSYNC = 62
23
+ DHCID = 49
24
+ DLV = 32769
25
+ DNAME = 39
26
+ DNSKEY = 48
27
+ DS = 43
28
+ EUI48 = 108
29
+ EUI64 = 109
30
+ HINFO = 13
31
+ HIP = 53
32
+ HIP_AC = 55
33
+ HTTPS = 65
34
+ IPSECKEY = 45
35
+ IXFR = 251
36
+ KEY = 25
37
+ KX = 36
38
+ LOC = 29
39
+ MX = 15
40
+ NAPTR = 35
41
+ NS = 2
42
+ NSEC = 47
43
+ NSEC3 = 50
44
+ NSEC3PARAM = 51
45
+ NULL = 10
46
+ OPENPGPKEY = 61
47
+ OPT = 41
48
+ PTR = 12
49
+ RP = 17
50
+ RRSIG = 46
51
+ SIG = 24
52
+ SOA = 6
53
+ SPF = 99
54
+ SRV = 33
55
+ SSHFP = 44
56
+ SVCB = 64
57
+ TA = 32768
58
+ TKEY = 249
59
+ TLSA = 52
60
+ TSIG = 250
61
+ TXT = 16
62
+ URI = 256
63
+ ZONEMD = 63
64
+
65
+
66
+ class DNSClass(enum.IntEnum):
67
+ IN = 1
68
+ CS = 2
69
+ CH = 3
70
+ Hesiod = 4
71
+ NONE = 254
72
+ ASTERISK = 255
73
+
74
+
75
+ class RD(dnslib.RD, Hashable, ABC):
76
+ def __hash__(self) -> int:
77
+ return hash(self.data)
78
+
79
+ @classmethod
80
+ @abstractmethod
81
+ def create(cls, *args: Any, **kwargs: Any) -> "DNSRecord":
82
+ raise NotImplementedError
83
+
84
+
85
+ @dataclass(frozen=True)
86
+ class DNSRecord:
87
+ name: str
88
+ type: RecordType
89
+ data: RD
90
+ cls: DNSClass = field(default=DNSClass.IN)
91
+ ttl: int = field(default=3600)
92
+
93
+ def rr(self, query_type: int) -> dnslib.RR:
94
+ return dnslib.RR(
95
+ rname=self.name,
96
+ rtype=query_type,
97
+ rclass=self.cls,
98
+ ttl=self.ttl,
99
+ rdata=self.data,
100
+ )
101
+
102
+
103
+ class A(dnslib.A, RD):
104
+ @classmethod
105
+ def create(cls, name: str, ip: str, ttl: int = 3600) -> DNSRecord:
106
+ return DNSRecord(
107
+ name=name,
108
+ type=RecordType.A,
109
+ data=cls(ip),
110
+ ttl=ttl,
111
+ )
112
+
113
+
114
+ class AAAA(dnslib.AAAA, RD):
115
+ @classmethod
116
+ def create(cls, name: str, ipv6: str, ttl: int = 3600) -> DNSRecord:
117
+ return DNSRecord(
118
+ name=name,
119
+ type=RecordType.AAAA,
120
+ data=cls(ipv6),
121
+ ttl=ttl,
122
+ )
123
+
124
+
125
+ class RDLabel(RD):
126
+ label: bytes
127
+ __type__: RecordType
128
+
129
+ @classmethod
130
+ def create(cls, name: str, label: str, ttl: int = 3600) -> DNSRecord:
131
+ return DNSRecord(
132
+ name=name,
133
+ type=cls.__type__,
134
+ data=cls(label),
135
+ ttl=ttl,
136
+ )
137
+
138
+ def __hash__(self) -> int:
139
+ return hash(self.label)
140
+
141
+
142
+ class CNAME(dnslib.CNAME, RDLabel):
143
+ __type__ = RecordType.CNAME
144
+
145
+
146
+ class NS(dnslib.NS, RDLabel):
147
+ __type__ = RecordType.NS
148
+
149
+
150
+ class PTR(dnslib.PTR, RDLabel):
151
+ __type__ = RecordType.PTR
152
+
153
+
154
+ class DNAME(dnslib.DNAME, RDLabel):
155
+ __type__ = RecordType.DNAME
156
+
157
+
158
+ class MX(dnslib.MX, RD):
159
+ @classmethod
160
+ def create(
161
+ cls, name: str, exchange: str, preference: int, ttl: int = 3600,
162
+ ) -> DNSRecord:
163
+ return DNSRecord(
164
+ name=name,
165
+ type=RecordType.MX,
166
+ data=cls(exchange, preference),
167
+ ttl=ttl,
168
+ )
169
+
170
+
171
+ class TXT(dnslib.TXT, RD):
172
+ @classmethod
173
+ def create(cls, name: str, text: str, ttl: int = 3600) -> DNSRecord:
174
+ return DNSRecord(
175
+ name=name,
176
+ type=RecordType.TXT,
177
+ data=cls(text),
178
+ ttl=ttl,
179
+ )
180
+
181
+
182
+ class SOA(dnslib.SOA, RD):
183
+ @classmethod
184
+ def create(
185
+ cls, name: str, mname: str, rname: str, serial: int,
186
+ refresh: int, retry: int, expire: int, minimum: int,
187
+ ttl: int = 3600,
188
+ ) -> DNSRecord:
189
+ return DNSRecord(
190
+ name=name,
191
+ type=RecordType.SOA,
192
+ data=cls(
193
+ mname, rname,
194
+ (serial, refresh, retry, expire, minimum),
195
+ ),
196
+ ttl=ttl,
197
+ )
198
+
199
+
200
+ class SRV(dnslib.SRV, RD):
201
+ @classmethod
202
+ def create(
203
+ cls, name: str, priority: int, weight: int, port: int,
204
+ target: str, ttl: int = 3600,
205
+ ) -> DNSRecord:
206
+ return DNSRecord(
207
+ name=name,
208
+ type=RecordType.SRV,
209
+ data=cls(priority, weight, port, target),
210
+ ttl=ttl,
211
+ )
212
+
213
+
214
+ class CAA(dnslib.CAA, RD):
215
+ @classmethod
216
+ def create(
217
+ cls, name: str, flags: int, tag: str, value: str,
218
+ ttl: int = 3600,
219
+ ) -> DNSRecord:
220
+ return DNSRecord(
221
+ name=name,
222
+ type=RecordType.CAA,
223
+ data=cls(flags, tag, value),
224
+ ttl=ttl,
225
+ )
226
+
227
+
228
+ class NAPTR(dnslib.NAPTR, RD):
229
+ @classmethod
230
+ def create(
231
+ cls, name: str, order: int, preference: int, flags: str,
232
+ service: str, regexp: str, replacement: str, ttl: int = 3600,
233
+ ) -> DNSRecord:
234
+ return DNSRecord(
235
+ name=name,
236
+ type=RecordType.NAPTR,
237
+ data=cls(order, preference, flags, service, regexp, replacement),
238
+ ttl=ttl,
239
+ )
240
+
241
+
242
+ class DS(dnslib.DS, RD):
243
+ @classmethod
244
+ def create(
245
+ cls, name: str, key_tag: int, algorithm: int, digest_type: int,
246
+ digest: str, ttl: int = 3600,
247
+ ) -> DNSRecord:
248
+ return DNSRecord(
249
+ name=name,
250
+ type=RecordType.DS,
251
+ data=cls(key_tag, algorithm, digest_type, digest),
252
+ ttl=ttl,
253
+ )
254
+
255
+
256
+ class DNSKEY(dnslib.DNSKEY, RD):
257
+ @classmethod
258
+ def create(
259
+ cls, name: str, flags: int, protocol: int, algorithm: int,
260
+ key: str, ttl: int = 3600,
261
+ ) -> DNSRecord:
262
+ return DNSRecord(
263
+ name=name,
264
+ type=RecordType.DNSKEY,
265
+ data=cls(flags, protocol, algorithm, key),
266
+ ttl=ttl,
267
+ )
268
+
269
+
270
+ class RRSIG(dnslib.RRSIG, RD):
271
+ @classmethod
272
+ def create(
273
+ cls, name: str, type_covered: int, algorithm: int, labels: int,
274
+ original_ttl: int, expiration: int, inception: int, key_tag: int,
275
+ signer: str, signature: str, ttl: int = 3600,
276
+ ) -> DNSRecord:
277
+ return DNSRecord(
278
+ name=name,
279
+ type=RecordType.RRSIG,
280
+ data=cls(
281
+ type_covered, algorithm, labels, original_ttl, expiration,
282
+ inception, key_tag, signer, signature,
283
+ ),
284
+ ttl=ttl,
285
+ )
286
+
287
+
288
+ class NSEC(dnslib.NSEC, RD):
289
+ @classmethod
290
+ def create(
291
+ cls, name: str, next_domain: str,
292
+ rrtypes: Iterable[int], ttl: int = 3600,
293
+ ) -> DNSRecord:
294
+ return DNSRecord(
295
+ name=name,
296
+ type=RecordType.NSEC,
297
+ data=cls(next_domain, list(rrtypes)),
298
+ ttl=ttl,
299
+ )
300
+
301
+
302
+ class HTTPS(dnslib.HTTPS, RD):
303
+ @classmethod
304
+ def create(
305
+ cls, name: str, priority: int, target: str,
306
+ params: List[str], ttl: int = 3600,
307
+ ) -> DNSRecord:
308
+ return DNSRecord(
309
+ name=name,
310
+ type=RecordType.HTTPS,
311
+ data=cls(priority, target, params),
312
+ ttl=ttl,
313
+ )
314
+
315
+
316
+ class LOC(dnslib.LOC, RD):
317
+ @classmethod
318
+ def create(
319
+ cls, name: str, latitude: float, longitude: float,
320
+ altitude: float, size: float, h_precision: float,
321
+ v_precision: float, ttl: int = 3600,
322
+ ) -> DNSRecord:
323
+ return DNSRecord(
324
+ name=name,
325
+ type=RecordType.LOC,
326
+ data=cls(
327
+ latitude, longitude, altitude, size, h_precision, v_precision,
328
+ ),
329
+ ttl=ttl,
330
+ )
331
+
332
+
333
+ class RP(dnslib.RP, RD):
334
+ @classmethod
335
+ def create(
336
+ cls, name: str, mbox: str, txt: str, ttl: int = 3600,
337
+ ) -> DNSRecord:
338
+ return DNSRecord(
339
+ name=name,
340
+ type=RecordType.RP,
341
+ data=cls(mbox, txt),
342
+ ttl=ttl,
343
+ )
344
+
345
+
346
+ class TLSA(dnslib.TLSA, RD):
347
+ @classmethod
348
+ def create(
349
+ cls, name: str, usage: int, selector: int, mtype: int, cert: str,
350
+ ttl: int = 3600,
351
+ ) -> DNSRecord:
352
+ return DNSRecord(
353
+ name=name,
354
+ type=RecordType.TLSA,
355
+ data=cls(usage, selector, mtype, cert),
356
+ ttl=ttl,
357
+ )
358
+
359
+
360
+ class SSHFP(dnslib.SSHFP, RD):
361
+ @classmethod
362
+ def create(
363
+ cls, name: str, algorithm: int, fptype: int,
364
+ fingerprint: str, ttl: int = 3600,
365
+ ) -> DNSRecord:
366
+ return DNSRecord(
367
+ name=name,
368
+ type=RecordType.SSHFP,
369
+ data=cls(algorithm, fptype, fingerprint),
370
+ ttl=ttl,
371
+ )
372
+
373
+
374
+ __all__ = (
375
+ "AAAA",
376
+ "CAA",
377
+ "CNAME",
378
+ "DNAME",
379
+ "DNSKEY",
380
+ "DS",
381
+ "HTTPS",
382
+ "LOC",
383
+ "MX",
384
+ "NAPTR",
385
+ "NS",
386
+ "NSEC",
387
+ "PTR",
388
+ "RD",
389
+ "RP",
390
+ "RRSIG",
391
+ "SOA",
392
+ "SRV",
393
+ "SSHFP",
394
+ "TLSA",
395
+ "TXT",
396
+ "A",
397
+ "DNSRecord",
398
+ "DNSClass",
399
+ "RecordType",
400
+ )
@@ -0,0 +1,131 @@
1
+ import asyncio
2
+ import logging
3
+ from asyncio import StreamReader, StreamWriter
4
+ from struct import Struct
5
+ from typing import Optional, Tuple
6
+
7
+ import dnslib # type: ignore[import-untyped]
8
+
9
+ from aiomisc.service import TCPServer, UDPServer
10
+
11
+ from .records import RecordType
12
+ from .store import DNSStore
13
+
14
+
15
+ log = logging.getLogger(__name__)
16
+
17
+
18
+ class DNSServer:
19
+ store: DNSStore
20
+
21
+ __proto__: str
22
+
23
+ @staticmethod
24
+ def get_edns_max_size(record: dnslib.DNSRecord) -> Optional[int]:
25
+ for rr in record.ar:
26
+ if rr.rtype != dnslib.QTYPE.OPT:
27
+ continue
28
+ return rr.edns_len
29
+ return None
30
+
31
+ def handle_request(
32
+ self, addr: tuple, data: bytes, should_truncate: bool = False,
33
+ ) -> Tuple[bytes, tuple]:
34
+ record = dnslib.DNSRecord.parse(data)
35
+ question: dnslib.DNSQuestion = record.get_q()
36
+ reply = record.reply()
37
+ query_name = str(question.get_qname())
38
+ query_type = question.qtype
39
+
40
+ try:
41
+ log.debug(
42
+ "Processing %s request from %r:\n%s\n",
43
+ self.__proto__, addr, record,
44
+ )
45
+ edns_max_size = self.get_edns_max_size(record)
46
+ if edns_max_size is not None:
47
+ reply.add_ar(dnslib.EDNS0(udp_len=edns_max_size))
48
+
49
+ records = self.store.query(query_name, RecordType(query_type))
50
+ for rec in records:
51
+ reply.add_answer(rec.rr(query_type))
52
+
53
+ if not records:
54
+ reply.header.rcode = dnslib.RCODE.NXDOMAIN
55
+
56
+ log.debug(
57
+ "Sending %s answer to %r:\n%s\n",
58
+ self.__proto__, addr, reply,
59
+ )
60
+
61
+ reply_body = reply.pack()
62
+
63
+ if len(reply_body) > 512:
64
+ reply.header.tc = 1
65
+ reply_body = reply.pack()
66
+
67
+ if should_truncate:
68
+ reply_body = reply_body[:edns_max_size]
69
+ except Exception as e:
70
+ log.exception(
71
+ "Failed to process %s request from %r: %s",
72
+ self.__proto__, addr, e,
73
+ )
74
+ reply = record.reply()
75
+ reply.header.set_rcode(dnslib.RCODE.SERVFAIL)
76
+ reply_body = reply.pack()
77
+
78
+ return reply_body, addr
79
+
80
+
81
+ class UDPDNSServer(DNSServer, UDPServer):
82
+ MAX_SIZE = 512
83
+
84
+ store: DNSStore
85
+
86
+ __required__ = ("store",)
87
+ __proto__ = "UDP"
88
+
89
+ async def handle_datagram(self, data: bytes, addr: tuple) -> None:
90
+ self.sendto(
91
+ *self.handle_request(
92
+ data=data, addr=addr, should_truncate=True,
93
+ ),
94
+ )
95
+
96
+
97
+ TCP_HEADER_STRUCT: Struct = Struct("!H")
98
+
99
+
100
+ class TCPDNSServer(DNSServer, TCPServer):
101
+ store: DNSStore
102
+
103
+ __required__ = ("store",)
104
+
105
+ __proto__ = "TCP"
106
+
107
+ async def handle_client(
108
+ self, reader: StreamReader, writer: StreamWriter,
109
+ ) -> None:
110
+ addr = writer.get_extra_info("peername")
111
+ try:
112
+ while True:
113
+ data = await reader.readexactly(TCP_HEADER_STRUCT.size)
114
+ if not data:
115
+ break
116
+ length = TCP_HEADER_STRUCT.unpack(data)[0]
117
+ packet = await reader.read(length)
118
+
119
+ if not data:
120
+ break
121
+
122
+ reply_body, _ = self.handle_request(data=packet, addr=addr)
123
+
124
+ writer.write(TCP_HEADER_STRUCT.pack(len(reply_body)))
125
+ writer.write(reply_body)
126
+ await writer.drain()
127
+ except asyncio.IncompleteReadError:
128
+ pass
129
+ finally:
130
+ writer.close()
131
+ await writer.wait_closed()
@@ -0,0 +1,51 @@
1
+ from typing import Optional, Sequence, Tuple
2
+
3
+ from .records import DNSRecord, RecordType
4
+ from .tree import RadixTree
5
+ from .zone import DNSZone
6
+
7
+
8
+ class DNSStore:
9
+ zones: RadixTree[DNSZone]
10
+
11
+ __slots__ = ("zones",)
12
+
13
+ def __init__(self) -> None:
14
+ self.zones = RadixTree()
15
+
16
+ def add_zone(self, zone: DNSZone) -> None:
17
+ zone_tuple = self.get_reverse_tuple(zone.name)
18
+ if self.zones.search(zone_tuple):
19
+ raise ValueError(f"Zone {zone.name} already exists.")
20
+ self.zones.insert(zone_tuple, zone)
21
+
22
+ def remove_zone(self, zone_name: str) -> None:
23
+ zone_tuple = self.get_reverse_tuple(zone_name)
24
+ if not self.zones.search(zone_tuple):
25
+ raise ValueError(
26
+ f"Zone {zone_name} does not exist.",
27
+ )
28
+ # Clear zone from RadixTree
29
+ self.zones.insert(zone_tuple, None)
30
+
31
+ def get_zone(self, zone_name: str) -> Optional[DNSZone]:
32
+ zone_tuple = self.get_reverse_tuple(zone_name)
33
+ return self.zones.search(zone_tuple)
34
+
35
+ def query(self, name: str, record_type: RecordType) -> Sequence[DNSRecord]:
36
+ if not name.endswith("."):
37
+ name += "."
38
+ zone_tuple = self.get_zone_for_name(name)
39
+ if not zone_tuple:
40
+ return ()
41
+ zone = self.zones.search(zone_tuple)
42
+ return zone.get_records(name, record_type) if zone is not None else ()
43
+
44
+ def get_zone_for_name(self, name: str) -> Optional[Tuple[str, ...]]:
45
+ labels = self.get_reverse_tuple(name)
46
+ result = self.zones.find_prefix(labels)
47
+ return result[0] if result else None
48
+
49
+ @staticmethod
50
+ def get_reverse_tuple(zone_name: str) -> Tuple[str, ...]:
51
+ return tuple(zone_name.strip(".").split("."))[::-1]
@@ -0,0 +1,59 @@
1
+ from typing import Any, Dict, Generic, Hashable, List, Optional, Tuple, TypeVar
2
+
3
+
4
+ K = Tuple[str, ...]
5
+ T = TypeVar("T", bound=Any)
6
+
7
+
8
+ class RadixNode(Generic[T]):
9
+ __slots__ = ("children", "value")
10
+
11
+ def __init__(self) -> None:
12
+ self.children: Dict[Hashable, RadixNode[T]] = {}
13
+ self.value: Optional[T] = None
14
+
15
+
16
+ class RadixTree(Generic[T]):
17
+ root: RadixNode[T]
18
+
19
+ __slots__ = ("root",)
20
+
21
+ def __init__(self) -> None:
22
+ self.root = RadixNode()
23
+
24
+ def insert(self, key: K, value: Optional[T]) -> None:
25
+ node = self.root
26
+ for part in key:
27
+ if part not in node.children:
28
+ node.children[part] = RadixNode()
29
+ node = node.children[part]
30
+ node.value = value
31
+
32
+ def search(self, key: K) -> Optional[T]:
33
+ node = self.root
34
+ for part in key:
35
+ if part not in node.children:
36
+ return None
37
+ node = node.children[part]
38
+ return node.value
39
+
40
+ def find_prefix(self, key: K) -> Optional[Tuple[K, T]]:
41
+ node = self.root
42
+ longest_prefix: List[str] = []
43
+ value = None
44
+ part: Hashable
45
+ if node.value is not None:
46
+ value = node.value
47
+ for part in key:
48
+ if part in node.children:
49
+ node = node.children[part]
50
+ longest_prefix.append(part)
51
+ if node.value is not None:
52
+ value = node.value
53
+ else:
54
+ break
55
+
56
+ if value is None:
57
+ return None
58
+
59
+ return tuple(longest_prefix), value
@@ -0,0 +1,44 @@
1
+ from collections import defaultdict
2
+ from typing import DefaultDict, Sequence, Set, Tuple
3
+
4
+ from .records import DNSRecord, RecordType
5
+
6
+
7
+ class DNSZone:
8
+ records: DefaultDict[Tuple[str, RecordType], Set[DNSRecord]]
9
+ name: str
10
+
11
+ __slots__ = ("name", "records")
12
+
13
+ def __init__(self, name: str):
14
+ if not name.endswith("."):
15
+ name += "."
16
+ self.name = name
17
+ self.records = defaultdict(set)
18
+
19
+ def add_record(self, record: DNSRecord) -> None:
20
+ if not self._is_valid_record(record):
21
+ raise ValueError(
22
+ f"Record {record.name} does not belong to zone {self.name}",
23
+ )
24
+ key = (record.name, record.type)
25
+ self.records[key].add(record)
26
+
27
+ def remove_record(self, record: DNSRecord) -> None:
28
+ key = (record.name, record.type)
29
+ if key in self.records:
30
+ self.records[key].discard(record)
31
+ if self.records[key]:
32
+ return
33
+ del self.records[key]
34
+
35
+ def get_records(
36
+ self, name: str, record_type: RecordType,
37
+ ) -> Sequence[DNSRecord]:
38
+ if not name.endswith("."):
39
+ name += "."
40
+ key = (name, record_type)
41
+ return tuple(self.records.get(key, ()))
42
+
43
+ def _is_valid_record(self, record: DNSRecord) -> bool:
44
+ return record.name.endswith(self.name)
@@ -2,9 +2,12 @@ import asyncio
2
2
  import logging
3
3
  import re
4
4
  import sys
5
+ from collections import defaultdict
5
6
  from concurrent.futures import Executor
6
7
  from types import MappingProxyType
7
- from typing import Any, Optional, Sequence, Set, Tuple
8
+ from typing import (
9
+ Any, DefaultDict, Dict, Mapping, Optional, Sequence, Set, Tuple,
10
+ )
8
11
 
9
12
  from .base import Service
10
13
 
@@ -36,6 +39,7 @@ class GRPCService(Service):
36
39
  _server_args: MappingProxyType
37
40
  _insecure_ports: Set[Tuple[str, PortFuture]]
38
41
  _secure_ports: Set[Tuple[str, grpc.ServerCredentials, PortFuture]]
42
+ _registered_services: DefaultDict[str, Dict[str, grpc.RpcMethodHandler]]
39
43
 
40
44
  def __init__(
41
45
  self, *,
@@ -60,6 +64,7 @@ class GRPCService(Service):
60
64
  self._insecure_ports = set()
61
65
  self._secure_ports = set()
62
66
  self._reflection = reflection
67
+ self._registered_services = defaultdict(dict)
63
68
  super().__init__(**kwds)
64
69
 
65
70
  @classmethod
@@ -90,6 +95,12 @@ class GRPCService(Service):
90
95
  service_names.append(reflection.SERVICE_NAME)
91
96
  reflection.enable_server_reflection(service_names, self._server)
92
97
 
98
+ for name, handlers in self._registered_services.items():
99
+ # noinspection PyUnresolvedReferences
100
+ self._server.add_registered_method_handlers( # type: ignore
101
+ name, handlers,
102
+ )
103
+
93
104
  self._server.add_generic_rpc_handlers(tuple(self._services))
94
105
  await self._server.start()
95
106
 
@@ -102,6 +113,11 @@ class GRPCService(Service):
102
113
  for service in generic_rpc_handlers:
103
114
  self._services.add(service)
104
115
 
116
+ def add_registered_method_handlers(
117
+ self, name: str, handlers: Mapping[str, grpc.RpcMethodHandler],
118
+ ) -> None:
119
+ self._registered_services[name].update(handlers)
120
+
105
121
  def add_insecure_port(self, address: str) -> PortFuture:
106
122
  future: PortFuture = asyncio.Future()
107
123
  self._insecure_ports.add((address, future))
@@ -224,9 +224,7 @@ class ThreadPoolExecutor(ThreadPoolExecutorBase):
224
224
  def submit( # type: ignore
225
225
  self, fn: F, *args: Any, **kwargs: Any,
226
226
  ) -> asyncio.Future:
227
- """
228
- Submit blocking function to the pool
229
- """
227
+ """Submit blocking function to the pool"""
230
228
  if fn is None or not callable(fn):
231
229
  raise ValueError("First argument must be callable")
232
230
 
@@ -2,5 +2,5 @@
2
2
  # BY: poem-plugins "git" plugin
3
3
  # NEVER EDIT THIS FILE MANUALLY
4
4
 
5
- version_info = (17, 5, 19)
6
- __version__ = "17.5.19"
5
+ version_info = (17, 5, 20)
6
+ __version__ = "17.5.20"
@@ -1,7 +1,7 @@
1
1
  [tool.poetry]
2
2
  name = "aiomisc"
3
3
  # This is a dummy version which will be rewritten with poem-plugins
4
- version = "17.5.19"
4
+ version = "17.5.20"
5
5
  description = "aiomisc - miscellaneous utils for asyncio"
6
6
  authors = ["Dmitry Orlov <me@mosquito.su>"]
7
7
  readme = "README.rst"
@@ -72,22 +72,28 @@ typing_extensions = [{ version = '*', python = "< 3.10" }]
72
72
  uvloop = { version = ">=0.19, <1", optional = true }
73
73
  uvicorn = { version = "^0.27", optional = true }
74
74
  asgiref = { version = "^3.7", optional = true }
75
+ dnslib = { version = "^0.9", optional = true }
75
76
 
76
77
  [tool.poetry.group.dev.dependencies]
78
+ aiocarbon = "^0.15.3"
77
79
  aiohttp = "^3.9"
78
80
  aiohttp-asgi = "~0.5.2"
79
81
  aiomisc-pytest = "^1.0.8"
82
+ asgiref = "^3.7"
80
83
  async-timeout = "^4.0.2"
81
84
  autodoc = "^0.5.0"
82
85
  autoflake = "1.4"
86
+ certifi = "^2024.6.2"
83
87
  collective-checkdocs = "^0.2"
84
88
  coveralls = "^3.3.1"
85
89
  croniter = "^2.0"
90
+ dnslib = "^0.9"
86
91
  fastapi = "^0.110"
87
92
  furo = "^2022"
88
- grpcio = "^1.56"
89
- grpcio-tools = "^1.56"
90
- grpcio-reflection = "^1.56"
93
+ grpc-stubs = "^1.53.0.2"
94
+ grpcio = "^1.64"
95
+ grpcio-reflection = "^1.64"
96
+ grpcio-tools = "^1.64"
91
97
  mypy = "^1.9"
92
98
  pre-commit = "^2.20.0"
93
99
  pylama = "^8.4.1"
@@ -105,9 +111,7 @@ sphinx-intl = "^2.0"
105
111
  timeout-decorator = "^0.5.0"
106
112
  types-croniter = "^1.3"
107
113
  types-setuptools = "^65.6.0.1"
108
- grpc-stubs = "^1.53.0.2"
109
114
  uvicorn = "^0.27"
110
- asgiref = "^3.7"
111
115
 
112
116
  [tool.poetry.group.uvloop.dependencies]
113
117
  uvloop = "^0.19.0"
@@ -117,11 +121,12 @@ aiohttp = ["aiohttp"]
117
121
  asgi = ["aiohttp-asgi"]
118
122
  carbon = ["aiocarbon"]
119
123
  cron = ["croniter"]
124
+ dns = ["dnslib"]
120
125
  grpc = ["grpcio", "grpcio-tools", "grpcio-reflection"]
121
126
  raven = ["aiohttp", "raven"]
122
127
  rich = ["rich"]
123
- uvloop = ["uvloop"]
124
128
  uvicorn = ["uvicorn", "asgiref"]
129
+ uvloop = ["uvloop"]
125
130
 
126
131
  [tool.poetry.plugins.aiomisc]
127
132
  systemd_watchdog = "aiomisc.service.sdwatchdog"
@@ -156,6 +161,11 @@ files = [
156
161
  "tests",
157
162
  ]
158
163
 
164
+ [[tool.mypy.overrides]]
165
+ module = "aiomisc.service.dns.records"
166
+ check_untyped_defs = true
167
+ disallow_subclassing_any = false
168
+
159
169
  [[tool.mypy.overrides]]
160
170
  module = ["tests.*"]
161
171
  check_untyped_defs = true
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes