anthropic-proxy/internal/proxy/sniff.go

package proxy

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/tidwall/gjson"
)

// SniffedProfile holds everything captured from a real Claude Code request.
// The proxy replays these verbatim — no hardcoded values needed.
type SniffedProfile struct {
	// Raw headers exactly as Claude Code sent them (name→value).
	// Excludes only host, content-length, and auth (we substitute our own token).
	Headers [][2]string

	// The full request body with system prompt, tools, metadata, thinking config, etc.
	// We swap out model + messages from the incoming client request.
	Body []byte

	// Parsed from User-Agent for billing header fingerprint computation.
	Version string
}

var skipHeaders = map[string]bool{
	"host":           true,
	"content-length": true,
	"authorization":  true,
	"x-api-key":      true,
	"connection":     true,
}

func SniffClaudeCode(claudeBinary string) (*SniffedProfile, error) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, fmt.Errorf("listen: %w", err)
	}
	port := listener.Addr().(*net.TCPAddr).Port

	var profile *SniffedProfile
	var mu sync.Mutex
	captured := make(chan struct{}, 1)

	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.Method == "HEAD" {
			w.WriteHeader(200)
			return
		}
		if r.Method != "POST" || !strings.Contains(r.URL.Path, "/v1/messages") {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(200)
			fmt.Fprint(w, `{"id":"msg_fake","type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"model":"claude-sonnet-4-6","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`)
			return
		}

		body, _ := io.ReadAll(r.Body)

		mu.Lock()
		if profile == nil {
			profile = extractProfile(r, body)
			select {
			case captured <- struct{}{}:
			default:
			}
		}
		mu.Unlock()

		if strings.Contains(string(body), `"stream":true`) {
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(200)
			fmt.Fprint(w, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_fake\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-sonnet-4-6\",\"stop_reason\":null,\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n")
			fmt.Fprint(w, "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n")
			fmt.Fprint(w, "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"ok\"}}\n\n")
			fmt.Fprint(w, "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n")
			fmt.Fprint(w, "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n")
			fmt.Fprint(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")
		} else {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(200)
			fmt.Fprint(w, `{"id":"msg_fake","type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"model":"claude-sonnet-4-6","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`)
		}
	})

	srv := &http.Server{Handler: mux}
	go srv.Serve(listener)
	defer srv.Close()

	cmd := exec.Command(claudeBinary, "--print", "say hi")
	cmd.Env = append(cmd.Environ(), fmt.Sprintf("ANTHROPIC_BASE_URL=http://127.0.0.1:%d", port))
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start claude: %w", err)
	}

	done := make(chan error, 1)
	go func() { done <- cmd.Wait() }()

	select {
	case <-captured:
		cmd.Process.Kill()
	case err := <-done:
		if err != nil && profile == nil {
			return nil, fmt.Errorf("claude exited: %w", err)
		}
	case <-time.After(30 * time.Second):
		cmd.Process.Kill()
		return nil, fmt.Errorf("sniff timed out after 30s")
	}

	if profile == nil {
		return nil, fmt.Errorf("no API request captured")
	}

	log.Printf("sniffed claude-code: version=%s headers=%d body=%d bytes",
		profile.Version, len(profile.Headers), len(profile.Body))
	return profile, nil
}

func extractProfile(r *http.Request, body []byte) *SniffedProfile {
	// Capture raw headers preserving original casing and order.
	var headers [][2]string
	for i := 0; i < len(r.Header); i++ {
		for name, vals := range r.Header {
			if skipHeaders[strings.ToLower(name)] {
				continue
			}
			for _, v := range vals {
				headers = append(headers, [2]string{name, v})
			}
		}
		break
	}

	// Deduplicate and strip subscription-specific betas.
	seen := map[string]bool{}
	var deduped [][2]string
	for _, h := range headers {
		key := strings.ToLower(h[0])
		if seen[key] {
			continue
		}
		seen[key] = true
		if key == "anthropic-beta" {
			var filtered []string
			for _, b := range strings.Split(h[1], ",") {
				if !strings.Contains(b, "context-1m") {
					filtered = append(filtered, b)
				}
			}
			h[1] = strings.Join(filtered, ",")
		}
		deduped = append(deduped, h)
	}

	ua := r.Header.Get("User-Agent")
	version := ""
	if i := strings.Index(ua, "/"); i > 0 {
		rest := ua[i+1:]
		if j := strings.IndexByte(rest, ' '); j > 0 {
			version = rest[:j]
		} else {
			version = rest
		}
	}

	// Extract the system prompt template from the body (everything except the billing header block).
	// The billing header is the first system block starting with "x-anthropic-billing-header:".
	systemBlocks := gjson.GetBytes(body, "system")
	var templateSystem []string
	if systemBlocks.IsArray() {
		for _, block := range systemBlocks.Array() {
			text := block.Get("text").String()
			if strings.HasPrefix(text, "x-anthropic-billing-header:") {
				continue
			}
			templateSystem = append(templateSystem, block.Raw)
		}
	}
	_ = templateSystem // stored in body for now

	return &SniffedProfile{
		Headers: deduped,
		Body:    body,
		Version: version,
	}
}