# New attempt at consolidating SSH config that wraps the step client, SSH certs, and all that in a top-level module
{ self, inputs, ... }:
let
  caPatterns = [ "*.john-stream.com" "192.168.1.*" "fded:fb16:653e:25da:be24:11ff:*" ];
in
{
  flake.modules.nixos.ssh-new = { config, pkgs, lib, ... }:
    let
      cfg = config.ssh-new;
      hostKeyFile = "${cfg.host.configDir}/${cfg.host.keyFile}";
      hostCertFile = "${hostKeyFile}-cert.pub";
      userCAFile = "${cfg.host.configDir}/${cfg.certificates.user.CAFile}";
      sshHostCAContent = lib.removeSuffix "\n" (builtins.readFile ../hosts/janus/public/ssh_host_ca_key.pub);
      CAknownHosts = (lib.genAttrs caPatterns (_: {
          certAuthority = true;
          publicKey = sshHostCAContent;
      }));

      wrappers = inputs.self.wrappers;
      provisionerPasswordFile = config.sops.secrets."janus/admin_jwk".path;
      sshHostCertSign = (wrappers.signHostWrapper.apply {
        inherit pkgs provisionerPasswordFile;
        inherit (cfg.certificates) provisioner;
        inherit (cfg.certificates.host) extraPrincipals;
      }).wrapper;
      sshHostCertRenew = (wrappers.renewHostWrapper.apply {
        inherit pkgs;
        sshHostKeyFile = hostKeyFile;
        overwrite = true;
      }).wrapper;
      sshHostCertCheck = (wrappers.hostCheckWrapper.apply {
        inherit pkgs;
        certPath = hostCertFile;
      }).wrapper;
      sshHostRenewalCheck = (wrappers.renewalCheck.apply {
        inherit pkgs;
        certPath = hostCertFile;
        expires-in = "4h";
      }).wrapper;
      sshUserCertSign = (wrappers.signUserWrapper.apply {
        inherit pkgs provisionerPasswordFile;
        inherit (cfg.certificates) provisioner;
        validUsers = [ "root" "john" "appdaemon" ];
      }).wrapper;
      sshUserCertCheck = (wrappers.userCheckWrapper.apply {
        inherit pkgs;
        certPath = "${cfg.user.keyFile}-cert.pub";
      }).wrapper;
    in
    {
      options.ssh-new = {
        user = {
          keyFile = lib.mkOption {
            type = lib.types.str;
            default = "$HOME/.ssh/id_ed25519";
          };
        };
        host = {
          configDir = lib.mkOption {
            type = lib.types.str;
            default = "/etc/ssh";
          };
          keyFile = lib.mkOption {
            description = "String path to the host private key file";
            type = lib.types.str;
            default = "ssh_host_ed25519_key";
          };
          keyType = lib.mkOption {
            description = "OpenSSH host key type for ssh.hostKey.";
            type = lib.types.enum [ "ed25519" "rsa" "ecdsa" ];
            default = "ed25519";
          };
          extraSettings = lib.mkOption {
            description = "Extra settings to merge";
            type = lib.types.attrs;
            default = { };
          };
        };
        certificates = {
          provisioner = lib.mkOption {
            type = lib.types.nullOr lib.types.str;
          };
          host = {
            enable = lib.mkEnableOption "Enable SSH host certs";
            autoRenew = lib.mkEnableOption "Auto-renew the SSH host certs with a systemd service/timer";
            extraPrincipals = lib.mkOption {
              type = lib.types.listOf lib.types.str;
              default = [ ];
            };
          };
          user = {
            enable = lib.mkEnableOption "Enable SSH user certs";
            CAFile = lib.mkOption {
              description = "Filename of the SSH user CA with the config directory";
              type = lib.types.str;
              default = "ssh_user_ca_key.pub";
            };
            extraPrincipals = lib.mkOption {
              type = lib.types.listOf lib.types.str;
              default = [ ];
            };
          };
        };
      };

      config = {
        services.openssh = {
          enable = true;
          openFirewall = true;
          hostKeys = [
            {
              path = hostKeyFile;
              type = cfg.host.keyType;
            }
          ];
          settings = lib.mkMerge [
            {
              PasswordAuthentication = false;
              KbdInteractiveAuthentication = false;
            }
            cfg.host.extraSettings
            (lib.mkIf cfg.certificates.host.enable {
              HostCertificate = "${hostKeyFile}-cert.pub";
            })
            (lib.mkIf cfg.certificates.user.enable {
              TrustedUserCAKeys = "${cfg.host.configDir}/${cfg.certificates.user.CAFile}";
            })
          ];
        };

        environment.etc."ssh/${cfg.certificates.user.CAFile}" = lib.mkIf cfg.certificates.user.enable {
          source = ../hosts/janus/public/ssh_user_ca_key.pub;
        };

        environment.systemPackages = (
          lib.optionals cfg.certificates.host.enable [
            sshHostCertSign
            sshHostCertRenew
            sshHostRenewalCheck
            sshHostCertCheck
          ]
        );

        sops = lib.mkIf cfg.certificates.host.autoRenew {
          secrets."janus/admin_jwk" = {
            sopsFile = ../../keys/secrets.yaml;
            owner = "root";
            group = "root";
            mode = "0400";
          };
        };

        systemd = lib.mkIf cfg.certificates.host.autoRenew {
          services.ssh-certs-renew = {
            description = "SSH host certificate renewal";
            wantedBy = [ "multi-user.target" ];
            after = [ "network-online.target" ];
            wants = [ "network-online.target" ];
            path = with pkgs; [ step-cli systemd ];
            serviceConfig = {
              Type = "oneshot";
              User = "root";
              Group = "root";
              ExecCondition = lib.getExe sshHostRenewalCheck;
              ExecStart = lib.getExe sshHostCertRenew;
            };
          };
          timers.ssh-certs-renew = {
            description = "Periodic Step SSH host certificate renewal";
            wantedBy = [ "timers.target" ];
            timerConfig = {
              OnBootSec = "5m";
              OnUnitActiveSec = "4h";
              RandomizedDelaySec = "15m";
              Persistent = true;
              Unit = "ssh-certs-renew.service";
            };
          };
        };

        # This winds up in /etc/ssh/ssh_known_hosts, which applies to all users
        programs.ssh.knownHosts = lib.mkIf cfg.certificates.user.enable CAknownHosts;

        home-manager.users.root = lib.mkIf cfg.certificates.user.enable {
          home.stateVersion = lib.mkDefault config.system.stateVersion;
          imports = with inputs.self.modules.homeManager; [
            ssh-new
          ];
          home.packages = [
            sshUserCertSign
            sshUserCertCheck
          ];
          ssh-new.certificates.enable = true;
        };
      };
    };

  flake.modules.homeManager.ssh-new = { config, pkgs, lib, ... }:
    let
      cfg = config.ssh-new;
      sshHostCAContent = lib.removeSuffix "\n" (builtins.readFile ../hosts/janus/public/ssh_host_ca_key.pub);
      knownHostsText = lib.concatMapStrings
        (pattern: "@cert-authority ${pattern} ${sshHostCAContent}\n")
        caPatterns;
    in
    {
      options.ssh-new = {
        keyFile = lib.mkOption {
          type = lib.types.str;
          default = "${config.home.homeDirectory}/.ssh/id_ed25519";
        };
        certificates = {
          enable = lib.mkEnableOption "Enable SSH client certificates";
        };
      };

      config = {
        home.file.".ssh/known_hosts" = lib.mkIf cfg.certificates.enable {
          text = knownHostsText;
        };
        programs.ssh = {
          enable = true;
          enableDefaultConfig = false;
          settings = {
            # These settings apply to all connections
            "*" = lib.mkMerge (
              [
                # Default settings
                {
                  Compression = false;
                  ServerAliveInterval = 60;
                  ServerAliveCountMax = 3;
                  TCPKeepAlive = "yes";

                  PubkeyAuthentication = "yes";
                  PasswordAuthentication = "no";
                  PreferredAuthentications = "publickey";
                  IdentitiesOnly = true;
                  IdentityFile = cfg.keyFile;

                  StrictHostKeyChecking = "accept-new";
                  UserKnownHostsFile = "${config.home.homeDirectory}/.ssh/known_hosts";

                  HashKnownHosts = lib.mkDefault false;
                  AddKeysToAgent = lib.mkDefault "yes";
                  ForwardAgent = lib.mkDefault false;

                  RequestTTY = lib.mkDefault "auto";
                  SetEnv.TERM = "xterm-256color";
                }
              ]
              # SSH certificate settings
              ++ lib.optionals cfg.certificates.enable [
                {
                  CertificateFile = "${cfg.keyFile}-cert.pub";
                }
              ]
            );
          };
        };
      };
    };
}