aiomisc 17.5.17__tar.gz → 17.5.20__tar.gz

{aiomisc-17.5.17 → aiomisc-17.5.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: aiomisc
-Version: 17.5.17
+Version: 17.5.20
 Summary: aiomisc - miscellaneous utils for asyncio
 Home-page: https://github.com/aiokitchen/aiomisc
 License: MIT
@@ -44,6 +44,7 @@ Provides-Extra: aiohttp
 Provides-Extra: asgi
 Provides-Extra: carbon
 Provides-Extra: cron
+Provides-Extra: dns
 Provides-Extra: grpc
 Provides-Extra: raven
 Provides-Extra: rich
@@ -55,6 +56,7 @@ Requires-Dist: aiohttp-asgi (>=0.5.2,<0.6.0) ; extra == "asgi"
 Requires-Dist: asgiref (>=3.7,<4.0) ; extra == "uvicorn"
 Requires-Dist: colorlog (>=6.0,<7.0)
 Requires-Dist: croniter (==2.0) ; extra == "cron"
+Requires-Dist: dnslib (>=0.9,<0.10) ; extra == "dns"
 Requires-Dist: grpcio (>=1.56,<2.0) ; extra == "grpc"
 Requires-Dist: grpcio-reflection (>=1.56,<2.0) ; extra == "grpc"
 Requires-Dist: grpcio-tools (>=1.56,<2.0) ; extra == "grpc"

{aiomisc-17.5.17 → aiomisc-17.5.20}/aiomisc/aggregate.py

@@ -7,8 +7,8 @@ from dataclasses import dataclass
 from inspect import Parameter
 from typing import (
     Any,
-    Awaitable,
     Callable,
+    Coroutine,
     Generic,
     Iterable,
     List,
@@ -246,7 +246,7 @@ class Aggregator(AggregatorAsync[V, R], Generic[V, R]):
 
 def aggregate(
     leeway_ms: float, max_count: Optional[int] = None
-) -> Callable[[AggregateFunc[V, R]], Callable[[V], Awaitable[R]]]:
+) -> Callable[[AggregateFunc[V, R]], Callable[[V], Coroutine[Any, Any, R]]]:
     """
     Parametric decorator that aggregates multiple
     (but no more than ``max_count`` defaulting to ``None``) single-argument
@@ -275,7 +275,9 @@ def aggregate(
 
     :return:
     """
-    def decorator(func: AggregateFunc[V, R]) -> Callable[[V], Awaitable[R]]:
+    def decorator(
+        func: AggregateFunc[V, R]
+    ) -> Callable[[V], Coroutine[Any, Any, R]]:
         aggregator = Aggregator(
             func, max_count=max_count, leeway_ms=leeway_ms,
         )
@@ -285,7 +287,10 @@ def aggregate(
 
 def aggregate_async(
     leeway_ms: float, max_count: Optional[int] = None,
-) -> Callable[[AggregateAsyncFunc[V, R]], Callable[[V], Awaitable[R]]]:
+) -> Callable[
+    [AggregateAsyncFunc[V, R]],
+    Callable[[V], Coroutine[Any, Any, R]]
+]:
     """
     Same as ``aggregate``, but with ``func`` arguments of type ``Arg``
     containing ``value`` and ``future`` attributes instead. In this setting
@@ -298,7 +303,7 @@ def aggregate_async(
     """
     def decorator(
         func: AggregateAsyncFunc[V, R]
-    ) -> Callable[[V], Awaitable[R]]:
+    ) -> Callable[[V], Coroutine[Any, Any, R]]:
         aggregator = AggregatorAsync(
             func, max_count=max_count, leeway_ms=leeway_ms,
         )

{aiomisc-17.5.17 → aiomisc-17.5.20}/aiomisc/backoff.py

