# This module enables Network Address Translation (NAT).
# XXX: todo: support multiple upstream links
# see http://yesican.chsoft.biz/lartc/MultihomedLinuxNetworking.html

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

with lib;

let
  cfg = config.networking.nat;

  mkDest =
    externalIP: if externalIP == null then "-j MASQUERADE" else "-j SNAT --to-source ${externalIP}";
  dest = mkDest cfg.externalIP;
  destIPv6 = mkDest cfg.externalIPv6;

  # Whether given IP (plus optional port) is an IPv6.
  isIPv6 = ip: builtins.length (lib.splitString ":" ip) > 2;

  helpers = import ./helpers.nix { inherit config lib; };

  flushNat = ''
    ${helpers}
    ip46tables -w -t nat -D PREROUTING -j nixos-nat-pre 2>/dev/null|| true
    ip46tables -w -t nat -F nixos-nat-pre 2>/dev/null || true
    ip46tables -w -t nat -X nixos-nat-pre 2>/dev/null || true
    ip46tables -w -t nat -D POSTROUTING -j nixos-nat-post 2>/dev/null || true
    ip46tables -w -t nat -F nixos-nat-post 2>/dev/null || true
    ip46tables -w -t nat -X nixos-nat-post 2>/dev/null || true
    ip46tables -w -t nat -D OUTPUT -j nixos-nat-out 2>/dev/null || true
    ip46tables -w -t nat -F nixos-nat-out 2>/dev/null || true
    ip46tables -w -t nat -X nixos-nat-out 2>/dev/null || true
    ip46tables -w -t filter -D FORWARD -j nixos-filter-forward 2>/dev/null || true
    ip46tables -w -t filter -F nixos-filter-forward 2>/dev/null || true
    ip46tables -w -t filter -X nixos-filter-forward 2>/dev/null || true

    ${cfg.extraStopCommands}
  '';

  mkSetupNat =
    {
      iptables,
      dest,
      internalIPs,
      forwardPorts,
      externalIp,
    }:
    ''
      # We can't match on incoming interface in POSTROUTING, so
      # mark packets coming from the internal interfaces.
      ${concatMapStrings (iface: ''
        ${iptables} -w -t nat -A nixos-nat-pre \
          -i '${iface}' -j MARK --set-mark 1
        ${iptables} -w -t filter -A nixos-filter-forward \
          -i '${iface}' ${
            optionalString (cfg.externalInterface != null) "-o ${cfg.externalInterface}"
          } -j ACCEPT
      '') cfg.internalInterfaces}

      # NAT the marked packets.
      ${optionalString (cfg.internalInterfaces != [ ]) ''
        ${iptables} -w -t nat -A nixos-nat-post -m mark --mark 1 \
          ${optionalString (cfg.externalInterface != null) "-o ${cfg.externalInterface}"} ${dest}
      ''}

      # NAT packets coming from the internal IPs.
      ${concatMapStrings (range: ''
        ${iptables} -w -t nat -A nixos-nat-post \
          -s '${range}' ${
            optionalString (cfg.externalInterface != null) "-o ${cfg.externalInterface}"
          } ${dest}
        ${iptables} -w -t filter -A nixos-filter-forward \
          -s '${range}' ${
            optionalString (cfg.externalInterface != null) "-o ${cfg.externalInterface}"
          } -j ACCEPT
      '') internalIPs}

      # Related connections are allowed
      ${iptables} -w -t filter -A nixos-filter-forward \
        -m state --state ESTABLISHED,RELATED -j ACCEPT

      # NAT from external ports to internal ports.
      ${concatMapStrings (fwd: ''
        ${iptables} -w -t nat -A nixos-nat-pre \
          -i ${toString cfg.externalInterface} -p ${fwd.proto} \
          ${optionalString (externalIp != null) "-d ${externalIp}"} --dport ${toString fwd.sourcePort} \
          -j DNAT --to-destination ${fwd.destination}
        ${iptables} -w -t filter -A nixos-filter-forward \
          -i ${toString cfg.externalInterface} -p ${fwd.proto} \
          --dport ${toString fwd.sourcePort} -j ACCEPT

        ${concatMapStrings (
          loopbackip:
          let
            matchIP = if isIPv6 fwd.destination then "[[]([0-9a-fA-F:]+)[]]" else "([0-9.]+)";
            m = builtins.match "${matchIP}:([0-9-]+)" fwd.destination;
            destinationIP = if m == null then throw "bad ip:ports `${fwd.destination}'" else elemAt m 0;
            destinationPorts =
              if m == null then
                throw "bad ip:ports `${fwd.destination}'"
              else
                builtins.replaceStrings [ "-" ] [ ":" ] (elemAt m 1);
          in
          ''
            # Allow connections to ${loopbackip}:${toString fwd.sourcePort} from the host itself
            ${iptables} -w -t nat -A nixos-nat-out \
              -d ${loopbackip} -p ${fwd.proto} \
              --dport ${toString fwd.sourcePort} \
              -j DNAT --to-destination ${fwd.destination}

            # Allow connections to ${loopbackip}:${toString fwd.sourcePort} from other hosts behind NAT
            ${concatMapStrings (range: ''
              ${iptables} -w -t nat -A nixos-nat-pre \
                -d ${loopbackip} -p ${fwd.proto} -s '${range}' \
                --dport ${toString fwd.sourcePort} \
                -j DNAT --to-destination ${fwd.destination}
              ${iptables} -w -t nat -A nixos-nat-post \
                -d ${destinationIP} -p ${fwd.proto} \
                -s '${range}' --dport ${destinationPorts} \
                -j SNAT --to-source ${loopbackip}
              ${iptables} -w -t filter -A nixos-filter-forward \
                -d ${destinationIP} -p ${fwd.proto} \
                -s '${range}' --dport ${destinationPorts} -j ACCEPT
            '') internalIPs}
            ${concatMapStrings (iface: ''
              ${iptables} -w -t nat -A nixos-nat-pre \
                -d ${loopbackip} -p ${fwd.proto} -i '${iface}' \
                --dport ${toString fwd.sourcePort} \
                -j DNAT --to-destination ${fwd.destination}
              ${iptables} -w -t nat -A nixos-nat-post \
                -d ${destinationIP} -p ${fwd.proto} \
                -i '${iface}' --dport ${destinationPorts} \
                -j SNAT --to-source ${loopbackip}
              ${iptables} -w -t filter -A nixos-filter-forward \
                -d ${destinationIP} -p ${fwd.proto} \
                -i '${iface}' --dport ${destinationPorts} -j ACCEPT
            '') cfg.internalInterfaces}
          ''
        ) fwd.loopbackIPs}
      '') forwardPorts}
    '';

  setupNat = ''
    ${helpers}
    # Create subchains where we store rules
    ip46tables -w -t nat -N nixos-nat-pre
    ip46tables -w -t nat -N nixos-nat-post
    ip46tables -w -t nat -N nixos-nat-out
    ip46tables -w -t filter -N nixos-filter-forward

    ${mkSetupNat {
      iptables = "iptables";
      inherit dest;
      inherit (cfg) internalIPs;
      forwardPorts = filter (x: !(isIPv6 x.destination)) cfg.forwardPorts;
      externalIp = cfg.externalIP;
    }}

    ${optionalString cfg.enableIPv6 (mkSetupNat {
      iptables = "ip6tables";
      dest = destIPv6;
      internalIPs = cfg.internalIPv6s;
      forwardPorts = filter (x: isIPv6 x.destination) cfg.forwardPorts;
      externalIp = cfg.externalIPv6;
    })}

    ${optionalString (cfg.dmzHost != null) ''
      iptables -w -t nat -A nixos-nat-pre \
        -i ${toString cfg.externalInterface} -j DNAT \
        --to-destination ${cfg.dmzHost}
    ''}

    ${cfg.extraCommands}

    # Append our chains to the nat tables
    ip46tables -w -t nat -A PREROUTING -j nixos-nat-pre
    ip46tables -w -t nat -A POSTROUTING -j nixos-nat-post
    ip46tables -w -t nat -A OUTPUT -j nixos-nat-out
    ip46tables -w -t filter -A FORWARD -j nixos-filter-forward
  '';

