From b826ae6b910ba71b95d42cbda4f8128e8681e60d Mon Sep 17 00:00:00 2001
From: nom3ad <19239479+nom3ad@users.noreply.github.com>
Date: Wed, 1 May 2024 20:08:15 +0530
Subject: [PATCH] windows: support automatic nameserver detection for --dns
 option

---
 sshuttle/helpers.py          | 43 +++++++++++++++++++++++++++++-------
 sshuttle/server.py           |  4 ++--
 tests/client/test_helpers.py |  4 ++--
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/sshuttle/helpers.py b/sshuttle/helpers.py
index 85ff057..2c98940 100644
--- a/sshuttle/helpers.py
+++ b/sshuttle/helpers.py
@@ -3,7 +3,9 @@ import socket
 import errno
 import os
 import threading
+import subprocess
 import traceback
+import re
 
 if sys.platform != "win32":
     import fcntl
@@ -114,18 +116,43 @@ def resolvconf_nameservers(systemd_resolved):
     return nsservers
 
 
-def resolvconf_random_nameserver(systemd_resolved):
+def windows_nameservers():
+    out = subprocess.check_output(["powershell", "-NonInteractive", "-NoProfile", "-Command", "Get-DnsClientServerAddress"],
+                                  encoding="utf-8")
+    servers = set()
+    for line in out.splitlines():
+        if line.startswith("Loopback "):
+            continue
+        m = re.search(r'{.+}', line)
+        if not m:
+            continue
+        for s in m.group().strip('{}').split(','):
+            s = s.strip()
+            if s.startswith('fec0:0:0:ffff'):
+                continue
+            servers.add(s)
+    debug2("Found DNS servers: %s" % servers)
+    return [(socket.AF_INET6 if ':' in s else socket.AF_INET, s) for s in servers]
+
+
+def get_random_nameserver():
     """Return a random nameserver selected from servers produced by
-    resolvconf_nameservers(). See documentation for
-    resolvconf_nameservers() for a description of the parameter.
+    resolvconf_nameservers()/windows_nameservers()
     """
-    lines = resolvconf_nameservers(systemd_resolved)
-    if lines:
-        if len(lines) > 1:
+    if sys.platform == "win32":
+        if globals().get('_nameservers') is None:
+            ns_list = windows_nameservers()
+            globals()['_nameservers'] = ns_list
+        else:
+            ns_list = globals()['_nameservers']
+    else:
+        ns_list = resolvconf_nameservers(systemd_resolved=False)
+    if ns_list:
+        if len(ns_list) > 1:
             # don't import this unless we really need it
             import random
-            random.shuffle(lines)
-        return lines[0]
+            random.shuffle(ns_list)
+        return ns_list[0]
     else:
         return (socket.AF_INET, '127.0.0.1')
 
diff --git a/sshuttle/server.py b/sshuttle/server.py
index 62c156f..4350dfb 100644
--- a/sshuttle/server.py
+++ b/sshuttle/server.py
@@ -14,7 +14,7 @@ import sshuttle.hostwatch as hostwatch
 import subprocess as ssubprocess
 from sshuttle.ssnet import Handler, Proxy, Mux, MuxWrapper
 from sshuttle.helpers import b, log, debug1, debug2, debug3, Fatal, \
-    resolvconf_random_nameserver, which, get_env, SocketRWShim
+    get_random_nameserver, which, get_env, SocketRWShim
 
 
 def _ipmatch(ipstr):
@@ -199,7 +199,7 @@ class DnsProxy(Handler):
         self.tries += 1
 
         if self.to_nameserver is None:
-            _, peer = resolvconf_random_nameserver(False)
+            _, peer = get_random_nameserver()
             port = 53
         else:
             peer = self.to_ns_peer
diff --git a/tests/client/test_helpers.py b/tests/client/test_helpers.py
index ca1aba3..bfbb145 100644
--- a/tests/client/test_helpers.py
+++ b/tests/client/test_helpers.py
@@ -143,7 +143,7 @@ nameserver 2404:6800:4004:80c::4
 
 
 @patch('sshuttle.helpers.open', create=True)
-def test_resolvconf_random_nameserver(mock_open):
+def test_get_random_nameserver(mock_open):
     mock_open.return_value = io.StringIO(u"""
 # Generated by NetworkManager
 search pri
@@ -156,7 +156,7 @@ nameserver 2404:6800:4004:80c::2
 nameserver 2404:6800:4004:80c::3
 nameserver 2404:6800:4004:80c::4
 """)
-    ns = sshuttle.helpers.resolvconf_random_nameserver(False)
+    ns = sshuttle.helpers.get_random_nameserver()
     assert ns in [
         (AF_INET, u'192.168.1.1'), (AF_INET, u'192.168.2.1'),
         (AF_INET, u'192.168.3.1'), (AF_INET, u'192.168.4.1'),