@@ -2,7 +2,7 @@ import asyncio
 import sys
 from functools import wraps
 from typing import (
-    Any, Awaitable, Callable, Optional, Tuple, Type, TypeVar, Union,
+    Any, Callable, Coroutine, Optional, Tuple, Type, TypeVar, Union,
 )
 
 from .counters import Statistic
@@ -42,7 +42,10 @@ def asyncbackoff(
     giveup: Optional[Callable[[Exception], bool]] = None,
     statistic_name: Optional[str] = None,
     statistic_class: Type[BackoffStatistic] = BackoffStatistic,
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     """
     Patametric decorator that ensures that ``attempt_timeout`` and
     ``deadline`` time limits are met by decorated function.
@@ -85,7 +88,9 @@ def asyncbackoff(
     exceptions = tuple(exceptions) or ()
     exceptions += asyncio.TimeoutError,
 
-    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+    def decorator(
+        func: Callable[P, Coroutine[Any, Any, T]]
+    ) -> Callable[P, Coroutine[Any, Any, T]]:
         if attempt_timeout is not None:
             func = timeout(attempt_timeout)(func)
 
@@ -145,7 +150,10 @@ def asyncretry(
     pause: Number = 0,
     giveup: Optional[Callable[[Exception], bool]] = None,
     statistic_name: Optional[str] = None,
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     """
     Shortcut of ``asyncbackoff(None, None, 0, Exception)``.
 

{aiomisc-17.5.17 → aiomisc-17.5.20}/aiomisc/process_pool.py

@@ -39,9 +39,7 @@ class ProcessPoolExecutor(ProcessPoolExecutorBase, EventLoopMixin):
         self._statistic.sum_time += loop.time() - start_time
 
     def submit(self, *args: Any, **kwargs: Any) -> Future:
-        """
-        Submit blocking function to the pool
-        """
+        """Submit blocking function to the pool"""
         loop = asyncio.get_running_loop()
         start_time = loop.time()
         future = super().submit(*args, **kwargs)

aiomisc-17.5.20/aiomisc/service/dns/__init__.py

@@ -0,0 +1,14 @@
+from . import records
+from .service import DNSServer, TCPDNSServer, UDPDNSServer
+from .store import DNSStore
+from .zone import DNSZone
+
+
+__all__ = (
+    "DNSServer",
+    "DNSStore",
+    "DNSZone",
+    "TCPDNSServer",
+    "UDPDNSServer",
+    "records",
+)

aiomisc-17.5.20/aiomisc/service/dns/records.py

@@ -0,0 +1,400 @@
+import enum
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Hashable, Iterable, List
+
+import dnslib  # type: ignore[import-untyped]
+
+
+class RecordType(enum.IntEnum):
+    A = 1
+    A6 = 38
+    AAAA = 28
+    AFSDB = 18
+    ANY = 255
+    APL = 42
+    AXFR = 252
+    CAA = 257
+    CDNSKEY = 60
+    CDS = 59
+    CERT = 37
+    CNAME = 5
+    CSYNC = 62
+    DHCID = 49
+    DLV = 32769
+    DNAME = 39
+    DNSKEY = 48
+    DS = 43
+    EUI48 = 108
+    EUI64 = 109
+    HINFO = 13
+    HIP = 53
+    HIP_AC = 55
+    HTTPS = 65
+    IPSECKEY = 45
+    IXFR = 251
+    KEY = 25
+    KX = 36
+    LOC = 29
+    MX = 15
+    NAPTR = 35
+    NS = 2
+    NSEC = 47
+    NSEC3 = 50
+    NSEC3PARAM = 51
+    NULL = 10
+    OPENPGPKEY = 61
+    OPT = 41
+    PTR = 12
+    RP = 17
+    RRSIG = 46
+    SIG = 24
+    SOA = 6
+    SPF = 99
+    SRV = 33
+    SSHFP = 44
+    SVCB = 64
+    TA = 32768
+    TKEY = 249
+    TLSA = 52
+    TSIG = 250
+    TXT = 16
+    URI = 256
+    ZONEMD = 63
+
+
+class DNSClass(enum.IntEnum):
+    IN = 1
+    CS = 2
+    CH = 3
+    Hesiod = 4
+    NONE = 254
+    ASTERISK = 255
+
+
+class RD(dnslib.RD, Hashable, ABC):
+    def __hash__(self) -> int:
+        return hash(self.data)
+
+    @classmethod
+    @abstractmethod
+    def create(cls, *args: Any, **kwargs: Any) -> "DNSRecord":
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class DNSRecord:
+    name: str
+    type: RecordType
+    data: RD
+    cls: DNSClass = field(default=DNSClass.IN)
+    ttl: int = field(default=3600)
+
+    def rr(self, query_type: int) -> dnslib.RR:
+        return dnslib.RR(
+            rname=self.name,
+            rtype=query_type,
+            rclass=self.cls,
+            ttl=self.ttl,
+            rdata=self.data,
+        )
+
+
+class A(dnslib.A, RD):
+    @classmethod
+    def create(cls, name: str, ip: str, ttl: int = 3600) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.A,
+            data=cls(ip),
+            ttl=ttl,
+        )
+
+
+class AAAA(dnslib.AAAA, RD):
+    @classmethod
+    def create(cls, name: str, ipv6: str, ttl: int = 3600) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.AAAA,
+            data=cls(ipv6),
+            ttl=ttl,
+        )
+
+
+class RDLabel(RD):
+    label: bytes
+    __type__: RecordType
+
+    @classmethod
+    def create(cls, name: str, label: str, ttl: int = 3600) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=cls.__type__,
+            data=cls(label),
+            ttl=ttl,
+        )
+
+    def __hash__(self) -> int:
+        return hash(self.label)
+
+
+class CNAME(dnslib.CNAME, RDLabel):
+    __type__ = RecordType.CNAME
+
+
+class NS(dnslib.NS, RDLabel):
+    __type__ = RecordType.NS
+
+
+class PTR(dnslib.PTR, RDLabel):
+    __type__ = RecordType.PTR
+
+
+class DNAME(dnslib.DNAME, RDLabel):
+    __type__ = RecordType.DNAME
+
+
+class MX(dnslib.MX, RD):
+    @classmethod
+    def create(
+        cls, name: str, exchange: str, preference: int, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.MX,
+            data=cls(exchange, preference),
+            ttl=ttl,
+        )
+
+
+class TXT(dnslib.TXT, RD):
+    @classmethod
+    def create(cls, name: str, text: str, ttl: int = 3600) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.TXT,
+            data=cls(text),
+            ttl=ttl,
+        )
+
+
+class SOA(dnslib.SOA, RD):
+    @classmethod
+    def create(
+        cls, name: str, mname: str, rname: str, serial: int,
+        refresh: int, retry: int, expire: int, minimum: int,
+        ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.SOA,
+            data=cls(
+                mname, rname,
+                (serial, refresh, retry, expire, minimum),
+            ),
+            ttl=ttl,
+        )
+
+
+class SRV(dnslib.SRV, RD):
+    @classmethod
+    def create(
+        cls, name: str, priority: int, weight: int, port: int,
+        target: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.SRV,
+            data=cls(priority, weight, port, target),
+            ttl=ttl,
+        )
+
+
+class CAA(dnslib.CAA, RD):
+    @classmethod
+    def create(
+        cls, name: str, flags: int, tag: str, value: str,
+        ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.CAA,
+            data=cls(flags, tag, value),
+            ttl=ttl,
+        )
+
+
+class NAPTR(dnslib.NAPTR, RD):
+    @classmethod
+    def create(
+        cls, name: str, order: int, preference: int, flags: str,
+        service: str, regexp: str, replacement: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.NAPTR,
+            data=cls(order, preference, flags, service, regexp, replacement),
+            ttl=ttl,
+        )
+
+
+class DS(dnslib.DS, RD):
+    @classmethod
+    def create(
+        cls, name: str, key_tag: int, algorithm: int, digest_type: int,
+        digest: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.DS,
+            data=cls(key_tag, algorithm, digest_type, digest),
+            ttl=ttl,
+        )
+
+
+class DNSKEY(dnslib.DNSKEY, RD):
+    @classmethod
+    def create(
+        cls, name: str, flags: int, protocol: int, algorithm: int,
+        key: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.DNSKEY,
+            data=cls(flags, protocol, algorithm, key),
+            ttl=ttl,
+        )
+
+
+class RRSIG(dnslib.RRSIG, RD):
+    @classmethod
+    def create(
+        cls, name: str, type_covered: int, algorithm: int, labels: int,
+        original_ttl: int, expiration: int, inception: int, key_tag: int,
+        signer: str, signature: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.RRSIG,
+            data=cls(
+                type_covered, algorithm, labels, original_ttl, expiration,
+                inception, key_tag, signer, signature,
+            ),
+            ttl=ttl,
+        )
+
+
+class NSEC(dnslib.NSEC, RD):
+    @classmethod
+    def create(
+        cls, name: str, next_domain: str,
+        rrtypes: Iterable[int], ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.NSEC,
+            data=cls(next_domain, list(rrtypes)),
+            ttl=ttl,
+        )
+
+
+class HTTPS(dnslib.HTTPS, RD):
+    @classmethod
+    def create(
+        cls, name: str, priority: int, target: str,
+        params: List[str], ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.HTTPS,
+            data=cls(priority, target, params),
+            ttl=ttl,
+        )
+
+
+class LOC(dnslib.LOC, RD):
+    @classmethod
+    def create(
+        cls, name: str, latitude: float, longitude: float,
+        altitude: float, size: float, h_precision: float,
+        v_precision: float, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.LOC,
+            data=cls(
+                latitude, longitude, altitude, size, h_precision, v_precision,
+            ),
+            ttl=ttl,
+        )
+
+
+class RP(dnslib.RP, RD):
+    @classmethod
+    def create(
+        cls, name: str, mbox: str, txt: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.RP,
+            data=cls(mbox, txt),
+            ttl=ttl,
+        )
+
+
+class TLSA(dnslib.TLSA, RD):
+    @classmethod
+    def create(
+        cls, name: str, usage: int, selector: int, mtype: int, cert: str,
+        ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.TLSA,
+            data=cls(usage, selector, mtype, cert),
+            ttl=ttl,
+        )
+
+
+class SSHFP(dnslib.SSHFP, RD):
+    @classmethod
+    def create(
+        cls, name: str, algorithm: int, fptype: int,
+        fingerprint: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.SSHFP,
+            data=cls(algorithm, fptype, fingerprint),
+            ttl=ttl,
+        )
+
+
+__all__ = (
+    "AAAA",
+    "CAA",
+    "CNAME",
+    "DNAME",
+    "DNSKEY",
+    "DS",
+    "HTTPS",
+    "LOC",
+    "MX",
+    "NAPTR",
+    "NS",
+    "NSEC",
+    "PTR",
+    "RD",
+    "RP",
+    "RRSIG",
+    "SOA",
+    "SRV",
+    "SSHFP",
+    "TLSA",
+    "TXT",
+    "A",
+    "DNSRecord",
+    "DNSClass",
+    "RecordType",
+)

