ezwg2.nix

{ config, lib, pkgs, ... }:
with lib;
let
  cfg = config.services.ezwg;

  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 {
      type = types.int;
      default = 24;
      description = "Size of your VLAN";
    };
    privateKeyFile = mkOption {
      type = types.str;
      description = "Private wg key";
    };
    vlanIP = mkOption {
      type = types.str;
      description = "The IP to use on the wg VLAN";
    };
  };
in
{
  options.services.ezwg = {
    enable = mkEnableOption "Enable simple Wireguard connection";
    instances = mkOption {
      description = "Configuration of instances of Wireguard";
      default = { };
      type = with types; attrsOf (submodule instanceOpts);
    };
  };

  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;

    };
}