dendritic/modules/features/mtls.nix

{ self, inputs, lib, ... }:
let
  # Options that will be in common between the nixos module and the home-manager module.
  mkOpts = config: let cfg = config.mtls; in {
    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;
    };
    caFile = lib.mkOption {
      description = "String path for the root CA file";
      type = lib.types.str;
      default = "${cfg.certDir}/root_ca.crt";
    };
    keyFile = lib.mkOption {
      description = "String path for the private key";
      type = lib.types.str;
      default = "${cfg.certDir}/key.pem";
    };
    certFile = lib.mkOption {
      description = "String path for the public cert";
      type = lib.types.str;
      default = "${cfg.certDir}/cert.pem";
    };
    bundleFile = lib.mkOption {
      description = "String path for the mTLS key bundle";
      type = lib.types.str;
      default = "${cfg.certDir}/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 = "24h";
    };
    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.";
        type = lib.types.nullOr lib.types.str;
        default = "mtls";
      };
      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,
    certFile,
    keyFile,
    bundleFile,
    lifetime,
    user,
    group,
  }:
  let
    sanArgs = lib.concatMapStringsSep " " (s: "--san \"${s}\"") san;
  in
    pkgs.writeShellApplication {
      name = "mtls-generate";
      runtimeInputs = with pkgs; [ coreutils step-cli ];
      text = ''
        set -euo pipefail
        step ca certificate ${subject} ${certFile} ${keyFile} \
          --provisioner ${provisioner} \
          --not-before=-5m --not-after=${lifetime} \
          ${sanArgs} \
          "$@"
        (umask 077; cat ${certFile} ${keyFile} > ${bundleFile})
        chown ${user}:${group} ${certFile} ${keyFile} ${bundleFile}
        chmod 640 ${certFile} ${keyFile} ${bundleFile}
        printf '\033[32m✔\033[0m \033[1mmTLS Bundle:\033[0m %s\n' ${lib.escapeShellArg bundleFile}
      '';
    };
  
  mkMtlsCheckScript = { pkgs, bundleFile }: pkgs.writeShellApplication {
    name = "mtls-check";
    runtimeInputs = with pkgs; [ openssl ];
    text = ''
      openssl x509 -noout -in ${bundleFile} \
        -subject -issuer \
        -ext subjectAltName,extendedKeyUsage \
        -enddate
    '';
  };

  mkMtlsRenewScript = {
    pkgs,
    cfg,
    systemctlArgs ? [ ],
  }:
  let
    systemctlCmd = "systemctl ${lib.escapeShellArgs systemctlArgs}";

    hasReloadUnits = cfg.renew.reloadUnits != [ ];
    renewReloadScript = lib.concatMapStringsSep "\n" (unit: ''
      if ${systemctlCmd} --quiet is-active "${unit}"; then
        ${systemctlCmd} try-reload-or-restart "${unit}"
      fi
    '') cfg.renew.reloadUnits;

    hasPostCommands = cfg.renew.postCommands != [ ];
    renewPostCommands = lib.concatStringsSep "\n" cfg.renew.postCommands;

    fileOwner = "${cfg.renew.user}:${cfg.renew.group}";
  in
    pkgs.writeShellApplication {
      name = "mtls-renew";
      runtimeInputs = with pkgs; [ coreutils step-cli systemd ];
      text = ''
      set -euo pipefail

      YELLOW_BANG="\e[33m!\e[0m"

      force=0
      while [[ $# -gt 0 ]]; do
      case $1 in
        --force)
        force=1
        shift
        ;;
        *)
        echo -e "$YELLOW_BANG Warning: ignoring unrecognized argument '$1'"
        exit 1
        ;;
      esac
      done

      if [[ $force -eq 0 ]] && ! step certificate needs-renewal "${cfg.certFile}"; then
        echo "Skipping renew"
        exit 0
      fi

      echo "Renewing mTLS certificate"
      step ca renew --force "${cfg.certFile}" "${cfg.keyFile}"
      (umask 077; cat "${cfg.certFile}" "${cfg.keyFile}" > "${cfg.bundleFile}")
      chown ${fileOwner} ${cfg.certFile} ${cfg.keyFile} ${cfg.bundleFile}
      chmod 640 ${cfg.certFile} ${cfg.keyFile} ${cfg.bundleFile}

      ${lib.optionalString hasReloadUnits ''
        echo "Reloading units: ${lib.concatStringsSep ", " cfg.renew.reloadUnits}"
        ${renewReloadScript}
      ''}

      ${lib.optionalString hasPostCommands ''echo "Post commands:" ${renewPostCommands}''}
      '';
    };

  mkNixosMtlsRenewService = { pkgs, cfg, ... }:
    {
      description = "Renew the mTLS certificate when Smallstep marks it ready";
      wantedBy = [ ];
      after = [ "network-online.target" ];
      wants = [ "network-online.target" ];
      serviceConfig = {
        Type = "oneshot";
        User = cfg.renew.user;
        Group = cfg.renew.group;
        ExecStart = lib.getExe (mkMtlsRenewScript { inherit pkgs cfg; });
      };
    };

  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, cfg, ... }:
  let
    renewScript = mkMtlsRenewScript {
      inherit pkgs cfg;
      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;
      sanArgs = lib.concatMapStringsSep " " (san: "--san \"${san}\"") cfg.san;
    in
    {
      options.mtls = (mkOpts config) // {
        certDir = lib.mkOption {
          description = "String path to where the mtls certs will be stored.";
          type = lib.types.str;
          default = "/etc/step-ca/certs";
        };
        certReaders = lib.mkOption {
          description = "";
          type = lib.types.listOf lib.types.str;
          default = [ ];
        };
      };

      config = lib.mkIf cfg.enable {
        users.groups.certReaders = {
          name = cfg.renew.group;
          members = cfg.certReaders;
        };

        environment.systemPackages = with pkgs; lib.optionals cfg.enable [
          # step-cli
          (mkMtlsGenerateScript {
            inherit pkgs;
            inherit (cfg) subject provisioner san certFile keyFile bundleFile lifetime;
            inherit (cfg.renew) user group;
          })
          (mkMtlsCheckScript { inherit pkgs; inherit (cfg) bundleFile; })
          (mkMtlsRenewScript { inherit pkgs cfg; })
        ];

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

        systemd.services.mtls-renew = lib.mkIf cfg.renew.enable
          (mkNixosMtlsRenewService { inherit pkgs cfg; });

        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;
      keyFile = cfg.keyFile;
      certFile = cfg.certFile;
      bundleFile = cfg.bundleFile;
      sanArgs = lib.concatMapStringsSep " " (san: "--san \"${san}\"") cfg.san;
    in
    {
      options.mtls = (mkOpts config) // {
        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 (cfg.renew) user group;
            inherit pkgs certFile keyFile bundleFile;
          })
          (mkMtlsCheckScript { inherit pkgs bundleFile; })
          (mkMtlsRenewScript { inherit pkgs cfg; systemctlArgs = [ "--user" ]; })
        ];

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

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