From: Ishai <ishai@iThinkPad.hitronhub.home>
Date: Wed, 5 Jan 2022 20:39:48 +0000 (-0800)
Subject: Add resolver edns options
X-Git-Tag: v2.3.0rc1~157^2~1
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=998cd46f33178fd4c87cb080bcd856e746179060;p=thirdparty%2Fdnspython.git

Add resolver edns options
---

diff --git a/dns/resolver.py b/dns/resolver.py
index 166f8492..4c3a8dfc 100644
--- a/dns/resolver.py
+++ b/dns/resolver.py
@@ -607,7 +607,8 @@ class _Resolution:
                 request.use_tsig(self.resolver.keyring, self.resolver.keyname,
                                  algorithm=self.resolver.keyalgorithm)
             request.use_edns(self.resolver.edns, self.resolver.ednsflags,
-                             self.resolver.payload)
+                             self.resolver.payload,
+                             options=self.resolver.ednsoptions)
             if self.resolver.flags is not None:
                 request.flags = self.resolver.flags
 
@@ -776,6 +777,7 @@ class BaseResolver:
         self.keyalgorithm = dns.tsig.default_algorithm
         self.edns = -1
         self.ednsflags = 0
+        self.ednsoptions = None
         self.payload = 0
         self.cache = None
         self.flags = None
@@ -931,7 +933,7 @@ class BaseResolver:
         self.keyalgorithm = algorithm
 
     def use_edns(self, edns=0, ednsflags=0,
-                 payload=dns.message.DEFAULT_EDNS_PAYLOAD):
+                 payload=dns.message.DEFAULT_EDNS_PAYLOAD, options=None):
         """Configure EDNS behavior.
 
         *edns*, an ``int``, is the EDNS level to use.  Specifying
@@ -944,6 +946,9 @@ class BaseResolver:
         *payload*, an ``int``, is the EDNS sender's payload field, which is the
         maximum size of UDP datagram the sender can handle.  I.e. how big
         a response to this message can be.
+
+        *options*, a list of ``dns.edns.Option`` objects or ``None``, the EDNS
+        options.
         """
 
         if edns is None or edns is False:
@@ -953,6 +958,7 @@ class BaseResolver:
         self.edns = edns
         self.ednsflags = ednsflags
         self.payload = payload
+        self.ednsoptions = options
 
     def set_flags(self, flags):
         """Overrides the default flags with your own.
diff --git a/examples/edns_resolver.py b/examples/edns_resolver.py
new file mode 100644
index 00000000..fe5cc0f9
--- /dev/null
+++ b/examples/edns_resolver.py
@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+
+import dns.edns
+import dns.message
+import dns.query
+import dns.resolver
+
+n = '.'
+t = dns.rdatatype.SOA
+l = 'google.com'  # Address of l.root-servers.net, '199.7.83.42'
+i = 'ns1.isc.org'  # Address of ns1.isc.org, for COOKIEs, '149.20.1.73'
+
+o_list = []
+
+# A query without options
+o_list.append((l, dict()))
+
+# The same query, but with empty options list
+o_list.append((l, dict(options=[])))
+
+# Use use_edns() to specify EDNS0 options, such as buffer size
+o_list.append((l, dict(payload=2000)))
+
+# With an NSID option, but with use_edns() to specify the options
+edns_kwargs = dict(edns=0, options=[
+    dns.edns.GenericOption(dns.edns.OptionType.NSID, b'')])
+o_list.append((l, edns_kwargs))
+
+# With a COOKIE
+o_list.append((i, dict(options=[
+    dns.edns.GenericOption(dns.edns.OptionType.COOKIE, b'0xfe11ac99bebe3322')])))
+
+# With an ECS option using cloudflare dns address
+o_list.append((l, dict(options=[dns.edns.ECSOption('1.1.1.1', 24)])))
+
+# With an ECS option using the current machine address
+import urllib.request
+
+external_ip = urllib.request.urlopen('https://ident.me').read().decode('utf8')
+
+o_list.append((l, dict(options=[dns.edns.ECSOption(external_ip, 24)])))
+
+aresolver = dns.resolver.Resolver()
+
+for (addr, edns_kwargs) in o_list:
+    if edns_kwargs:
+        aresolver.use_edns(**edns_kwargs)
+    aresolver.nameservers = ['8.8.8.8']
+    print(list(aresolver.resolve(addr, 'A')))
diff --git a/tests/test_resolver.py b/tests/test_resolver.py
index 16d9065e..4a6bbb92 100644
--- a/tests/test_resolver.py
+++ b/tests/test_resolver.py
@@ -21,8 +21,8 @@ import sys
 import socket
 import time
 import unittest
-
 import pytest
+from unittest.mock import patch
 
 import dns.e164
 import dns.message
@@ -647,6 +647,14 @@ class LiveResolverTests(unittest.TestCase):
         dnsgoogle = dns.name.from_text('dns.google.')
         self.assertEqual(answer[0].target, dnsgoogle)
 
+    @patch.object(dns.message.Message, 'use_edns')
+    def testResolveEdnsOptions(self, message_use_edns_mock):
+        resolver = dns.resolver.Resolver()
+        options = [dns.edns.ECSOption('1.1.1.1')]
+        resolver.use_edns(True, options=options)
+        resolver.resolve('dns.google.', 'A')
+        assert message_use_edns_mock.call_args.kwargs == {'options': options}
+
     def testResolveNodataException(self):
         def bad():
             dns.resolver.resolve('dnspython.org.', 'SRV')