Layer that handles resolving DNS queries.
| 54 | |
| 55 | |
| 56 | class DNSLayer(layer.Layer): |
| 57 | """ |
| 58 | Layer that handles resolving DNS queries. |
| 59 | """ |
| 60 | |
| 61 | flows: dict[int, dns.DNSFlow] |
| 62 | req_buf: bytearray |
| 63 | resp_buf: bytearray |
| 64 | |
| 65 | def __init__(self, context: Context): |
| 66 | super().__init__(context) |
| 67 | self.flows = {} |
| 68 | self.req_buf = bytearray() |
| 69 | self.resp_buf = bytearray() |
| 70 | |
| 71 | def handle_request( |
| 72 | self, flow: dns.DNSFlow, msg: dns.DNSMessage |
| 73 | ) -> layer.CommandGenerator[None]: |
| 74 | flow.request = msg # if already set, continue and query upstream again |
| 75 | yield DnsRequestHook(flow) |
| 76 | if flow.response: |
| 77 | yield from self.handle_response(flow, flow.response) |
| 78 | elif flow.error: |
| 79 | yield from self.handle_error(flow, flow.error.msg) |
| 80 | elif not self.context.server.address: |
| 81 | yield from self.handle_error( |
| 82 | flow, "No hook has set a response and there is no upstream server." |
| 83 | ) |
| 84 | else: |
| 85 | if not self.context.server.connected: |
| 86 | err = yield commands.OpenConnection(self.context.server) |
| 87 | if err: |
| 88 | yield from self.handle_error(flow, str(err)) |
| 89 | # cannot recover from this |
| 90 | return |
| 91 | packed = pack_message(flow.request, flow.server_conn.transport_protocol) |
| 92 | yield commands.SendData(self.context.server, packed) |
| 93 | |
| 94 | def handle_response( |
| 95 | self, flow: dns.DNSFlow, msg: dns.DNSMessage |
| 96 | ) -> layer.CommandGenerator[None]: |
| 97 | flow.response = msg |
| 98 | yield DnsResponseHook(flow) |
| 99 | if flow.response: |
| 100 | packed = pack_message(flow.response, flow.client_conn.transport_protocol) |
| 101 | yield commands.SendData(self.context.client, packed) |
| 102 | |
| 103 | def handle_error(self, flow: dns.DNSFlow, err: str) -> layer.CommandGenerator[None]: |
| 104 | flow.error = mflow.Error(err) |
| 105 | yield DnsErrorHook(flow) |
| 106 | servfail = flow.request.fail(response_codes.SERVFAIL) |
| 107 | yield commands.SendData( |
| 108 | self.context.client, |
| 109 | pack_message(servfail, flow.client_conn.transport_protocol), |
| 110 | ) |
| 111 | |
| 112 | def unpack_message(self, data: bytes, from_client: bool) -> List[dns.DNSMessage]: |
| 113 | msgs: List[dns.DNSMessage] = [] |
no outgoing calls
no test coverage detected
searching dependent graphs…