lightscale/nixos/modules/lightscale-client.nix

{ defaultPackage }:
{ config, lib, pkgs, ... }:
let
  cfg = config.services.lightscale-client;

  stateProfileDir = "${cfg.stateDir}/${cfg.profile}";

  baseArgs =
    [ "--profile" cfg.profile ]
    ++ lib.optionals (cfg.configFile != null) [ "--config" cfg.configFile ]
    ++ [ "--state-dir" stateProfileDir ]
    ++ lib.optionals (cfg.controlUrls != [ ]) [ "--control-url" (lib.concatStringsSep "," cfg.controlUrls) ]
    ++ lib.optionals (cfg.tlsPin != null) [ "--tls-pin" cfg.tlsPin ]
    ++ lib.optionals (cfg.adminToken != null) [ "--admin-token" cfg.adminToken ];

  agentArgs =
    [ "agent" "--listen-port" (toString cfg.listenPort) ]
    ++ lib.optionals (cfg.interface != null) [ "--interface" cfg.interface ]
    ++ lib.optionals cfg.applyRoutes [ "--apply-routes" ]
    ++ lib.optionals cfg.acceptExitNode [ "--accept-exit-node" ]
    ++ lib.optionals (cfg.exitNodeId != null) [ "--exit-node-id" cfg.exitNodeId ]
    ++ lib.optionals (cfg.exitNodeName != null) [ "--exit-node-name" cfg.exitNodeName ]
    ++ lib.optionals (cfg.exitNodeTag != null) [ "--exit-node-tag" cfg.exitNodeTag ]
    ++ [ "--exit-node-policy" cfg.exitNodePolicy ]
    ++ lib.optionals (cfg.exitNodeMetricBase != null) [ "--exit-node-metric-base" (toString cfg.exitNodeMetricBase) ]
    ++ lib.optionals (cfg.exitNodeUidRange != null) [ "--exit-node-uid-range" cfg.exitNodeUidRange ]
    ++ lib.optionals cfg.allowRouteConflicts [ "--allow-route-conflicts" ]
    ++ lib.optionals (cfg.routeTable != null) [ "--route-table" (toString cfg.routeTable) ]
    ++ lib.concatMap (endpoint: [ "--endpoint" endpoint ]) cfg.endpoints
    ++ lib.concatMap (prefix: [ "--advertise-route" prefix ]) cfg.advertiseRoutes
    ++ lib.concatMap (mapping: [ "--advertise-map" mapping ]) cfg.advertiseMaps
    ++ lib.optionals cfg.advertiseExitNode [ "--advertise-exit-node" ]
    ++ [ "--heartbeat-interval" (toString cfg.heartbeatInterval) ]
    ++ [ "--longpoll-timeout" (toString cfg.longpollTimeout) ]
    ++ [ "--backend" cfg.backend ]
    ++ lib.optionals cfg.stun [ "--stun" ]
    ++ lib.optionals (cfg.stunServers != [ ]) [ "--stun-server" (lib.concatStringsSep "," cfg.stunServers) ]
    ++ lib.optionals (cfg.stunPort != null) [ "--stun-port" (toString cfg.stunPort) ]
    ++ [ "--stun-timeout" (toString cfg.stunTimeout) ]
    ++ lib.optionals cfg.probePeers [ "--probe-peers" ]
    ++ [ "--probe-timeout" (toString cfg.probeTimeout) ]
    ++ lib.optionals cfg.streamRelay [ "--stream-relay" ]
    ++ lib.optionals (cfg.streamRelayServers != [ ]) [ "--stream-relay-server" (lib.concatStringsSep "," cfg.streamRelayServers) ]
    ++ [ "--endpoint-stale-after" (toString cfg.endpointStaleAfter) ]
    ++ [ "--endpoint-max-rotations" (toString cfg.endpointMaxRotations) ]
    ++ [ "--relay-reprobe-after" (toString cfg.relayReprobeAfter) ]
    ++ lib.optionals (cfg.dnsHostsPath != null) [ "--dns-hosts-path" cfg.dnsHostsPath ]
    ++ lib.optionals cfg.dnsServe [ "--dns-serve" ]
    ++ lib.optionals (cfg.dnsListen != null) [ "--dns-listen" cfg.dnsListen ]
    ++ lib.optionals cfg.dnsApplyResolver [ "--dns-apply-resolver" ]
    ++ lib.optionals cfg.l2Relay [ "--l2-relay" ]
    ++ cfg.extraArgs;

  args = baseArgs ++ agentArgs;
  startCmd = "${lib.getExe' cfg.package "lightscale-client"} ${lib.escapeShellArgs args}";

  registerArgs =
    baseArgs
    ++ [ "register" ]
    ++ lib.optionals (cfg.registerNodeName != null) [ "--node-name" cfg.registerNodeName ]
    ++ cfg.registerExtraArgs;

  registerScript = pkgs.writeShellScript "lightscale-client-register-${cfg.profile}" ''
    set -euo pipefail

    state_file=${lib.escapeShellArg "${stateProfileDir}/state.json"}
    if [[ -s "$state_file" ]]; then
      exit 0
    fi

    token_file=${lib.escapeShellArg cfg.enrollmentTokenFile}
    if [[ ! -r "$token_file" ]]; then
      echo "enrollment token file not readable: $token_file" >&2
      exit 1
    fi

    token="$(tr -d '\r\n' < "$token_file")"
    if [[ -z "$token" ]]; then
      echo "enrollment token is empty" >&2
      exit 1
    fi

    exec ${lib.getExe' cfg.package "lightscale-client"} ${lib.escapeShellArgs registerArgs} -- "$token"
  '';

  defaultUdpPorts = lib.unique ([ cfg.listenPort ] ++ cfg.firewallUDPPorts);
in
{
  options.services.lightscale-client = {
    enable = lib.mkEnableOption "lightscale client agent";

    package = lib.mkOption {
      type = lib.types.package;
      default = defaultPackage;
      description = "lightscale-client package.";
    };

    profile = lib.mkOption {
      type = lib.types.str;
      default = "default";
      description = "Profile name passed to lightscale-client.";
    };

    stateDir = lib.mkOption {
      type = lib.types.str;
      default = "/var/lib/lightscale-client";
      description = "Base state directory; per-profile state is stored under stateDir/profile.";
    };

    configFile = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Optional config file path for profiles and TLS pin.";
    };

    controlUrls = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Control plane URLs passed via --control-url (comma-separated failover supported).";
    };

    tlsPin = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Optional SHA-256 TLS pin (same format as `lightscale-client pin`).";
    };

    adminToken = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Optional admin token flag. Prefer environmentFiles for secrets.";
    };

    autoRegister = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Run a one-shot registration service when state file is missing.";
    };

    enrollmentTokenFile = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Path to enrollment token file for autoRegister.";
    };

    registerNodeName = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Optional node name used by autoRegister.";
    };

    registerExtraArgs = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Extra args appended to autoRegister `register` command.";
    };

    environmentFiles = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "systemd EnvironmentFile entries (for LIGHTSCALE_ADMIN_TOKEN, etc.).";
    };

    environment = lib.mkOption {
      type = lib.types.attrsOf lib.types.str;
      default = { };
      description = "Additional environment variables for the client service.";
    };

    interface = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "WireGuard interface name. Default is derived from profile.";
    };

    listenPort = lib.mkOption {
      type = lib.types.port;
      default = 51820;
      description = "WireGuard listen port for the agent.";
    };

    backend = lib.mkOption {
      type = lib.types.enum [ "kernel" "boringtun" ];
      default = "kernel";
      description = "WireGuard backend used by the agent.";
    };

    endpoints = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Advertised endpoints (HOST:PORT), passed as repeated --endpoint flags.";
    };

    applyRoutes = lib.mkOption {
      type = lib.types.bool;
      default = true;
      description = "Apply advertised subnet/exit routes on netmap updates.";
    };

    acceptExitNode = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Accept exit routes from selected peers.";
    };

    exitNodeId = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Specific exit node ID to use.";
    };

    exitNodeName = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Specific exit node name to use.";
    };

    exitNodePolicy = lib.mkOption {
      type = lib.types.enum [ "first" "latest" "multi" ];
      default = "first";
      description = "Policy for selecting exit nodes.";
    };

    exitNodeTag = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Only consider exit nodes with this tag.";
    };

    exitNodeMetricBase = lib.mkOption {
      type = lib.types.nullOr lib.types.ints.unsigned;
      default = null;
      description = "Base route metric used for selected exit routes.";
    };

    exitNodeUidRange = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "UID or UID range for per-user exit routing (e.g. 1000-1999).";
    };

    allowRouteConflicts = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Force route installation even when conflicts are detected.";
    };

    routeTable = lib.mkOption {
      type = lib.types.nullOr lib.types.ints.unsigned;
      default = null;
      description = "Optional policy routing table ID.";
    };

    advertiseRoutes = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Advertised subnet routes (CIDR).";
    };

    advertiseMaps = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Route map entries in REAL=MAPPED form.";
    };

    advertiseExitNode = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Advertise this node as an exit node.";
    };

    heartbeatInterval = lib.mkOption {
      type = lib.types.ints.positive;
      default = 30;
      description = "Heartbeat interval in seconds.";
    };

    longpollTimeout = lib.mkOption {
      type = lib.types.ints.positive;
      default = 30;
      description = "Netmap long-poll timeout in seconds.";
    };

    stun = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Enable STUN endpoint discovery in the agent.";
    };

    stunServers = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Explicit STUN servers (HOST:PORT).";
    };

    stunPort = lib.mkOption {
      type = lib.types.nullOr lib.types.port;
      default = null;
      description = "Optional local UDP port for STUN probe socket.";
    };

    stunTimeout = lib.mkOption {
      type = lib.types.ints.positive;
      default = 3;
      description = "STUN timeout in seconds.";
    };

    probePeers = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Send UDP probes to peers when netmap changes.";
    };

    probeTimeout = lib.mkOption {
      type = lib.types.ints.positive;
      default = 1;
      description = "UDP probe timeout in seconds.";
    };

    streamRelay = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Enable stream-relay endpoint fallback support.";
    };

    streamRelayServers = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Explicit stream relay servers (HOST:PORT).";
    };

    endpointStaleAfter = lib.mkOption {
      type = lib.types.ints.positive;
      default = 15;
      description = "Seconds before endpoint is considered stale.";
    };

    endpointMaxRotations = lib.mkOption {
      type = lib.types.ints.positive;
      default = 2;
      description = "Maximum endpoint rotations before relay fallback.";
    };

    relayReprobeAfter = lib.mkOption {
      type = lib.types.ints.positive;
      default = 60;
      description = "Seconds between direct re-probes while relay is active.";
    };

    dnsHostsPath = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "Optional hosts file path updated from netmap.";
    };

    dnsServe = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Serve MagicDNS-like answers for the overlay domain.";
    };

    dnsListen = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = "DNS listen address (defaults to 127.0.0.1:53 when dnsServe=true).";
    };

    dnsApplyResolver = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Apply per-domain resolver routing via resolvectl.";
    };

    l2Relay = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Enable experimental L2 relay behavior.";
    };

    extraArgs = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      description = "Additional CLI arguments appended to the agent command.";
    };

    openFirewall = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = "Open firewall ports listed in firewallTCPPorts and listen UDP port(s).";
    };

    firewallTCPPorts = lib.mkOption {
      type = lib.types.listOf lib.types.port;
      default = [ ];
      description = "Extra TCP ports to open when openFirewall=true.";
    };

    firewallUDPPorts = lib.mkOption {
      type = lib.types.listOf lib.types.port;
      default = [ ];
      description = "Extra UDP ports to open when openFirewall=true (listenPort is always included).";
    };
  };

  config = lib.mkIf cfg.enable {
    assertions = [
      {
        assertion = (!cfg.autoRegister) || (cfg.enrollmentTokenFile != null);
        message = "services.lightscale-client: autoRegister=true requires enrollmentTokenFile.";
      }
    ];

    systemd.services.lightscale-client-register = lib.mkIf cfg.autoRegister {
      description = "lightscale client one-shot registration (${cfg.profile})";
      wantedBy = [ "multi-user.target" ];
      before = [ "lightscale-client.service" ];
      after = [ "network-online.target" ];
      wants = [ "network-online.target" ];
      path = [ pkgs.coreutils ];
      unitConfig = {
        ConditionPathExists = "!${stateProfileDir}/state.json";
      };
      serviceConfig = {
        Type = "oneshot";
        ExecStart = registerScript;
        EnvironmentFile = cfg.environmentFiles;
      };
    };

    systemd.services.lightscale-client = {
      description = "lightscale client agent";
      wantedBy = [ "multi-user.target" ];
      after = [ "network-online.target" ] ++ lib.optionals cfg.autoRegister [ "lightscale-client-register.service" ];
      wants = [ "network-online.target" ] ++ lib.optionals cfg.autoRegister [ "lightscale-client-register.service" ];
      unitConfig = {
        ConditionPathExists = "${stateProfileDir}/state.json";
      };
      path = [
        pkgs.iproute2
        pkgs.nftables
        pkgs.systemd
      ];
      environment = cfg.environment;
      preStart = ''
        install -d -m 0750 ${lib.escapeShellArg cfg.stateDir}
        install -d -m 0750 ${lib.escapeShellArg stateProfileDir}
      '';
      serviceConfig = {
        Type = "simple";
        ExecStart = startCmd;
        Restart = "on-failure";
        RestartSec = 2;
        EnvironmentFile = cfg.environmentFiles;
        CapabilityBoundingSet = [
          "CAP_NET_ADMIN"
          "CAP_NET_RAW"
          "CAP_NET_BIND_SERVICE"
        ];
        AmbientCapabilities = [
          "CAP_NET_ADMIN"
          "CAP_NET_RAW"
          "CAP_NET_BIND_SERVICE"
        ];
      };
    };

    networking.firewall.allowedTCPPorts = lib.mkIf cfg.openFirewall cfg.firewallTCPPorts;
    networking.firewall.allowedUDPPorts = lib.mkIf cfg.openFirewall defaultUdpPorts;
  };
}