in

{

  options = {

    networking.nat.extraCommands = mkOption {
      type = types.lines;
      default = "";
      example = "iptables -A INPUT -p icmp -j ACCEPT";
      description = ''
        Additional shell commands executed as part of the nat
        initialisation script.

        This option is incompatible with the nftables based nat module.
      '';
    };

    networking.nat.extraStopCommands = mkOption {
      type = types.lines;
      default = "";
      example = "iptables -D INPUT -p icmp -j ACCEPT || true";
      description = ''
        Additional shell commands executed as part of the nat
        teardown script.

        This option is incompatible with the nftables based nat module.
      '';
    };

  };

  config = mkIf (!config.networking.nftables.enable) (mkMerge [
    { networking.firewall.extraCommands = mkBefore flushNat; }
    (mkIf config.networking.nat.enable {

      networking.firewall = mkIf config.networking.firewall.enable {
        extraCommands = setupNat;
        extraStopCommands = flushNat;
      };

      systemd.services = mkIf (!config.networking.firewall.enable) {
        nat = {
          description = "Network Address Translation";
          wantedBy = [ "network.target" ];
          after = [
            "network-pre.target"
            "systemd-modules-load.service"
          ];
          path = [ config.networking.firewall.package ];
          unitConfig.ConditionCapability = "CAP_NET_ADMIN";

          serviceConfig = {
            Type = "oneshot";
            RemainAfterExit = true;
          };

          script = flushNat + setupNat;

          postStop = flushNat;
        };
      };
    })
  ]);
}