aiomisc-17.5.20/aiomisc/service/dns/service.py

@@ -0,0 +1,131 @@
+import asyncio
+import logging
+from asyncio import StreamReader, StreamWriter
+from struct import Struct
+from typing import Optional, Tuple
+
+import dnslib  # type: ignore[import-untyped]
+
+from aiomisc.service import TCPServer, UDPServer
+
+from .records import RecordType
+from .store import DNSStore
+
+
+log = logging.getLogger(__name__)
+
+
+class DNSServer:
+    store: DNSStore
+
+    __proto__: str
+
+    @staticmethod
+    def get_edns_max_size(record: dnslib.DNSRecord) -> Optional[int]:
+        for rr in record.ar:
+            if rr.rtype != dnslib.QTYPE.OPT:
+                continue
+            return rr.edns_len
+        return None
+
+    def handle_request(
+        self, addr: tuple, data: bytes, should_truncate: bool = False,
+    ) -> Tuple[bytes, tuple]:
+        record = dnslib.DNSRecord.parse(data)
+        question: dnslib.DNSQuestion = record.get_q()
+        reply = record.reply()
+        query_name = str(question.get_qname())
+        query_type = question.qtype
+
+        try:
+            log.debug(
+                "Processing %s request from %r:\n%s\n",
+                self.__proto__, addr, record,
+            )
+            edns_max_size = self.get_edns_max_size(record)
+            if edns_max_size is not None:
+                reply.add_ar(dnslib.EDNS0(udp_len=edns_max_size))
+
+            records = self.store.query(query_name, RecordType(query_type))
+            for rec in records:
+                reply.add_answer(rec.rr(query_type))
+
+            if not records:
+                reply.header.rcode = dnslib.RCODE.NXDOMAIN
+
+            log.debug(
+                "Sending %s answer to %r:\n%s\n",
+                self.__proto__, addr, reply,
+            )
+
+            reply_body = reply.pack()
+
+            if len(reply_body) > 512:
+                reply.header.tc = 1
+                reply_body = reply.pack()
+
+            if should_truncate:
+                reply_body = reply_body[:edns_max_size]
+        except Exception as e:
+            log.exception(
+                "Failed to process %s request from %r: %s",
+                self.__proto__, addr, e,
+            )
+            reply = record.reply()
+            reply.header.set_rcode(dnslib.RCODE.SERVFAIL)
+            reply_body = reply.pack()
+
+        return reply_body, addr
+
+
+class UDPDNSServer(DNSServer, UDPServer):
+    MAX_SIZE = 512
+
+    store: DNSStore
+
+    __required__ = ("store",)
+    __proto__ = "UDP"
+
+    async def handle_datagram(self, data: bytes, addr: tuple) -> None:
+        self.sendto(
+            *self.handle_request(
+                data=data, addr=addr, should_truncate=True,
+            ),
+        )
+
+
+TCP_HEADER_STRUCT: Struct = Struct("!H")
+
+
+class TCPDNSServer(DNSServer, TCPServer):
+    store: DNSStore
+
+    __required__ = ("store",)
+
+    __proto__ = "TCP"
+
+    async def handle_client(
+        self, reader: StreamReader, writer: StreamWriter,
+    ) -> None:
+        addr = writer.get_extra_info("peername")
+        try:
+            while True:
+                data = await reader.readexactly(TCP_HEADER_STRUCT.size)
+                if not data:
+                    break
+                length = TCP_HEADER_STRUCT.unpack(data)[0]
+                packet = await reader.read(length)
+
+                if not data:
+                    break
+
+                reply_body, _ = self.handle_request(data=packet, addr=addr)
+
+                writer.write(TCP_HEADER_STRUCT.pack(len(reply_body)))
+                writer.write(reply_body)
+                await writer.drain()
+        except asyncio.IncompleteReadError:
+            pass
+        finally:
+            writer.close()
+            await writer.wait_closed()

