ssh.nix

{ lib, config, ... }:
let
  inherit (lib) types;

  /*
  connections :: ListOf { from.host; from.user; to.host; to.user }

  A list of SSH connections that should be allowed as given by the ssh.access option
  */
  connections = lib.flatten (lib.mapAttrsToList (fromHost: fromHostValue:
    lib.mapAttrsToList (fromKey: fromUserValue:
      lib.mapAttrsToList (toHost: toHostValue:
        lib.mapAttrsToList (toUser: toUserValue:
          lib.optional toUserValue {
            from.host = fromHost;
            from.key = fromKey;
            to.host = toHost;
            to.user = toUser;
          }
        ) toHostValue
      ) fromUserValue.hasAccessTo
    ) fromHostValue.keys
  ) config.ssh.access);

  userConfig = lib.mapAttrs (host: hostConnections: {
    configuration.users.users = lib.mapAttrs (user: userConnections: {
      openssh.authorizedKeys.keys = lib.mkIf (userConnections != []) (map (conn:
        config.ssh.access.${conn.from.host}.keys.${conn.from.key}.publicKey
      ) userConnections);
    }) (lib.groupBy (conn: conn.to.user) hostConnections);
  }) (lib.groupBy (conn: conn.to.host) connections);

  knownHostsConfig = lib.mapAttrs (fromHost: fromHostConnections: {
    configuration.programs.ssh.knownHosts = lib.mkMerge (lib.mapAttrsToList (toHost: toHostConnections:
      if config.ssh.access.${toHost}.hostKeys == {}
      then throw "No host keys defined with ssh.access.${toHost}.hostKeys, but we need one to give secure access from ${fromHost}."
      else lib.mapAttrs' (hostKeyName: publicKey: {
        name = "${toHost}-${hostKeyName}";
        value = {
          hostNames = [ toHost ] ++ config.ssh.access.${toHost}.hostNames ++ lib.optional (fromHost == toHost) "localhost";
          inherit publicKey;
        };
      }) config.ssh.access.${toHost}.hostKeys
    ) (lib.groupBy (conn: conn.to.host) fromHostConnections));
  }) (lib.groupBy (conn: conn.from.host) connections);


in {

  options.ssh = {

    access = lib.mkOption {
      description = ''
        A specification for which host/key pair should have access to which
        other host/user pair. An entry here essentially makes `ssh user@host`
        work smoothly.

        This works by adding the source key to the target users authorized
        keys and by adding the target host key to the source hosts known hosts.

        Note that hosts are specified by Nixus node name.
      '';
      example = lib.literalExample ''
        {
          sourceHost = {
            keys.sourceKey = {
              # Generate this with ssh-keygen
              publicKey = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDOJYDb9isFfFog88Lzvs1CEfAmcVB7F9NUFzC7XXXXX";
              hasAccessTo = {
                # This allows you to `ssh targetUser@targetHost`
                # from a user having sourceKey on sourceHost
                targetHost.targetUser = true;
              };
            };
          };

          # The target hosts key needs to be specified
          targetHost = {
            # Usually autogenerated in /etc/ssh/ssh_host_*_key.pub
            hostKeys.ed25519 = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIElTpklo4GfIwPG2/HxUzsov9eW7Z0au0hF9HAzXXXXX";
          };
        }
      '';
      default = {};
      type = types.attrsOf (types.submodule {

        options.hostKeys = lib.mkOption {
          description = ''
            The host keys, usually autogenerated in /etc/ssh/ssh_host_*_key.pub.
            Generally only a single key is needed. The attribute name can be
            arbitrary and doesn't have any effect.

            This needs to be specified if any host/user needs to have SSH
            access to this host
          '';
          default = {};
          type = types.attrsOf types.str;
        };

        options.hostNames = lib.mkOption {
          description = ''
            The host names this host is reachable from. Among others, this can include
            - local subnet IP addresses
            - public IP addresses
            - VPN IP addresses
            - DNS domains

            Note that the node name itself and localhost (if applicable) are
            implicitly in this list.
          '';
          type = types.listOf types.str;
          default = [];
        };

        options.keys = lib.mkOption {
          description = ''
            Which keys this host has and the access they should have to which
            hosts. The attribute name can be arbitrary and has no effect on the
            result.
          '';
          default = {};
          type = types.attrsOf (types.submodule {
            options.publicKey = lib.mkOption {
              description = ''
                The public key. This can be generated with `ssh-keygen`
                or `ssh-keygen -t ed25519`. This needs to be specified if the
                keys owner wants to SSH into another host.

                If you have multiple keys for the same user, specify another
                attribute for the `keys` option.
              '';
              type = types.str;
            };
            options.hasAccessTo = lib.mkOption {
              description = ''
                Which host/user this key should have access to. A value
                of `<host>.<user> = true` allows `ssh <user>@<host>` to work.
              '';
              example = lib.literalExample ''
                {
                  targetHost.targetUser = true;
                }
              '';
              type = types.attrsOf (types.attrsOf types.bool);
              default = {};
            };
          });
        };

      });
    };

  };

  config = lib.mkMerge [
    { nodes = userConfig; }
    { nodes = knownHostsConfig; }
  ];

}