dendritic/modules/features/mtls.nix

{ self, inputs, lib, ... }:
let
  # Options that will be in common between
  opts = {
    enable = lib.mkEnableOption "Enable mTLS";
    subject = lib.mkOption {
      description = "The Common Name, DNS Name, or IP address that will be set as the Subject Common Name for the certificate. If no Subject Alternative Names (SANs) are configured (via the --san flag) then the subject will be set as the only SAN.";
      type = lib.types.str;
    };
    certDir = lib.mkOption {
      description = "String path to the directory where the certs will be stored";
      type = lib.types.str;
    };
    keyFilename = lib.mkOption {
      description = "String filename for the private key";
      type = lib.types.str;
      default = "key.pem";
    };
    certFilename = lib.mkOption {
      description = "String filename for the public certificate";
      type = lib.types.str;
      default = "cert.pem";
    };
    bundleFilename = lib.mkOption {
      description = "String filename for the mTLS key bundle";
      type = lib.types.str;
      default = "mtls.pem";
    };
    san = lib.mkOption {
      description = "List of SAN to give the mTLS cert";
      type = lib.types.listOf lib.types.str;
      default = [ ];
    };
    provisioner = lib.mkOption {
      type = lib.types.str;
      default = "admin";
    };
    lifetime = lib.mkOption {
      type = lib.types.str;
      default = "6h";
    };
    renew = {
      enable = lib.mkOption {
        description = "Enable automatic mTLS certificate renewal using a systemd timer.";
        type = lib.types.bool;
        default = true;
      };
      onCalendar = lib.mkOption {
        description = "systemd OnCalendar schedule for mTLS certificate renewal checks.";
        type = lib.types.str;
        default = "*:1/15";
      };
      randomizedDelaySec = lib.mkOption {
        description = "Randomized delay added to renewal timer runs to avoid synchronized renewals.";
        type = lib.types.str;
        default = "5m";
      };
      user = lib.mkOption {
        description = "User account to run the mTLS renewal service as.";
        type = lib.types.str;
        default = "root";
      };
      group = lib.mkOption {
        description = "Group to run the mTLS renewal service as. Defaults to the configured renewal user when null.";
        type = lib.types.nullOr lib.types.str;
        default = null;
      };
      reloadUnits = lib.mkOption {
        description = "systemd units to try-reload-or-restart after a successful certificate renewal.";
        type = lib.types.listOf lib.types.str;
        default = [ ];
      };
      postCommands = lib.mkOption {
        description = "Shell commands to run after a successful certificate renewal.";
        type = lib.types.listOf lib.types.lines;
        default = [ ];
      };
    };
  };

  mkMtlsGenerateScript = {
    pkgs,
    subject,
    provisioner,
    san,
    tlsCert,
    tlsKey,
    mtlsBundle,
    lifetime,
  }:
  let
    stepCmd = lib.getExe pkgs.step-cli;
    sanArgs = lib.concatMapStringsSep " " (s: "--san \"${s}\"") san;
  in
    pkgs.writeShellScriptBin "mtls-generate" ''
      set -euo pipefail
      ${stepCmd} ca certificate \
        ${subject} ${tlsCert} ${tlsKey} \
        --not-before=-5m --not-after=${lifetime} \
        --provisioner ${provisioner} \
        ${sanArgs} \
        "$@"
      cat ${tlsCert} ${tlsKey} > ${mtlsBundle}
    '';
  
  mkMtlsCheckScript = { pkgs, mtlsBundle }: pkgs.writeShellScriptBin "mtls-check" ''
    ${lib.getExe pkgs.openssl} x509 \
    -noout -subject -issuer \
    -ext subjectAltName,extendedKeyUsage \
    -enddate -in ${mtlsBundle}
  '';

  mkMtlsRenewScript = {
    pkgs,
    tlsCert,
    tlsKey,
    mtlsBundle,
    reloadUnits ? [ ],
    postCommands ? [ ],
    systemctlArgs ? [ ],
  }:
  let
    echoCmd = lib.getExe' pkgs.coreutils "echo";
    systemctl = lib.getExe' pkgs.systemd "systemctl";
    escapedArgs = lib.escapeShellArgs systemctlArgs;
    systemctlCommand = "${systemctl} ${escapedArgs}";
    renewReloadScript = lib.concatMapStringsSep "\n" (unit: ''
      if ${systemctlCommand} --quiet is-active "${unit}"; then
        ${systemctlCommand} try-reload-or-restart "${unit}"
      fi
    '') reloadUnits;
    renewPostCommands = lib.concatStringsSep "\n" postCommands;
  in
    pkgs.writeShellScriptBin "mtls-renew" ''
      set -euo pipefail

      if ${lib.getExe pkgs.step-cli} certificate needs-renewal "${tlsCert}"; then
        ${echoCmd} "Renewing mTLS certificate"
      else
        ${echoCmd} "Skipping renew"
        exit "$?"
      fi

      ${lib.getExe pkgs.step-cli} ca renew --force "${tlsCert}" "${tlsKey}"

      umask 077
      ${lib.getExe' pkgs.coreutils "cat"} "${tlsCert}" "${tlsKey}" > "${mtlsBundle}"

      ${echoCmd} "Reloading units:"
      ${renewReloadScript}

      ${echoCmd} "Post commands:"
      ${renewPostCommands}
    '';

  mkNixosMtlsRenewService = {
    pkgs,
    tlsCert,
    tlsKey,
    mtlsBundle,
    reloadUnits ? [ ],
    postCommands ? [ ],
    user ? "root",
    group ? null,
  }:
  let
    serviceGroup = if group == null then user else group;
    renewScript = mkMtlsRenewScript {
      inherit pkgs tlsCert tlsKey mtlsBundle reloadUnits postCommands;
    };
  in
  {
    description = "Renew the mTLS certificate when Smallstep marks it ready";
    wantedBy = [ ];
    after = [ "network-online.target" ];
    wants = [ "network-online.target" ];
    serviceConfig = {
      Type = "oneshot";
      User = user;
      Group = serviceGroup;
      ExecStart = lib.getExe renewScript;
    };
  };

  mkNixosMtlsRenewTimer = {
    onCalendar,
    randomizedDelaySec,
    unit ? "mtls-renew.service",
  }: {
    description = "Periodic Smallstep renewal for the mTLS certificate";
    wantedBy = [ "timers.target" ];
    timerConfig = {
      Persistent = true;
      OnCalendar = onCalendar;
      AccuracySec = "1us";
      RandomizedDelaySec = randomizedDelaySec;
      Unit = unit;
    };
  };

  mkHomeManagerMtlsRenewService = {
    pkgs,
    tlsCert,
    tlsKey,
    mtlsBundle,
    reloadUnits ? [ ],
    postCommands ? [ ],
  }:
  let
    renewScript = mkMtlsRenewScript {
      inherit pkgs tlsCert tlsKey mtlsBundle reloadUnits postCommands;
      systemctlArgs = [ "--user" ];
    };
  in
  {
    Unit = {
      Description = "Renew the mTLS certificate when Smallstep marks it ready";
      After = [ "network-online.target" ];
      Wants = [ "network-online.target" ];
    };
    Service = {
      Type = "oneshot";
      ExecStart = lib.getExe renewScript;
    };
  };

  mkHomeManagerMtlsRenewTimer = {
    onCalendar,
    randomizedDelaySec,
    unit ? "mtls-renew.service",
  }: {
    Unit = {
      Description = "Periodic Smallstep renewal for the mTLS certificate";
    };
    Timer = {
      Persistent = true;
      OnCalendar = onCalendar;
      AccuracySec = "1us";
      RandomizedDelaySec = randomizedDelaySec;
      Unit = unit;
    };
    Install = {
      WantedBy = [ "timers.target" ];
    };
  };

in
{
  flake.modules.nixos.mtls = { config, lib, pkgs, ... }:
    let
      cfg = config.mtls;
      tlsKey = "${cfg.certDir}/${cfg.keyFilename}";
      tlsCert = "${cfg.certDir}/${cfg.certFilename}";
      mtlsBundle = "${cfg.certDir}/${cfg.bundleFilename}";
      sanArgs = lib.concatMapStringsSep " " (san: "--san \"${san}\"") cfg.san;
    in
    {
      options.mtls = opts // {
        certDir = lib.mkOption {
          description = "String path to where the mtls certs will be stored.";
          type = lib.types.str;
          default = "/etc/step/certs";
        };
      };

      config = lib.mkIf cfg.enable {
        environment.systemPackages = with pkgs; lib.optionals cfg.enable [
          # step-cli
          (mkMtlsGenerateScript {
            inherit (cfg) subject provisioner san lifetime;
            inherit pkgs tlsCert tlsKey mtlsBundle;
          })
          (mkMtlsCheckScript { inherit pkgs mtlsBundle; })
          (mkMtlsRenewScript { inherit pkgs tlsCert tlsKey mtlsBundle; })
        ];

        systemd.tmpfiles.rules = [
          "d ${cfg.certDir} 0750 ${cfg.renew.user} ${if cfg.renew.group == null then cfg.renew.user else cfg.renew.group} -"
        ];

        systemd.services.mtls-renew = lib.mkIf cfg.renew.enable (mkNixosMtlsRenewService {
          inherit pkgs tlsCert tlsKey mtlsBundle;
          inherit (cfg.renew) reloadUnits postCommands user group;
        });

        systemd.timers.mtls-renew = lib.mkIf cfg.renew.enable (mkNixosMtlsRenewTimer {
          inherit (cfg.renew) onCalendar randomizedDelaySec;
        });
      };
    };

  flake.modules.homeManager.mtls = { config, lib, pkgs, ... }:
    let
      cfg = config.mtls;
      tlsKey = "${cfg.certDir}/${cfg.keyFilename}";
      tlsCert = "${cfg.certDir}/${cfg.certFilename}";
      mtlsBundle = "${cfg.certDir}/${cfg.bundleFilename}";
      sanArgs = lib.concatMapStringsSep " " (san: "--san \"${san}\"") cfg.san;
    in
    {
      options.mtls = opts // {
        certDir = lib.mkOption {
          description = "String path to where the mtls certs will be stored.";
          type = lib.types.str;
          default = "${config.home.homeDirectory}/.step/certs";
        };
      };

      config = {
        home.packages = with pkgs; lib.optionals cfg.enable [
          # step-cli
          (mkMtlsGenerateScript {
            inherit (cfg) subject provisioner san lifetime;
            inherit pkgs tlsCert tlsKey mtlsBundle;
          })
          (mkMtlsCheckScript { inherit pkgs mtlsBundle; })
          (mkMtlsRenewScript { inherit pkgs tlsCert tlsKey mtlsBundle; })
        ];

        systemd.user.tmpfiles.rules = lib.mkIf cfg.enable [
          "d ${cfg.certDir} 0700 - - -"
        ];

        systemd.user.services.mtls-renew = lib.mkIf cfg.renew.enable (mkHomeManagerMtlsRenewService {
          inherit pkgs tlsCert tlsKey mtlsBundle;
          inherit (cfg.renew) reloadUnits postCommands;
        });

        systemd.user.timers.mtls-renew = lib.mkIf cfg.renew.enable (mkHomeManagerMtlsRenewTimer {
          inherit (cfg.renew) onCalendar randomizedDelaySec;
        });
      };
    };
}