IRC-kosmi-relay/cmd/test-websocket-direct/main.go

package main

import (
	"context"
	"encoding/base64"
	"flag"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/gorilla/websocket"
)

const (
	userAgent  = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
	appVersion = "4364"
	wsURL      = "wss://engine.kosmi.io/gql-ws"
)

func main() {
	token := flag.String("token", "", "JWT token from captured session")
	roomID := flag.String("room", "@hyperspaceout", "Room ID")
	flag.Parse()

	if *token == "" {
		fmt.Fprintf(os.Stderr, "Error: -token is required\n")
		fmt.Fprintf(os.Stderr, "Usage: %s -token <JWT_TOKEN> -room <ROOM_ID>\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "\nTo get a token, run: ./capture-auth and extract it from auth-data.json\n")
		os.Exit(1)
	}

	fmt.Println("🔌 Testing direct WebSocket connection with JWT token")
	fmt.Printf("   Token: %s...\n", truncate(*token, 50))
	fmt.Printf("   Room: %s\n\n", *roomID)

	// Connect WebSocket
	fmt.Println("1️⃣  Connecting to WebSocket...")
	conn, err := connectWebSocket(*token)
	if err != nil {
		fmt.Fprintf(os.Stderr, "❌ Failed to connect: %v\n", err)
		os.Exit(1)
	}
	defer conn.Close()
	fmt.Println("✅ WebSocket connected!")

	// Wait for connection_ack
	fmt.Println("\n2️⃣  Waiting for connection_ack...")
	if err := waitForAck(conn); err != nil {
		fmt.Fprintf(os.Stderr, "❌ Failed to receive ack: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("✅ Received connection_ack!")

	// Subscribe to messages
	fmt.Printf("\n3️⃣  Subscribing to messages in room %s...\n", *roomID)
	if err := subscribeToMessages(conn, *roomID); err != nil {
		fmt.Fprintf(os.Stderr, "❌ Failed to subscribe: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("✅ Subscribed!")

	// Listen for messages
	fmt.Println("\n👂 Listening for messages (press Ctrl+C to exit)...\n")
	
	messageCount := 0
	for {
		var msg map[string]interface{}
		if err := conn.ReadJSON(&msg); err != nil {
			fmt.Fprintf(os.Stderr, "Error reading message: %v\n", err)
			break
		}

		msgType, _ := msg["type"].(string)
		
		switch msgType {
		case "next":
			payload, _ := msg["payload"].(map[string]interface{})
			data, _ := payload["data"].(map[string]interface{})
			
			if newMessage, ok := data["newMessage"].(map[string]interface{}); ok {
				messageCount++
				body, _ := newMessage["body"].(string)
				user, _ := newMessage["user"].(map[string]interface{})
				username, _ := user["displayName"].(string)
				if username == "" {
					username, _ = user["username"].(string)
				}
				timestamp, _ := newMessage["time"].(float64)
				
				t := time.Unix(int64(timestamp), 0)
				fmt.Printf("[%s] %s: %s\n", t.Format("15:04:05"), username, body)
			}
		case "complete":
			id, _ := msg["id"].(string)
			fmt.Printf("   [Subscription %s completed]\n", id)
		case "error":
			fmt.Printf("   ⚠️  Error: %+v\n", msg)
		default:
			// Ignore other message types (ka, etc)
		}
	}

	fmt.Printf("\n📊 Total messages received: %d\n", messageCount)
}

func connectWebSocket(token string) (*websocket.Conn, error) {
	dialer := websocket.Dialer{
		Subprotocols:    []string{"graphql-ws"},
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	
	headers := http.Header{
		"Origin":                []string{"https://app.kosmi.io"},
		"User-Agent":            []string{userAgent},
		"Sec-WebSocket-Protocol": []string{"graphql-ws"},
		"Cache-Control":         []string{"no-cache"},
		"Pragma":                []string{"no-cache"},
	}
	
	conn, resp, err := dialer.Dial(wsURL, headers)
	if err != nil {
		if resp != nil {
			fmt.Printf("   Response status: %d\n", resp.StatusCode)
		}
		return nil, err
	}

	// Send connection_init with token
	uaEncoded := base64.StdEncoding.EncodeToString([]byte(userAgent))
	
	initMsg := map[string]interface{}{
		"type": "connection_init",
		"payload": map[string]interface{}{
			"token": token,
			"ua":    uaEncoded,
			"v":     appVersion,
			"r":     "",
		},
	}

	if err := conn.WriteJSON(initMsg); err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to send connection_init: %w", err)
	}

	return conn, nil
}

func waitForAck(conn *websocket.Conn) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	done := make(chan error, 1)
	
	go func() {
		for {
			var msg map[string]interface{}
			if err := conn.ReadJSON(&msg); err != nil {
				done <- err
				return
			}

			msgType, _ := msg["type"].(string)
			fmt.Printf("   Received: %s\n", msgType)
			
			if msgType == "connection_ack" {
				done <- nil
				return
			}
			
			// Keep reading other messages (like ka)
		}
	}()

	select {
	case err := <-done:
		return err
	case <-ctx.Done():
		return fmt.Errorf("timeout waiting for connection_ack")
	}
}

func subscribeToMessages(conn *websocket.Conn, roomID string) error {
	query := fmt.Sprintf(`
		subscription {
			newMessage(roomId: "%s") {
				body
				time
				user {
					displayName
					username
				}
			}
		}
	`, roomID)

	subMsg := map[string]interface{}{
		"id":   "newMessage-subscription",
		"type": "subscribe",
		"payload": map[string]interface{}{
			"query":     query,
			"variables": map[string]interface{}{},
		},
	}

	return conn.WriteJSON(subMsg)
}

func truncate(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}