nixos-servers/system/pipe-to-loki.nix

{
  config,
  pkgs,
  lib,
  ...
}:
let
  pipe-to-loki = pkgs.writeShellApplication {
    name = "pipe-to-loki";
    runtimeInputs = with pkgs; [
      curl
      jq
      util-linux
      coreutils
    ];
    text = ''
      set -euo pipefail

      LOKI_URL="https://loki.home.2rjus.net/loki/api/v1/push"
      LOKI_AUTH_FILE="/run/secrets/promtail-loki-auth"
      HOSTNAME=$(hostname)
      SESSION_ID=""
      RECORD_MODE=false

      usage() {
        echo "Usage: pipe-to-loki [--id ID] [--record]"
        echo ""
        echo "Send command output or interactive sessions to Loki."
        echo ""
        echo "Options:"
        echo "  --id ID      Set custom session ID (default: auto-generated)"
        echo "  --record     Start interactive recording session"
        echo ""
        echo "Examples:"
        echo "  command | pipe-to-loki           # Pipe command output"
        echo "  command | pipe-to-loki --id foo  # Pipe with custom ID"
        echo "  pipe-to-loki --record            # Start recording session"
        exit 1
      }

      generate_id() {
        local random_chars
        random_chars=$(head -c 2 /dev/urandom | od -An -tx1 | tr -d ' \n')
        echo "''${HOSTNAME}-$(date +%s)-''${random_chars}"
      }

      send_to_loki() {
        local content="$1"
        local type="$2"
        local timestamp_ns
        timestamp_ns=$(date +%s%N)

        local payload
        payload=$(jq -n \
          --arg job "pipe-to-loki" \
          --arg host "$HOSTNAME" \
          --arg type "$type" \
          --arg id "$SESSION_ID" \
          --arg ts "$timestamp_ns" \
          --arg content "$content" \
          '{
            streams: [{
              stream: {
                job: $job,
                hostname: $host,
                type: $type,
                id: $id
              },
              values: [[$ts, $content]]
            }]
          }')

        local auth_args=()
        if [[ -f "$LOKI_AUTH_FILE" ]]; then
          auth_args=(-u "promtail:$(cat "$LOKI_AUTH_FILE")")
        fi

        if curl -s -X POST "$LOKI_URL" \
          "''${auth_args[@]}" \
          -H "Content-Type: application/json" \
          -d "$payload" > /dev/null; then
          return 0
        else
          echo "Error: Failed to send to Loki" >&2
          return 1
        fi
      }

      # Parse arguments
      while [[ $# -gt 0 ]]; do
        case $1 in
          --id)
            SESSION_ID="$2"
            shift 2
            ;;
          --record)
            RECORD_MODE=true
            shift
            ;;
          --help|-h)
            usage
            ;;
          *)
            echo "Unknown option: $1" >&2
            usage
            ;;
        esac
      done

      # Generate ID if not provided
      if [[ -z "$SESSION_ID" ]]; then
        SESSION_ID=$(generate_id)
      fi

      if $RECORD_MODE; then
        # Session recording mode
        SCRIPT_FILE=$(mktemp)
        trap 'rm -f "$SCRIPT_FILE"' EXIT

        echo "Recording session $SESSION_ID... (exit to send)"

        # Use script to record the session
        script -q "$SCRIPT_FILE"

        # Read the transcript and send to Loki
        content=$(cat "$SCRIPT_FILE")
        if send_to_loki "$content" "session"; then
          echo "Session $SESSION_ID sent to Loki"
        fi
      else
        # Pipe mode - read from stdin
        if [[ -t 0 ]]; then
          echo "Error: No input provided. Pipe a command or use --record for interactive mode." >&2
          exit 1
        fi

        content=$(cat)
        if send_to_loki "$content" "command"; then
          echo "Sent to Loki with id: $SESSION_ID"
        fi
      fi
    '';
  };
in
{
  environment.systemPackages = [ pipe-to-loki ];
}