dotfiles/nixos/modules/notnft.nix

{ pkgs, config, lib, notnft, ... }:
let
  inherit (lib)
    types
    mkOption
    mkDefault
    mkEnableOption
    flip
    concatMapStringsSep
    optionalAttrs
    listToAttrs
    optional
    filter
    optionalString
    ;
  cfg = config.networking.notnft;
  jsonFormat = pkgs.formats.json {};
in {
  options.networking.notnft = {
    enable = mkEnableOption "notnft";

    preStart = mkOption {
      type = types.lines;
      default = "";
    };

    postStart = mkOption {
      type = types.lines;
      default = "";
    };

    preStop = mkOption {
      type = types.lines;
      default = "";
    };

    postStop = mkOption {
      type = types.lines;
      default = "";
    };

    preRules = mkOption {
      type = types.listOf jsonFormat.type;
      default = [];
    };

    rules = mkOption {
      type = notnft.types.ruleset;
      default = {};
    };

    postRules = mkOption {
      type = types.listOf jsonFormat.type;
      default = [];
    };

    json = mkOption {
      type = jsonFormat.type;
      readOnly = true;
    };

    jsonFile = mkOption {
      type = types.path;
      readOnly = true;
    };

    flush = mkOption {
      type = types.bool;
      default = true;
    };

    chains = {
      dnsDrop = {
        enable = mkEnableOption "Add dns-drop chain";

        rule = mkOption {
          type = notnft.type.rule;
          readOnly = true;
          default = with notnft.dsl; with payload;
            [ jump "dns-drop" ];
        };
      };
    };
  };

  config = {
    networking.notnft.rules = with notnft.dsl; with payload; ruleset {
      filter = add table { family = f: f.inet; }
        (listToAttrs (filter (x: x != {}) [
          (optionalAttrs cfg.chains.dnsDrop.enable {
            name = "dns-drop";
            value = add chain
              [ (is.ne ip.daddr "10.64.2.1") (is.eq ip.protocol (f: with f; set [ tcp udp ])) (is.eq th.dport 53) drop ];
          })
        ]));
    };

    networking.notnft.json = builtins.toJSON {
      nftables = (optional cfg.flush { flush.ruleset = null; }) ++ cfg.preRules ++ cfg.rules.nftables ++ cfg.postRules;
    };
    networking.notnft.jsonFile = pkgs.writeText "rules.json" cfg.json;

    boot.blacklistedKernelModules = [ "ip_tables" ];
    environment.systemPackages = [ pkgs.nftables ];
    # networking.networkmanager.firewallBackend = mkDefault "nftables";
    systemd.services.notnftables = {
      description = "notnftables firewall";
      before = [ "network-pre.target" ];
      wants = [ "network-pre.target" ];
      wantedBy = [ "multi-user.target" ];
      reloadIfChanged = true;
      serviceConfig = let
        startScript = pkgs.writeShellScript "start-nft.sh" ''
          ${cfg.preStart}
          ${pkgs.buildPackages.nftables}/bin/nft -j -f ${cfg.jsonFile}
          ${cfg.postStart}
        '';

        stopScript = pkgs.writeShellScript "stop-nft.sh" ''
          ${cfg.preStop}
          ${optionalString cfg.flush "${pkgs.nftables}/bin/nft flush ruleset"}
          ${cfg.postStop}
        '';
        # rulesScript = pkgs.writeTextFile {
        #   name =  "nftables-rules";
        #   executable = true;
        #   text = ''
        #     #! ${pkgs.nftables}/bin/nft -f
        #     flush ruleset
        #     ${if cfg.rulesetFile != null then ''
        #       include "${cfg.rulesetFile}"
        #     '' else cfg.ruleset}
        #   '';
        #   checkPhase = lib.optionalString cfg.checkRuleset ''
        #     cp $out ruleset.conf
        #     ${cfg.preCheckRuleset}
        #     export NIX_REDIRECTS=/etc/protocols=${pkgs.buildPackages.iana-etc}/etc/protocols:/etc/services=${pkgs.buildPackages.iana-etc}/etc/services
        #     LD_PRELOAD="${pkgs.buildPackages.libredirect}/lib/libredirect.so ${pkgs.buildPackages.lklWithFirewall.lib}/lib/liblkl-hijack.so" \
        #     ${pkgs.buildPackages.nftables}/bin/nft --check -j < ${cfg.jsonFile}
        #   '';
        # };
      in {
        Type = "oneshot";
        RemainAfterExit = true;
        ExecStart = startScript;
        ExecReload = startScript;
        ExecStop = stopScript;
      };
    };
  };
}