desktop/wireguard.nix

{
  config,
  lib,
  pkgs,
  ...
}:

with lib;

let
  cfg = config.services.wireguard-netns;

  vpnOptions =
    { name, ... }:
    {
      options = {
        name = mkOption {
          type = types.str;
          default = name;
          description = "Name of the network namespace and WireGuard interface";
        };

        dns = mkOption {
          type = types.str;
          description = "Network DNS";
          example = "1.1.1.1";
        };

        address = mkOption {
          type = types.str;
          description = "Address";
          example = "10.2.0.1/32";
        };

        conf = mkOption {
          type = types.str;
          description = "Path to sops secret containing WireGuard configuration";
          example = "wg0/conf";
        };
      };
    };

in
{
  options.services.wireguard-netns = {
    enable = mkEnableOption "WireGuard network namespaces";

    namespaces = mkOption {
      type = types.attrsOf (types.submodule vpnOptions);
      default = { };
      description = "WireGuard VPN configurations in separate network namespaces";
      example = literalExpression ''
        {
          wg0ns = {
            dns = "10.2.0.1";
            address = "10.2.0.2/32";
            conf = "wg0/conf";
          };
          wg1ns = {
            dns = "10.3.0.1";
            address = "10.3.0.2/32";
            conf = "wg1/conf";
          };
        }
      '';
    };
  };

  config = mkIf cfg.enable {
    systemd.services = listToAttrs (
      flatten (
        imap1 (
          index: elem:
          let
            name = elem.name;
            vpnCfg = elem.value;
            ns = "${name}ns";
            nsIP = "10.200.${toString index}.2/24";
            hostIP = "10.200.${toString index}.1/24";
            vethHost = "veth-${name}";
            vethNS = "veth-${name}-ns";
            proxyIP = "10.200.${toString index}.2"; # IP without CIDR for binding
            proxyPort = 1080;
          in
          [
            {
              name = "netns@${ns}";
              value = {
                description = "${ns} network namespace";
                before = [ "network.target" ];
                serviceConfig = {
                  Type = "oneshot";
                  RemainAfterExit = true;
                  ExecStart = pkgs.writers.writeBash "${ns}-up" ''
                    ${pkgs.coreutils}/bin/mkdir -p /etc/netns/${ns}
                    echo "nameserver ${vpnCfg.dns}" > /etc/netns/${ns}/resolv.conf
                    ${pkgs.iproute2}/bin/ip netns add ${ns}

                    ${pkgs.iproute2}/bin/ip link add ${vethHost} type veth peer name ${vethNS}
                    ${pkgs.iproute2}/bin/ip link set ${vethNS} netns ${ns}

                    ${pkgs.iproute2}/bin/ip addr add ${hostIP} dev ${vethHost}
                    ${pkgs.iproute2}/bin/ip link set ${vethHost} up

                    ${pkgs.iproute2}/bin/ip -n ${ns} addr add ${nsIP} dev ${vethNS}
                    ${pkgs.iproute2}/bin/ip -n ${ns} link set ${vethNS} up
                  '';
                  ExecStop = pkgs.writers.writeBash "${ns}-down" ''
                    ${pkgs.iproute2}/bin/ip link del ${vethHost} || true
                    ${pkgs.iproute2}/bin/ip netns del ${ns}
                  '';
                };
              };
            }

            {
              name = name;
              value = {
                description = "${name} network interface";
                bindsTo = [ "netns@${ns}.service" ];
                requires = [ "network-online.target" ];
                after = [ "netns@${ns}.service" ];
                wants = [ "network-online.target" ];
                wantedBy = [ "multi-user.target" ];
                serviceConfig = {
                  Type = "oneshot";
                  RemainAfterExit = true;
                  ExecStart = pkgs.writers.writeBash "${name}-up" ''
                    ${pkgs.iproute2}/bin/ip link add ${name} type wireguard
                    ${pkgs.iproute2}/bin/ip link set ${name} netns ${ns}
                    ${pkgs.iproute2}/bin/ip -n ${ns} address add ${vpnCfg.address} dev ${name}
                    ${pkgs.iproute2}/bin/ip netns exec ${ns} \
                        ${pkgs.wireguard-tools}/bin/wg setconf ${name} ${config.sops.secrets.${vpnCfg.conf}.path}
                    ${pkgs.iproute2}/bin/ip -n ${ns} link set lo up
                    ${pkgs.iproute2}/bin/ip -n ${ns} link set ${name} up
                    ${pkgs.iproute2}/bin/ip -n ${ns} route add default dev ${name}
                  '';
                  ExecStop = pkgs.writers.writeBash "${name}-down" ''
                    ${pkgs.iproute2}/bin/ip -n ${ns} route del default dev ${name} || true
                    ${pkgs.iproute2}/bin/ip -n ${ns} link del ${name} || true
                  '';
                };
              };
            }

            {
              name = "socks-${name}";
              value = {
                description = "SOCKS5 proxy for ${name} namespace";
                after = [ "${name}.service" ];
                requires = [ "${name}.service" ];
                wantedBy = [ "multi-user.target" ];
                serviceConfig = {
                  Type = "simple";
                  Restart = "always";
                  RestartSec = "5s";
                  ExecStart = "${pkgs.iproute2}/bin/ip netns exec ${ns} ${pkgs.microsocks}/bin/microsocks -i ${proxyIP} -p ${toString proxyPort}";
                };
              };
            }
          ]
        ) (lib.attrsToList cfg.namespaces)
      )
    );
  };
}