aiomisc-17.5.20/aiomisc/service/dns/store.py

@@ -0,0 +1,51 @@
+from typing import Optional, Sequence, Tuple
+
+from .records import DNSRecord, RecordType
+from .tree import RadixTree
+from .zone import DNSZone
+
+
+class DNSStore:
+    zones: RadixTree[DNSZone]
+
+    __slots__ = ("zones",)
+
+    def __init__(self) -> None:
+        self.zones = RadixTree()
+
+    def add_zone(self, zone: DNSZone) -> None:
+        zone_tuple = self.get_reverse_tuple(zone.name)
+        if self.zones.search(zone_tuple):
+            raise ValueError(f"Zone {zone.name} already exists.")
+        self.zones.insert(zone_tuple, zone)
+
+    def remove_zone(self, zone_name: str) -> None:
+        zone_tuple = self.get_reverse_tuple(zone_name)
+        if not self.zones.search(zone_tuple):
+            raise ValueError(
+                f"Zone {zone_name} does not exist.",
+            )
+        # Clear zone from RadixTree
+        self.zones.insert(zone_tuple, None)
+
+    def get_zone(self, zone_name: str) -> Optional[DNSZone]:
+        zone_tuple = self.get_reverse_tuple(zone_name)
+        return self.zones.search(zone_tuple)
+
+    def query(self, name: str, record_type: RecordType) -> Sequence[DNSRecord]:
+        if not name.endswith("."):
+            name += "."
+        zone_tuple = self.get_zone_for_name(name)
+        if not zone_tuple:
+            return ()
+        zone = self.zones.search(zone_tuple)
+        return zone.get_records(name, record_type) if zone is not None else ()
+
+    def get_zone_for_name(self, name: str) -> Optional[Tuple[str, ...]]:
+        labels = self.get_reverse_tuple(name)
+        result = self.zones.find_prefix(labels)
+        return result[0] if result else None
+
+    @staticmethod
+    def get_reverse_tuple(zone_name: str) -> Tuple[str, ...]:
+        return tuple(zone_name.strip(".").split("."))[::-1]

