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;
      };
    };
  };
}