multiple instance and multiple server support for ezwg

modules/ezwg.nix

@@ -2,102 +2,148 @@
 with lib;
 let
   cfg = config.services.ezwg;
-in
-{
-  options.services.ezwg = {
-    enable = mkEnableOption "Enable simple Wireguard connection";
 
+  peerNameReplacement = lib.replaceChars [ "/" "-" " " "+" "=" ] [
+    "-"
+    "\\x2d"
+    "\\x20"
+    "\\x2b"
+    "\\x3d"
+  ];
+
+  ranges = serverIPs:
+    let
+      generateRangesScript =
+        builtins.toFile "exclusionary-wildcard-ranges-generator.py" ''
+          import ipaddress
+          serverNetworks = [${map (ip: "ip_network('${ip}/32')") serverIPs}]
+          ranges = [ipaddress.ip_network('0.0.0.0/0')]
+          for serverNetwork in serverNetworks:
+            ranges = map(lambda r: list(r.address_exclude(serverNetwork)), ranges)
+          print(':'.join(ranges))
+        '';
+      rangesOutput = pkgs.runCommandNoCC "exclusionary-wildcard-ranges" { } ''
+        ${pkgs.python3}/bin/python3 ${generateRangesScript} > $out
+      '';
+    in
+    lib.splitString ":" (builtins.readFile "${rangesOutput}");
+
+  subnet = vlanIP: vlanSize:
+    let
+      generateSubnetScript =
+        builtins.toFile "subnet-without-host-bits-generator.py" ''
+          import ipaddress
+          n1 = ipaddress.ip_network('${vlanIP}/${toString vlanSize}', False)
+          print(n1, end="")
+        '';
+      subnetOutput = pkgs.runCommandNoCC "subnet-without-host-bits" { } ''
+        ${pkgs.python3}/bin/python3 ${generateSubnetScript} > $out
+      '';
+    in
+    builtins.readFile "${subnetOutput}";
+
+  serverOpts.options = {
+    ip = mkOption {
+      type = types.str;
+      description = "The IP of the wg server";
+    };
+    port = mkOption {
+      type = types.int;
+      default = 51820;
+      description = "The port of the wg server";
+    };
+    publicKey = mkOption {
+      type = types.str;
+      description = "The public key of the wg server";
+    };
+  };
+
+  instanceOpts.options = {
+    servers = mkOption {
+      description = "Configuration of servers to connect to";
+      default = { };
+      type = with types; listOf (submodule serverOpts);
+    };
+    autoStart = mkOption {
+      type = types.bool;
+      default = true;
+      description = "Start this wireguard tunnel by default";
+    };
     proxy = mkOption {
       type = types.bool;
       default = true;
       description = "Route all your traffic through this connection";
     };
-
+    vlanSize = mkOption {
-    lanSize = mkOption {
       type = types.int;
       default = 24;
-      description = "Size of your VLAN (only relevant if proxy is false)";
+      description = "Size of your VLAN";
     };
-
-    serverIP = mkOption {
-      type = types.str;
-      description = "The IP of the wg server";
-    };
-
-    excludeIP = mkOption {
-      type = types.str;
-      default = cfg.serverIP;
-      description = "The IP to _not_ route through the proxy, you normally want this to be the same as `serverIP` when not tunneling";
-    };
-
-    serverPort = mkOption {
-      type = types.int;
-      default = 51820;
-      description = "The port of the wg server";
-    };
-
-    serverKey = mkOption {
-      type = types.str;
-      description = "The public key of the wg server";
-    };
-
     privateKeyFile = mkOption {
       type = types.str;
       description = "Private wg key";
     };
-
     vlanIP = mkOption {
       type = types.str;
       description = "The IP to use on the wg VLAN";
     };
   };
-
+in
-  config = mkIf cfg.enable {
+{
-    networking.firewall.checkReversePath = false;
+  options.services.ezwg = {
-    networking.wireguard.interfaces.wg0 =
+    enable = mkEnableOption "Enable simple Wireguard connection";
-      let
+    instances = mkOption {
-        generateRangesScript =
+      description = "Configuration of instances of Wireguard";
-          builtins.toFile "exclusionary-wildcard-ranges-generator.py" ''
+      default = { };
-            import ipaddress
+      type = with types; attrsOf (submodule instanceOpts);
-            n1 = ipaddress.ip_network('0.0.0.0/0')
+    };
-            n2 = ipaddress.ip_network('${cfg.excludeIP}/32')
-            print(':'.join(list(map(lambda x: str(x), list(n1.address_exclude(n2))))), end="")
-          '';
-
-        rangesOutput =
-          pkgs.runCommandNoCC "exclusionary-wildcard-ranges"
-            { } ''
-            ${pkgs.python3}/bin/python3 ${generateRangesScript} > $out
-          '';
-
-        generateSubnetScript =
-          builtins.toFile "subnet-without-host-bits-generator.py" ''
-            import ipaddress
-            n1 = ipaddress.ip_network('${cfg.vlanIP}/${toString cfg.lanSize}', False)
-            print(n1, end="")
-          '';
-
-        subnetOutput =
-          pkgs.runCommandNoCC "subnet-without-host-bits"
-            { } ''
-            ${pkgs.python3}/bin/python3 ${generateSubnetScript} > $out
-          '';
-
-        ranges = lib.splitString ":" (builtins.readFile "${rangesOutput}");
-        subnet = builtins.readFile "${subnetOutput}";
-      in
-      {
-        ips = [ "${cfg.vlanIP}/32" ];
-        privateKeyFile = cfg.privateKeyFile;
-
-        peers = [
-          {
-            publicKey = cfg.serverKey;
-            allowedIPs = if cfg.proxy then ranges else [ subnet ];
-            endpoint = "${cfg.serverIP}:${toString cfg.serverPort}";
-            persistentKeepalive = 25;
-          }
-        ];
-      };
   };
+
+  config = mkIf cfg.enable
+    {
+      networking.firewall.checkReversePath = false;
+
+      systemd.paths = mapAttrs'
+        (instName: inst: {
+          name = "wireguard-${instName}";
+          value = if inst.autoStart then { } else { wantedBy = mkForce [ ]; };
+        })
+        cfg.instances;
+
+      systemd.services = lib.listToAttrs (flatten
+        (mapAttrsToList
+          (instName: inst: [
+            {
+              name = "wireguard-${instName}";
+              value = if inst.autoStart then { } else { wantedBy = mkForce [ ]; };
+            }
+          ] ++ map
+            (server: {
+              name = "wireguard-${instName}-peer${peerNameReplacement server.publicKey}";
+              value = if inst.autoStart then { } else { wantedBy = mkForce [ ]; };
+            })
+            inst.servers)
+          cfg.instances));
+
+
+      networking.wireguard.interfaces = mapAttrs
+        (instName: inst:
+          let
+            allowedIPs = if inst.proxy then ranges (map (s: s.ip) inst.servers) else [ (subnet inst.vlanIP inst.vlanSize) ];
+          in
+          {
+            ips = [ "${inst.vlanIP}/${toString inst.vlanSize}" ];
+            privateKeyFile = inst.privateKeyFile;
+            peers = map
+              (server: {
+                inherit allowedIPs;
+                publicKey = server.publicKey;
+                endpoint = "${server.ip}:${toString server.port}";
+                persistentKeepalive = 25;
+              })
+              inst.servers;
+          })
+        cfg.instances;
+
+    };
 }