e916d5942b
All checks were successful
Build & Sign Wraith / Build Windows + Sign (push) Successful in 1m5s
Reads %USERPROFILE%\.claude\.credentials.json (or ~/.claude/.credentials.json), extracts the access and refresh tokens, stores them encrypted in Wraith's vault. Works when Wraith's own OAuth exchange fails. If Claude Code is authenticated on the same machine, Wraith piggybacks on the existing token. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
474 lines
13 KiB
Go
474 lines
13 KiB
Go
package ai
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/vstockwell/wraith/internal/settings"
|
|
"github.com/vstockwell/wraith/internal/vault"
|
|
)
|
|
|
|
const (
|
|
oauthClientID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
|
|
oauthAuthorizeURL = "https://claude.ai/oauth/authorize"
|
|
oauthTokenURL = "https://platform.claude.com/v1/oauth/token"
|
|
)
|
|
|
|
// OAuthManager handles OAuth PKCE authentication for Claude Max subscriptions.
|
|
type OAuthManager struct {
|
|
settings *settings.SettingsService
|
|
vault *vault.VaultService
|
|
clientID string
|
|
authorizeURL string
|
|
tokenURL string
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// tokenResponse is the JSON body returned by the token endpoint.
|
|
type tokenResponse struct {
|
|
AccessToken string `json:"access_token"`
|
|
RefreshToken string `json:"refresh_token"`
|
|
ExpiresIn int `json:"expires_in"`
|
|
TokenType string `json:"token_type"`
|
|
}
|
|
|
|
// NewOAuthManager creates an OAuthManager wired to settings and vault services.
|
|
func NewOAuthManager(s *settings.SettingsService, v *vault.VaultService) *OAuthManager {
|
|
return &OAuthManager{
|
|
settings: s,
|
|
vault: v,
|
|
clientID: oauthClientID,
|
|
authorizeURL: oauthAuthorizeURL,
|
|
tokenURL: oauthTokenURL,
|
|
}
|
|
}
|
|
|
|
// StartLogin begins the OAuth PKCE flow. It starts a local HTTP server,
|
|
// opens the authorization URL in the user's browser, and returns a channel
|
|
// that receives nil on success or an error on failure.
|
|
func (o *OAuthManager) StartLogin(openURL func(string) error) (<-chan error, error) {
|
|
verifier, err := generateCodeVerifier()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("generate code verifier: %w", err)
|
|
}
|
|
challenge := generateCodeChallenge(verifier)
|
|
state, err := generateState()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("generate state: %w", err)
|
|
}
|
|
|
|
// Start a local HTTP server on a random port for the callback
|
|
listener, err := net.Listen("tcp", "localhost:0")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("listen for callback: %w", err)
|
|
}
|
|
port := listener.Addr().(*net.TCPAddr).Port
|
|
redirectURI := fmt.Sprintf("http://localhost:%d/callback", port)
|
|
|
|
done := make(chan error, 1)
|
|
|
|
mux := http.NewServeMux()
|
|
server := &http.Server{Handler: mux}
|
|
|
|
mux.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) {
|
|
defer func() {
|
|
go func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
server.Shutdown(ctx)
|
|
}()
|
|
}()
|
|
|
|
// Verify state
|
|
if r.URL.Query().Get("state") != state {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
fmt.Fprint(w, "State mismatch — please try logging in again.")
|
|
done <- fmt.Errorf("state mismatch")
|
|
return
|
|
}
|
|
|
|
if errParam := r.URL.Query().Get("error"); errParam != "" {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
fmt.Fprintf(w, "Authorization error: %s", errParam)
|
|
done <- fmt.Errorf("authorization error: %s", errParam)
|
|
return
|
|
}
|
|
|
|
code := r.URL.Query().Get("code")
|
|
if code == "" {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
fmt.Fprint(w, "No authorization code received.")
|
|
done <- fmt.Errorf("no authorization code")
|
|
return
|
|
}
|
|
|
|
// Exchange code for tokens
|
|
if err := o.exchangeCode(code, verifier, redirectURI); err != nil {
|
|
slog.Error("oauth token exchange failed", "error", err)
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
fmt.Fprintf(w, "<html><body><h2>Authentication Failed</h2><pre>%s</pre><p>Check Wraith logs for details.</p></body></html>", err.Error())
|
|
done <- fmt.Errorf("exchange code: %w", err)
|
|
return
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "text/html")
|
|
fmt.Fprint(w, `<html><body><h2>Authenticated!</h2><p>You can close this window and return to Wraith.</p></body></html>`)
|
|
done <- nil
|
|
})
|
|
|
|
go func() {
|
|
if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
|
|
slog.Error("oauth callback server error", "error", err)
|
|
done <- fmt.Errorf("callback server: %w", err)
|
|
}
|
|
}()
|
|
|
|
// Build the authorize URL
|
|
params := url.Values{
|
|
"code": {"true"},
|
|
"response_type": {"code"},
|
|
"client_id": {o.clientID},
|
|
"redirect_uri": {redirectURI},
|
|
"scope": {"org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers"},
|
|
"state": {state},
|
|
"code_challenge": {challenge},
|
|
"code_challenge_method": {"S256"},
|
|
}
|
|
authURL := o.authorizeURL + "?" + params.Encode()
|
|
|
|
if openURL != nil {
|
|
if err := openURL(authURL); err != nil {
|
|
listener.Close()
|
|
return nil, fmt.Errorf("open browser: %w", err)
|
|
}
|
|
}
|
|
|
|
return done, nil
|
|
}
|
|
|
|
// exchangeCode exchanges an authorization code for access and refresh tokens.
|
|
func (o *OAuthManager) exchangeCode(code, verifier, redirectURI string) error {
|
|
// Try JSON format first (what Claude Code appears to use), fall back to form-encoded
|
|
payload := map[string]string{
|
|
"grant_type": "authorization_code",
|
|
"code": code,
|
|
"redirect_uri": redirectURI,
|
|
"client_id": o.clientID,
|
|
"code_verifier": verifier,
|
|
}
|
|
|
|
jsonBody, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return fmt.Errorf("marshal token request: %w", err)
|
|
}
|
|
|
|
slog.Info("exchanging auth code", "tokenURL", o.tokenURL, "redirectURI", redirectURI)
|
|
|
|
req, err := http.NewRequest("POST", o.tokenURL, io.NopCloser(
|
|
bytes.NewReader(jsonBody),
|
|
))
|
|
if err != nil {
|
|
return fmt.Errorf("create token request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
return fmt.Errorf("post token request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return fmt.Errorf("read token response: %w", err)
|
|
}
|
|
|
|
slog.Info("token endpoint response", "status", resp.StatusCode, "body", string(body)[:min(len(body), 500)])
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
// If JSON failed, try form-encoded
|
|
slog.Info("JSON token exchange failed, trying form-encoded")
|
|
data := url.Values{
|
|
"grant_type": {"authorization_code"},
|
|
"code": {code},
|
|
"redirect_uri": {redirectURI},
|
|
"client_id": {o.clientID},
|
|
"code_verifier": {verifier},
|
|
}
|
|
resp2, err := http.PostForm(o.tokenURL, data)
|
|
if err != nil {
|
|
return fmt.Errorf("post form token request: %w", err)
|
|
}
|
|
defer resp2.Body.Close()
|
|
body, err = io.ReadAll(resp2.Body)
|
|
if err != nil {
|
|
return fmt.Errorf("read form token response: %w", err)
|
|
}
|
|
slog.Info("form-encoded token response", "status", resp2.StatusCode, "body", string(body)[:min(len(body), 500)])
|
|
if resp2.StatusCode != http.StatusOK {
|
|
return fmt.Errorf("token endpoint returned %d (json) and %d (form): %s", resp.StatusCode, resp2.StatusCode, string(body))
|
|
}
|
|
}
|
|
|
|
var tokens tokenResponse
|
|
if err := json.Unmarshal(body, &tokens); err != nil {
|
|
return fmt.Errorf("unmarshal token response: %w", err)
|
|
}
|
|
|
|
return o.storeTokens(tokens)
|
|
}
|
|
|
|
func min(a, b int) int {
|
|
if a < b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|
|
|
|
// storeTokens encrypts and persists OAuth tokens in settings.
|
|
func (o *OAuthManager) storeTokens(tokens tokenResponse) error {
|
|
o.mu.Lock()
|
|
defer o.mu.Unlock()
|
|
|
|
if o.vault == nil {
|
|
return fmt.Errorf("vault is not unlocked — cannot store tokens")
|
|
}
|
|
|
|
encAccess, err := o.vault.Encrypt(tokens.AccessToken)
|
|
if err != nil {
|
|
return fmt.Errorf("encrypt access token: %w", err)
|
|
}
|
|
if err := o.settings.Set("ai_access_token", encAccess); err != nil {
|
|
return fmt.Errorf("store access token: %w", err)
|
|
}
|
|
|
|
encRefresh, err := o.vault.Encrypt(tokens.RefreshToken)
|
|
if err != nil {
|
|
return fmt.Errorf("encrypt refresh token: %w", err)
|
|
}
|
|
if err := o.settings.Set("ai_refresh_token", encRefresh); err != nil {
|
|
return fmt.Errorf("store refresh token: %w", err)
|
|
}
|
|
|
|
expiresAt := time.Now().Add(time.Duration(tokens.ExpiresIn) * time.Second).Unix()
|
|
if err := o.settings.Set("ai_token_expires_at", strconv.FormatInt(expiresAt, 10)); err != nil {
|
|
return fmt.Errorf("store token expiry: %w", err)
|
|
}
|
|
|
|
slog.Info("oauth tokens stored successfully")
|
|
return nil
|
|
}
|
|
|
|
// GetAccessToken returns a valid access token, refreshing if expired.
|
|
func (o *OAuthManager) GetAccessToken() (string, error) {
|
|
o.mu.RLock()
|
|
vlt := o.vault
|
|
o.mu.RUnlock()
|
|
|
|
if vlt == nil {
|
|
return "", fmt.Errorf("vault is not unlocked")
|
|
}
|
|
|
|
// Check expiry
|
|
expiresStr, _ := o.settings.Get("ai_token_expires_at")
|
|
if expiresStr != "" {
|
|
expiresAt, _ := strconv.ParseInt(expiresStr, 10, 64)
|
|
if time.Now().Unix() >= expiresAt {
|
|
// Token expired, try to refresh
|
|
if err := o.refreshToken(); err != nil {
|
|
return "", fmt.Errorf("refresh expired token: %w", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
encAccess, err := o.settings.Get("ai_access_token")
|
|
if err != nil || encAccess == "" {
|
|
return "", fmt.Errorf("no access token stored")
|
|
}
|
|
|
|
accessToken, err := vlt.Decrypt(encAccess)
|
|
if err != nil {
|
|
return "", fmt.Errorf("decrypt access token: %w", err)
|
|
}
|
|
|
|
return accessToken, nil
|
|
}
|
|
|
|
// refreshToken uses the stored refresh token to obtain new tokens.
|
|
func (o *OAuthManager) refreshToken() error {
|
|
o.mu.RLock()
|
|
vlt := o.vault
|
|
o.mu.RUnlock()
|
|
|
|
if vlt == nil {
|
|
return fmt.Errorf("vault is not unlocked")
|
|
}
|
|
|
|
encRefresh, err := o.settings.Get("ai_refresh_token")
|
|
if err != nil || encRefresh == "" {
|
|
return fmt.Errorf("no refresh token stored")
|
|
}
|
|
|
|
refreshTok, err := vlt.Decrypt(encRefresh)
|
|
if err != nil {
|
|
return fmt.Errorf("decrypt refresh token: %w", err)
|
|
}
|
|
|
|
data := url.Values{
|
|
"grant_type": {"refresh_token"},
|
|
"refresh_token": {refreshTok},
|
|
"client_id": {o.clientID},
|
|
}
|
|
|
|
resp, err := http.PostForm(o.tokenURL, data)
|
|
if err != nil {
|
|
return fmt.Errorf("post refresh request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return fmt.Errorf("read refresh response: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return fmt.Errorf("token refresh returned %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
var tokens tokenResponse
|
|
if err := json.Unmarshal(body, &tokens); err != nil {
|
|
return fmt.Errorf("unmarshal refresh response: %w", err)
|
|
}
|
|
|
|
// If the refresh endpoint doesn't return a new refresh token, keep the old one
|
|
if tokens.RefreshToken == "" {
|
|
tokens.RefreshToken = refreshTok
|
|
}
|
|
|
|
return o.storeTokens(tokens)
|
|
}
|
|
|
|
// IsAuthenticated checks if we have stored tokens.
|
|
func (o *OAuthManager) IsAuthenticated() bool {
|
|
encAccess, err := o.settings.Get("ai_access_token")
|
|
return err == nil && encAccess != ""
|
|
}
|
|
|
|
// Logout clears all stored OAuth tokens.
|
|
func (o *OAuthManager) Logout() error {
|
|
for _, key := range []string{"ai_access_token", "ai_refresh_token", "ai_token_expires_at"} {
|
|
if err := o.settings.Delete(key); err != nil {
|
|
return fmt.Errorf("delete %s: %w", key, err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// generateCodeVerifier creates a PKCE code verifier (32 random bytes, base64url no padding).
|
|
func generateCodeVerifier() (string, error) {
|
|
b := make([]byte, 32)
|
|
if _, err := rand.Read(b); err != nil {
|
|
return "", err
|
|
}
|
|
return base64.RawURLEncoding.EncodeToString(b), nil
|
|
}
|
|
|
|
// generateCodeChallenge creates a PKCE S256 code challenge from a verifier.
|
|
func generateCodeChallenge(verifier string) string {
|
|
h := sha256.Sum256([]byte(verifier))
|
|
return base64.RawURLEncoding.EncodeToString(h[:])
|
|
}
|
|
|
|
// generateState creates a random state parameter (32 random bytes, base64url no padding).
|
|
func generateState() (string, error) {
|
|
b := make([]byte, 32)
|
|
if _, err := rand.Read(b); err != nil {
|
|
return "", err
|
|
}
|
|
return base64.RawURLEncoding.EncodeToString(b), nil
|
|
}
|
|
|
|
// ImportClaudeCodeToken reads the Claude Code credentials.json file and imports
|
|
// the OAuth tokens. This is the fallback when Wraith's own OAuth flow fails.
|
|
// Path: %USERPROFILE%\.claude\.credentials.json (Windows) or ~/.claude/.credentials.json (macOS/Linux)
|
|
func (o *OAuthManager) ImportClaudeCodeToken() error {
|
|
home := os.Getenv("USERPROFILE")
|
|
if home == "" {
|
|
var err error
|
|
home, err = os.UserHomeDir()
|
|
if err != nil {
|
|
return fmt.Errorf("find home directory: %w", err)
|
|
}
|
|
}
|
|
|
|
credPath := filepath.Join(home, ".claude", ".credentials.json")
|
|
data, err := os.ReadFile(credPath)
|
|
if err != nil {
|
|
return fmt.Errorf("read credentials.json at %s: %w", credPath, err)
|
|
}
|
|
|
|
var creds struct {
|
|
ClaudeAiOAuth struct {
|
|
AccessToken string `json:"accessToken"`
|
|
RefreshToken string `json:"refreshToken"`
|
|
ExpiresAt int64 `json:"expiresAt"`
|
|
} `json:"claudeAiOauth"`
|
|
}
|
|
|
|
if err := json.Unmarshal(data, &creds); err != nil {
|
|
return fmt.Errorf("parse credentials.json: %w", err)
|
|
}
|
|
|
|
if creds.ClaudeAiOAuth.AccessToken == "" {
|
|
return fmt.Errorf("no access token found in credentials.json")
|
|
}
|
|
|
|
// Convert millisecond timestamp to seconds-based ExpiresIn
|
|
expiresIn := int((creds.ClaudeAiOAuth.ExpiresAt / 1000) - time.Now().Unix())
|
|
if expiresIn < 0 {
|
|
expiresIn = 0
|
|
}
|
|
|
|
tokens := tokenResponse{
|
|
AccessToken: creds.ClaudeAiOAuth.AccessToken,
|
|
RefreshToken: creds.ClaudeAiOAuth.RefreshToken,
|
|
ExpiresIn: expiresIn,
|
|
}
|
|
|
|
slog.Info("imported Claude Code token", "expiresInSeconds", expiresIn)
|
|
return o.storeTokens(tokens)
|
|
}
|
|
|
|
// SetVault updates the vault reference (called after vault unlock).
|
|
func (o *OAuthManager) SetVault(v *vault.VaultService) {
|
|
o.mu.Lock()
|
|
defer o.mu.Unlock()
|
|
o.vault = v
|
|
}
|
|
|
|
// isBase64URL checks if a string contains only base64url characters (no padding).
|
|
func isBase64URL(s string) bool {
|
|
for _, c := range s {
|
|
if !((c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-' || c == '_') {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|