chatgpt-codex-proxy/main.go at main · minikomi/chatgpt-codex-proxy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
package main

import (
	"bufio"
	"context"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

func runAuth() error {
	pkce, err := generatePKCE()
	if err != nil {
		return err
	}
	state := generateState()
	authURL := buildAuthURL(pkce, state)

	// Try to start callback server on port 1455
	codeCh := make(chan string, 1)
	srv := &http.Server{Addr: ":1455"}

	mux := http.NewServeMux()
	mux.HandleFunc("/auth/callback", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Query().Get("state") != state {
			http.Error(w, "State mismatch", 400)
			return
		}
		code := r.URL.Query().Get("code")
		if code == "" {
			http.Error(w, "No code", 400)
			return
		}
		w.Header().Set("Content-Type", "text/html")
		fmt.Fprintln(w, "<html><body><h1>✓ Success!</h1><p>You can close this window and return to the terminal.</p></body></html>")
		codeCh <- code
	})
	srv.Handler = mux

	go srv.ListenAndServe()

	// Give server a moment to start
	time.Sleep(100 * time.Millisecond)

	fmt.Println("=== ChatGPT OAuth Authentication ===")
	fmt.Println()
	fmt.Println("Open this URL in your browser:")
	fmt.Println()
	fmt.Println(authURL)
	fmt.Println()
	fmt.Println("Waiting for browser callback on localhost:1455...")
	fmt.Println("(Or paste the redirect URL here if callback fails)")
	fmt.Println()

	inputCh := make(chan string, 1)
	go func() {
		reader := bufio.NewReader(os.Stdin)
		input, _ := reader.ReadString('\n')
		inputCh <- strings.TrimSpace(input)
	}()

	var code string
	select {
	case code = <-codeCh:
		fmt.Println("Got callback!")
	case input := <-inputCh:
		code = extractCode(input)
		if code == "" {
			return fmt.Errorf("could not extract code from: %s", input)
		}
	case <-time.After(5 * time.Minute):
		return fmt.Errorf("timeout waiting for authentication")
	}

	srv.Shutdown(context.Background())

	fmt.Println("Exchanging code for tokens...")
	tokens, err := exchangeCode(code, pkce.Verifier)
	if err != nil {
		return err
	}

	if err := saveTokens(tokens); err != nil {
		return err
	}

	fmt.Println()
	fmt.Println("Success! Tokens saved to:", tokensPath())
	return nil
}

func extractCode(input string) string {
	if u, err := url.Parse(input); err == nil && u.Query().Get("code") != "" {
		return u.Query().Get("code")
	}
	if len(input) > 10 && !strings.Contains(input, " ") {
		return input
	}
	return ""
}

func main() {
	addr := ":8080"

	if len(os.Args) > 1 {
		switch os.Args[1] {
		case "auth":
			if err := runAuth(); err != nil {
				log.Fatal(err)
			}
			return
		case "serve":
			if len(os.Args) > 2 {
				addr = os.Args[2]
			}
		default:
			addr = os.Args[1]
		}
	}

	// Check if we have tokens, if not run auth
	if _, err := loadTokens(); err != nil {
		fmt.Println("No tokens found. Starting authentication...")
		fmt.Println()
		if err := runAuth(); err != nil {
			log.Fatal(err)
		}
		fmt.Println()
	}

	if err := startProxy(addr); err != nil {
		log.Fatal(err)
	}
}