aiomisc-17.5.20/aiomisc/service/dns/tree.py

@@ -0,0 +1,59 @@
+from typing import Any, Dict, Generic, Hashable, List, Optional, Tuple, TypeVar
+
+
+K = Tuple[str, ...]
+T = TypeVar("T", bound=Any)
+
+
+class RadixNode(Generic[T]):
+    __slots__ = ("children", "value")
+
+    def __init__(self) -> None:
+        self.children: Dict[Hashable, RadixNode[T]] = {}
+        self.value: Optional[T] = None
+
+
+class RadixTree(Generic[T]):
+    root: RadixNode[T]
+
+    __slots__ = ("root",)
+
+    def __init__(self) -> None:
+        self.root = RadixNode()
+
+    def insert(self, key: K, value: Optional[T]) -> None:
+        node = self.root
+        for part in key:
+            if part not in node.children:
+                node.children[part] = RadixNode()
+            node = node.children[part]
+        node.value = value
+
+    def search(self, key: K) -> Optional[T]:
+        node = self.root
+        for part in key:
+            if part not in node.children:
+                return None
+            node = node.children[part]
+        return node.value
+
+    def find_prefix(self, key: K) -> Optional[Tuple[K, T]]:
+        node = self.root
+        longest_prefix: List[str] = []
+        value = None
+        part: Hashable
+        if node.value is not None:
+            value = node.value
+        for part in key:
+            if part in node.children:
+                node = node.children[part]
+                longest_prefix.append(part)
+                if node.value is not None:
+                    value = node.value
+            else:
+                break
+
+        if value is None:
+            return None
+
+        return tuple(longest_prefix), value

