NixMachines/machines/hippocampus/modules/headscale.nix

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

with lib;
let
  cfg = config.services.headscale;

  userOptions = { config, ... }: {
    options = {
      enablePreAuth = mkOption {
        type = types.bool;
        default = true;
        description = lib.mdDoc ''
          Generate Pre Authorized Token with User
        '';
      };
      preAuthEphemeral = mkOption {
        type = types.bool;
        default = false;
        description = lib.mdDoc ''
          Should the token be ephemeral, making the user
          dissappear after no usage
        '';
      };
      preAuthExpiration = mkOption {
        type = types.str;
        default = "1h";
        description = lib.mdDoc ''
          How long should the token be active for
        '';
      };
      preAuthReusable = mkOption {
        type = types.bool;
        default = false;
        description = lib.mdDoc ''
          Should the token be able to be used more than once
        '';
      };
      preAuthTags = mkOption {
        type = types.listOf types.str;
        default = [];
        description = lib.mdDoc ''
          How should this login token be tagged
        '';
      };
      preAuthRegen = mkOption {
        type = types.bool;
        default = false;
        description = lib.mdDoc ''
          Should a timer be built to regenerate the token
        '';
      };
      path = mkOption {
        type = types.path;
        default = "/run/headscale_${config._module.args.name}_auth/auth";
        description = mdDoc ''
          Path of generated token
        '';
      };
    };
  };

  mkUserService = name: options: {
    "headscale-ensureUser-${name}" = {
      description = "Ensure '${name}' user exists for headscale";
      wantedBy = ["multi-user.target" "headscale.service"];
      after = ["headscale.service"];
      requires = ["headscale.service"];
      partOf = ["headscale.service"];

      script = ''
        ${pkgs.bash}/bin/bash -c '${cfg.package}/bin/headscale users create ${name} || true'
      '';

      serviceConfig = {
        Type = "oneshot";
        User = cfg.user;
        Group = cfg.user;
      };
    };
  };

  # Make Auth token readable with given set of perms and also user?
  # also generate on timer

  mkUserPreAuth = name: options: optionalAttrs options.enablePreAuth {
    "headscale-preauth-${name}" = {
      description = "Generate Headscale Preauth Token for '${name}'";
      wantedBy = ["multi-user.target"];
      after = ["headscale-ensureUser-${name}.service"];
      requires = ["headscale-ensureUser-${name}.service"];
      partOf = ["headscale.service"];

      script = ''
        ${cfg.package}/bin/headscale preauthkeys -u $(${cfg.package}/bin/headscale users -o json-line list | ${pkgs.jq}/bin/jq '.[] | select(.name=="${name}").id') create \
          ${lib.optionalString options.preAuthEphemeral "--ephemeral"} \
          ${lib.optionalString options.preAuthReusable "--reusable"} \
          --expiration ${options.preAuthExpiration} \
          ${lib.optionalString (options.preAuthTags != []) "--tags ${toString options.preAuthTags}"} \
          > /run/headscale_${name}_auth/auth; echo /run/headscale_${name}_auth/auth created
      '';

      # TODO: Use Activation Script to Generate On Boot
      # with option to regenerate with timer

      serviceConfig = {
        Type = "oneshot";
        PermissionsStartOnly = "true";
        # Directory Creation Issues Don't Belong To User/Aren't Passed
        # with the specified user, resulting in %t resolving to root's
        # runtime directory
        # https://github.com/containers/podman/issues/12778
        RuntimeDirectory = "headscale_${name}_auth";
        RuntimeDirectoryMode = "0775";
        User = cfg.user;
        Group = cfg.user;
        RemainAfterExit = "yes";
      };
    };
  };

  mkUserPreAuthRegen = name: options: {
    "headscale-preauth-regen-${name}" = {
      wantedBy = ["multi-user.target" "headscale-ensureUser-${name}.service"];
      after = ["headscale-ensureUser-${name}.service"];

      script = ''
        ${pkgs.systemd}/bin/systemctl restart headscale-preauth-${name}.service
      '';
      serviceConfig = {
        Type = "oneshot";
      };
    };
  };

  mkServices = name: options: (mkUserService name options) //
                              (mkUserPreAuth name options) //
                              (mkUserPreAuthRegen name options);

  mkTimer = name: options: optionalAttrs options.preAuthRegen {
    wantedBy = ["multi-user.target" "headscale-ensureUser-${name}.service"];
    after = ["headscale-ensureUser-${name}.service"];
    requires = ["headscale-ensureUser-${name}.service"];
    bindsTo = ["headscale.service"];
    partOf = ["headscale.service"];

    timerConfig = {
      OnUnitActiveSec = options.preAuthExpiration;
      Unit = "headscale-preauth-regen-${name}.service";
    };
  };

in
{
  options.services.headscale.ensureUsers = mkOption {
    default = {};
    type = types.attrsOf (types.submodule userOptions);
    description = lib.mdDoc "Ensure these users exist in headscale";
  };
  config = lib.mkIf (cfg.enable && cfg.ensureUsers != {}) {
    systemd.services = (foldl' (a: b: a // b) {} (attrValues ((mapAttrs
      (n: v: (mkServices n v)) cfg.ensureUsers))));

    systemd.timers = filterAttrs (n: v: v !={}) (mapAttrs'
      (n: v: nameValuePair "headscale-preauth-${n}" (mkTimer n v)) cfg.ensureUsers);
  };
}