Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/aiodns/__init__.py: 29%

1 

2import asyncio 

3import functools 

4import pycares 

5import socket 

6 

7from typing import ( 

8 Any, 

9 List, 

10 Optional, 

11 Set 

12) 

13 

14from . import error 

15 

16 

17__version__ = '3.0.0' 

18 

19__all__ = ('DNSResolver', 'error') 

20 

21 

22READ = 1 

23WRITE = 2 

24 

25query_type_map = {'A' : pycares.QUERY_TYPE_A, 

26 'AAAA' : pycares.QUERY_TYPE_AAAA, 

27 'ANY' : pycares.QUERY_TYPE_ANY, 

28 'CAA' : pycares.QUERY_TYPE_CAA, 

29 'CNAME' : pycares.QUERY_TYPE_CNAME, 

30 'MX' : pycares.QUERY_TYPE_MX, 

31 'NAPTR' : pycares.QUERY_TYPE_NAPTR, 

32 'NS' : pycares.QUERY_TYPE_NS, 

33 'PTR' : pycares.QUERY_TYPE_PTR, 

34 'SOA' : pycares.QUERY_TYPE_SOA, 

35 'SRV' : pycares.QUERY_TYPE_SRV, 

36 'TXT' : pycares.QUERY_TYPE_TXT 

37 } 

38 

39query_class_map = {'IN' : pycares.QUERY_CLASS_IN, 

40 'CHAOS' : pycares.QUERY_CLASS_CHAOS, 

41 'HS' : pycares.QUERY_CLASS_HS, 

42 'NONE' : pycares.QUERY_CLASS_NONE, 

43 'ANY' : pycares.QUERY_CLASS_ANY 

44 } 

45 

46class DNSResolver: 

47 def __init__(self, nameservers: Optional[List[str]] = None, 

48 loop: Optional[asyncio.AbstractEventLoop] = None, 

49 **kwargs: Any) -> None: 

50 self.loop = loop or asyncio.get_event_loop() 

51 assert self.loop is not None 

52 kwargs.pop('sock_state_cb', None) 

53 self._channel = pycares.Channel(sock_state_cb=self._sock_state_cb, **kwargs) 

54 if nameservers: 

55 self.nameservers = nameservers 

56 self._read_fds = set() # type: Set[int] 

57 self._write_fds = set() # type: Set[int] 

58 self._timer = None # type: Optional[asyncio.TimerHandle] 

59 

60 @property 

61 def nameservers(self) -> pycares.Channel: 

62 return self._channel.servers 

63 

64 @nameservers.setter 

65 def nameservers(self, value: List[str]) -> None: 

66 self._channel.servers = value 

67 

68 @staticmethod 

69 def _callback(fut: asyncio.Future, result: Any, errorno: int) -> None: 

70 if fut.cancelled(): 

71 return 

72 if errorno is not None: 

73 fut.set_exception(error.DNSError(errorno, pycares.errno.strerror(errorno))) 

74 else: 

75 fut.set_result(result) 

76 

77 def query(self, host: str, qtype: str, qclass: str=None) -> asyncio.Future: 

78 try: 

79 qtype = query_type_map[qtype] 

80 except KeyError: 

81 raise ValueError('invalid query type: {}'.format(qtype)) 

82 if qclass is not None: 

83 try: 

84 qclass = query_class_map[qclass] 

85 except KeyError: 

86 raise ValueError('invalid query class: {}'.format(qclass)) 

87 

88 fut = asyncio.Future(loop=self.loop) # type: asyncio.Future 

89 cb = functools.partial(self._callback, fut) 

90 self._channel.query(host, qtype, cb, query_class=qclass) 

91 return fut 

92 

93 def gethostbyname(self, host: str, family: socket.AddressFamily) -> asyncio.Future: 

94 fut = asyncio.Future(loop=self.loop) # type: asyncio.Future 

95 cb = functools.partial(self._callback, fut) 

96 self._channel.gethostbyname(host, family, cb) 

97 return fut 

98 

99 def gethostbyaddr(self, name: str) -> asyncio.Future: 

100 fut = asyncio.Future(loop=self.loop) # type: asyncio.Future 

101 cb = functools.partial(self._callback, fut) 

102 self._channel.gethostbyaddr(name, cb) 

103 return fut 

104 

105 def cancel(self) -> None: 

106 self._channel.cancel() 

107 

108 def _sock_state_cb(self, fd: int, readable: bool, writable: bool) -> None: 

109 if readable or writable: 

110 if readable: 

111 self.loop.add_reader(fd, self._handle_event, fd, READ) 

112 self._read_fds.add(fd) 

113 if writable: 

114 self.loop.add_writer(fd, self._handle_event, fd, WRITE) 

115 self._write_fds.add(fd) 

116 if self._timer is None: 

117 self._timer = self.loop.call_later(1.0, self._timer_cb) 

118 else: 

119 # socket is now closed 

120 if fd in self._read_fds: 

121 self._read_fds.discard(fd) 

122 self.loop.remove_reader(fd) 

123 

124 if fd in self._write_fds: 

125 self._write_fds.discard(fd) 

126 self.loop.remove_writer(fd) 

127 

128 if not self._read_fds and not self._write_fds and self._timer is not None: 

129 self._timer.cancel() 

130 self._timer = None 

131 

132 def _handle_event(self, fd: int, event: Any) -> None: 

133 read_fd = pycares.ARES_SOCKET_BAD 

134 write_fd = pycares.ARES_SOCKET_BAD 

135 if event == READ: 

136 read_fd = fd 

137 elif event == WRITE: 

138 write_fd = fd 

139 self._channel.process_fd(read_fd, write_fd) 

140 

141 def _timer_cb(self) -> None: 

142 if self._read_fds or self._write_fds: 

143 self._channel.process_fd(pycares.ARES_SOCKET_BAD, pycares.ARES_SOCKET_BAD) 

144 self._timer = self.loop.call_later(1.0, self._timer_cb) 

145 else: 

146 self._timer = None