aiomisc-17.5.20/aiomisc/service/dns/zone.py

@@ -0,0 +1,44 @@
+from collections import defaultdict
+from typing import DefaultDict, Sequence, Set, Tuple
+
+from .records import DNSRecord, RecordType
+
+
+class DNSZone:
+    records: DefaultDict[Tuple[str, RecordType], Set[DNSRecord]]
+    name: str
+
+    __slots__ = ("name", "records")
+
+    def __init__(self, name: str):
+        if not name.endswith("."):
+            name += "."
+        self.name = name
+        self.records = defaultdict(set)
+
+    def add_record(self, record: DNSRecord) -> None:
+        if not self._is_valid_record(record):
+            raise ValueError(
+                f"Record {record.name} does not belong to zone {self.name}",
+            )
+        key = (record.name, record.type)
+        self.records[key].add(record)
+
+    def remove_record(self, record: DNSRecord) -> None:
+        key = (record.name, record.type)
+        if key in self.records:
+            self.records[key].discard(record)
+            if self.records[key]:
+                return
+            del self.records[key]
+
+    def get_records(
+        self, name: str, record_type: RecordType,
+    ) -> Sequence[DNSRecord]:
+        if not name.endswith("."):
+            name += "."
+        key = (name, record_type)
+        return tuple(self.records.get(key, ()))
+
+    def _is_valid_record(self, record: DNSRecord) -> bool:
+        return record.name.endswith(self.name)

