nix/shared/commons/nftables.nix

{ config, lib, ... }:

let
    # Import nodes
    nodes = import ./../../nodes.nix;
    myName = config.hostName;

    # Import mapping
    mapping = import ./../../mapping.nix;

    generatePortRules = protocol: ports: let
        publicPorts = lib.filter (p: p.public) ports;
        privatePorts = lib.filter (p: !p.public) ports;
        
        publicRules = map (p: "${protocol} dport ${toString p.port} accept") publicPorts;
        privateRules = map (p: "iifname mesh ${protocol} dport ${toString p.port} accept") privatePorts;
    in
        publicRules ++ privateRules;

    cfg = config.fwtables;
in
{
    options.fwtables = with lib; {
        allowedTCPPorts = mkOption {
            type = types.listOf (types.submodule {
                options = {
                    port = mkOption {
                        type = types.port;
                        description = "The TCP port number";
                    };
                    public = mkOption {
                        type = types.bool;
                        default = false;
                        description = "Whether the port should be accessible from public internet (true) or only from mesh network (false)";
                    };
                };
            });
            default = [];
            description = "List of allowed TCP ports with their accessibility settings";
            example = [
                { port = 80; public = true; }
                { port = 8080; public = false; }
            ];
        };
        
        allowedUDPPorts = mkOption {
            type = types.listOf (types.submodule {
                options = {
                    port = mkOption {
                        type = types.port;
                        description = "The UDP port number";
                    };
                    public = mkOption {
                        type = types.bool;
                        default = false;
                        description = "Whether the port should be accessible from public internet (true) or only from mesh network (false)";
                    };
                };
            });
            default = [];
            description = "List of allowed UDP ports with their accessibility settings";
            example = [
                { port = 53; public = true; }
                { port = 1234; public = false; }
            ];
        };

        allowedMgmtFwdToMesh = mkOption {
            type = types.bool;
            default = false;
            description = "Allow traffic to jump from mgmt if to mesh if";
        };
    };
    
    config = {
        networking = {
            nat.enable = false;
            firewall.enable = false;
            nftables = {
                enable = true;
                checkRuleset = true;
                flushRuleset = true;
                tables = {
                    filter = {
                        family = "inet";
                        content = ''
                            chain preroute {
                                type filter hook prerouting priority -150; policy accept;
                            }
                            chain input {
                                type filter hook input priority 0; policy drop;

                                # Authorized already setup connection
                                ct state related,established accept

                                # Reject sus stuff
                                ct state invalid counter drop
                                tcp flags & (fin|syn|rst|ack) != syn ct state new counter drop

                                # Loopback
                                iif lo accept

                                # ICMP
                                icmp type { echo-request } limit rate 4/second accept
                                icmpv6 type { echo-request } limit rate 4/second accept
                                ip protocol icmp accept
                                icmpv6 type { nd-neighbor-solicit, nd-router-advert, nd-neighbor-advert } accept

                                # Temporary SSH access for convenience during early stages
                                tcp dport 22 accept

                                ${lib.concatStringsSep "\n" (generatePortRules "tcp" cfg.allowedTCPPorts)}
                                ${lib.concatStringsSep "\n" (generatePortRules "udp" cfg.allowedUDPPorts)}

                                # Log anything else
                                ip protocol tcp counter log prefix "tcp.in.dropped: "
                                ip protocol udp counter log prefix "udp.in.dropped: "

                            }
                            chain forward {
                                type filter hook forward priority 0; policy drop;

                                ${lib.optionalString cfg.allowedMgmtFwdToMesh ''
                                    ct state related,established accept
                                    ct state invalid counter drop

                                    iifname mgmt oifname mesh accept
                                ''}
                            }
                            chain output {
                                type filter hook output priority 0; policy accept;
                            }
                            chain postroute {
                                type filter hook postrouting priority 50; policy accept;
                            }
                        '';
                    };
                    # We'll not nat the outgoing wg packet
                    # Instead custom routing is done on all
                    # other node to sent it back correclty
                    # That allow for custom ACL for the mgmt
                    # ip addresses
                    nat = {
                        family = "inet";
                        content = ''
                            chain prerouting {
                                type nat hook prerouting priority -100; policy accept;
                            }
                            chain postrouting {
                                type nat hook postrouting priority 100; policy accept;
                            }
                        '';
                    };
                };
            };
        };
    };
}