add mode spec for WireGuard mode

mitmproxy · mhils · Sep 18, 2022 · Sep 4, 2022 · Sep 6, 2022 · Sep 11, 2022
commit 151dbd69dfe66946f69ee24a55ca1f7f891896ea
diff --git a/mitmproxy/proxy/mode_specs.py b/mitmproxy/proxy/mode_specs.py
@@ -245,3 +245,59 @@ class DnsMode(ProxyMode):
 
     def __post_init__(self) -> None:
         _check_empty(self.data)
+
+
+class WireGuardMode(ProxyMode):
+    """Proxy Server based on WireGuard"""
+    description = "WireGuard server"
+    default_port = 51820
+    transport_protocol = UDP
+    custom_listen_port: int | None = None
+
+    wireguard_cfg_gen: bool | None = None
+    wireguard_cfg_path: str | None = None
+    wireguard_peer_num: int | None = None
+
+    # noinspection PyDataclass
+    def __post_init__(self) -> None:
+        # use default settings
+        if self.data in [""]:
+            return
+
+        settings = self.data.split(",")
+        for setting in settings:
+            # handle empty strings after trailing commas
+            if not setting:
+                break
+
+            # handle settings with and without default values
+            if "=" not in setting:
+                key, value = setting, None
+            else:
+                key, value = setting.split("=", 1)
+
+            if key not in {"generate", "load", "peers"}:
+                raise ValueError(f"invalid wireguard mode spec (unexpected {key!r} setting)")
+
+            if key == "generate":
+                # generate new configuration files
+                if self.wireguard_cfg_gen is False:
+                    raise ValueError("invalid wireguard mode spec (cannot set both 'load' and 'generate')")
+                self.wireguard_cfg_gen = True
+                self.wireguard_cfg_path = value
+
+            if key == "load":
+                # load existing configuration file
+                if self.wireguard_cfg_gen is True:
+                    raise ValueError("invalid wireguard mode spec (cannot set both 'load' and 'generate')")
+                self.wireguard_cfg_gen = False
+                self.wireguard_cfg_path = value
+
+            if key == "peers":
+                # generate configuration for specified number of peers
+                if self.wireguard_cfg_gen is False:
+                    raise ValueError(f"invalid wireguard mode spec (unexpected {setting!r} setting)")
+                try:
+                    self.wireguard_peer_num = int(value) if value is not None else None
+                except ValueError:
+                    raise ValueError(f"invalid wireguard mode spec (invalid peer number {value!r})")
diff --git a/test/mitmproxy/proxy/test_mode_specs.py b/test/mitmproxy/proxy/test_mode_specs.py
@@ -1,6 +1,6 @@
 import pytest
 
-from mitmproxy.proxy.mode_specs import ProxyMode, Socks5Mode
+from mitmproxy.proxy.mode_specs import ProxyMode, Socks5Mode, WireGuardMode
 
 
 def test_parse():
@@ -58,6 +58,8 @@ def test_parse_specific_modes():
     assert ProxyMode.parse("dns")
     assert ProxyMode.parse("reverse:dns://8.8.8.8")
     assert ProxyMode.parse("reverse:dtls://127.0.0.1:8004")
+    assert ProxyMode.parse("wireguard")
+    assert ProxyMode.parse("wireguard@51821").listen_port() == 51821
 
     with pytest.raises(ValueError, match="invalid port"):
         ProxyMode.parse("regular@invalid-port")
@@ -73,3 +75,25 @@ def test_parse_specific_modes():
 
     with pytest.raises(ValueError, match="Port specification missing."):
         ProxyMode.parse("reverse:dtls://127.0.0.1")
+
+
+def test_parse_wireguard_mode():
+    assert WireGuardMode.parse("wireguard:load,")
+    assert WireGuardMode.parse("wireguard:generate,peers=2").wireguard_peer_num == 2
+    assert WireGuardMode.parse("wireguard:load=~/.mitmproxy/wg.json").wireguard_cfg_path == "~/.mitmproxy/wg.json"
+
+    mode = WireGuardMode.parse("wireguard:generate,peers=2@51821")
+    assert mode.listen_port() == 51821
+    assert mode.wireguard_cfg_gen is True
+    assert mode.wireguard_peer_num == 2
+
+    with pytest.raises(ValueError, match="cannot set both 'load' and 'generate'"):
+        WireGuardMode.parse("wireguard:load,generate")
+    with pytest.raises(ValueError, match="cannot set both 'load' and 'generate'"):
+        WireGuardMode.parse("wireguard:generate,load")
+    with pytest.raises(ValueError, match=f"unexpected 'peers=2' setting"):
+        WireGuardMode.parse("wireguard:load,peers=2")
+    with pytest.raises(ValueError, match="unexpected 'foobar' setting"):
+        WireGuardMode.parse("wireguard:foobar")
+    with pytest.raises(ValueError, match=f"invalid peer number 'foo'"):
+        WireGuardMode.parse("wireguard:generate,peers=foo")