{aiomisc-17.5.17 → aiomisc-17.5.20}/aiomisc/service/grpc_server.py

@@ -2,9 +2,12 @@ import asyncio
 import logging
 import re
 import sys
+from collections import defaultdict
 from concurrent.futures import Executor
 from types import MappingProxyType
-from typing import Any, Optional, Sequence, Set, Tuple
+from typing import (
+    Any, DefaultDict, Dict, Mapping, Optional, Sequence, Set, Tuple,
+)
 
 from .base import Service
 
@@ -36,6 +39,7 @@ class GRPCService(Service):
     _server_args: MappingProxyType
     _insecure_ports: Set[Tuple[str, PortFuture]]
     _secure_ports: Set[Tuple[str, grpc.ServerCredentials, PortFuture]]
+    _registered_services: DefaultDict[str, Dict[str, grpc.RpcMethodHandler]]
 
     def __init__(
         self, *,
@@ -60,6 +64,7 @@ class GRPCService(Service):
         self._insecure_ports = set()
         self._secure_ports = set()
         self._reflection = reflection
+        self._registered_services = defaultdict(dict)
         super().__init__(**kwds)
 
     @classmethod
@@ -90,6 +95,12 @@ class GRPCService(Service):
             service_names.append(reflection.SERVICE_NAME)
             reflection.enable_server_reflection(service_names, self._server)
 
+        for name, handlers in self._registered_services.items():
+            # noinspection PyUnresolvedReferences
+            self._server.add_registered_method_handlers(   # type: ignore
+                name, handlers,
+            )
+
         self._server.add_generic_rpc_handlers(tuple(self._services))
         await self._server.start()
 
@@ -102,6 +113,11 @@ class GRPCService(Service):
         for service in generic_rpc_handlers:
             self._services.add(service)
 
+    def add_registered_method_handlers(
+        self, name: str, handlers: Mapping[str, grpc.RpcMethodHandler],
+    ) -> None:
+        self._registered_services[name].update(handlers)
+
     def add_insecure_port(self, address: str) -> PortFuture:
         future: PortFuture = asyncio.Future()
         self._insecure_ports.add((address, future))

{aiomisc-17.5.17 → aiomisc-17.5.20}/aiomisc/thread_pool.py

@@ -224,9 +224,7 @@ class ThreadPoolExecutor(ThreadPoolExecutorBase):
     def submit(  # type: ignore
         self, fn: F, *args: Any, **kwargs: Any,
     ) -> asyncio.Future:
-        """
-        Submit blocking function to the pool
-        """
+        """Submit blocking function to the pool"""
         if fn is None or not callable(fn):
             raise ValueError("First argument must be callable")
 

{aiomisc-17.5.17 → aiomisc-17.5.20}/aiomisc/timeout.py

@@ -1,7 +1,7 @@
 import asyncio
 import sys
 from functools import wraps
-from typing import Awaitable, Callable, TypeVar, Union
+from typing import Any, Callable, Coroutine, TypeVar, Union
 
 
 if sys.version_info >= (3, 10):
@@ -17,10 +17,13 @@ Number = Union[int, float]
 
 def timeout(
     value: Number
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     def decorator(
-        func: Callable[P, Awaitable[T]],
-    ) -> Callable[P, Awaitable[T]]:
+        func: Callable[P, Coroutine[Any, Any, T]],
+    ) -> Callable[P, Coroutine[Any, Any, T]]:
         if not asyncio.iscoroutinefunction(func):
             raise TypeError("Function is not a coroutine function")
 

{aiomisc-17.5.17 → aiomisc-17.5.20}/aiomisc/version.py

@@ -2,5 +2,5 @@
 # BY: poem-plugins "git" plugin
 # NEVER EDIT THIS FILE MANUALLY
 
-version_info = (17, 5, 17)
-__version__ = "17.5.17"
+version_info = (17, 5, 20)
+__version__ = "17.5.20"

{aiomisc-17.5.17 → aiomisc-17.5.20}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "aiomisc"
 # This is a dummy version which will be rewritten with poem-plugins
-version = "17.5.17"
+version = "17.5.20"
 description = "aiomisc - miscellaneous utils for asyncio"
 authors = ["Dmitry Orlov <me@mosquito.su>"]
 readme = "README.rst"
@@ -72,22 +72,28 @@ typing_extensions = [{ version = '*', python = "< 3.10" }]
 uvloop = { version = ">=0.19, <1", optional = true }
 uvicorn = { version = "^0.27", optional = true }
 asgiref = { version = "^3.7", optional = true }
+dnslib = { version = "^0.9", optional = true }
 
 [tool.poetry.group.dev.dependencies]
+aiocarbon = "^0.15.3"
 aiohttp = "^3.9"
 aiohttp-asgi = "~0.5.2"
 aiomisc-pytest = "^1.0.8"
+asgiref = "^3.7"
 async-timeout = "^4.0.2"
 autodoc = "^0.5.0"
 autoflake = "1.4"
+certifi = "^2024.6.2"
 collective-checkdocs = "^0.2"
 coveralls = "^3.3.1"
 croniter = "^2.0"
+dnslib = "^0.9"
 fastapi = "^0.110"
 furo = "^2022"
-grpcio = "^1.56"
-grpcio-tools = "^1.56"
-grpcio-reflection = "^1.56"
+grpc-stubs = "^1.53.0.2"
+grpcio = "^1.64"
+grpcio-reflection = "^1.64"
+grpcio-tools = "^1.64"
 mypy = "^1.9"
 pre-commit = "^2.20.0"
 pylama = "^8.4.1"
@@ -105,9 +111,7 @@ sphinx-intl = "^2.0"
 timeout-decorator = "^0.5.0"
 types-croniter = "^1.3"
 types-setuptools = "^65.6.0.1"
-grpc-stubs = "^1.53.0.2"
 uvicorn = "^0.27"
-asgiref = "^3.7"
 
 [tool.poetry.group.uvloop.dependencies]
 uvloop = "^0.19.0"
@@ -117,11 +121,12 @@ aiohttp = ["aiohttp"]
 asgi = ["aiohttp-asgi"]
 carbon = ["aiocarbon"]
 cron = ["croniter"]
+dns = ["dnslib"]
 grpc = ["grpcio", "grpcio-tools", "grpcio-reflection"]
 raven = ["aiohttp", "raven"]
 rich = ["rich"]
-uvloop = ["uvloop"]
 uvicorn = ["uvicorn", "asgiref"]
+uvloop = ["uvloop"]
 
 [tool.poetry.plugins.aiomisc]
 systemd_watchdog = "aiomisc.service.sdwatchdog"
@@ -156,6 +161,11 @@ files = [
     "tests",
 ]
 
+[[tool.mypy.overrides]]
+module = "aiomisc.service.dns.records"
+check_untyped_defs = true
+disallow_subclassing_any = false
+
 [[tool.mypy.overrides]]
 module = ["tests.*"]
 check_untyped_defs = true