Merge pull request #65 from router-for-me/gemini-web

feat(auth): Improve Gemini web auth with email label detection

internal/cmd/gemini-web_auth.go

@@ -6,42 +6,116 @@ import (
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
+	"net/http"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/auth/gemini"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
 	sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
 	log "github.com/sirupsen/logrus"
 )
 
 // DoGeminiWebAuth handles the process of creating a Gemini Web token file.
-// It prompts the user for their cookie values and saves them to a JSON file.
+// New flow:
+//  1. Prompt user to paste the full cookie string.
+//  2. Extract __Secure-1PSID and __Secure-1PSIDTS from the cookie string.
+//  3. Call https://accounts.google.com/ListAccounts with the cookie to obtain email.
+//  4. Save auth file with the same structure, and set Label to the email.
 func DoGeminiWebAuth(cfg *config.Config) {
 	reader := bufio.NewReader(os.Stdin)
 
-	fmt.Print("Enter your __Secure-1PSID cookie value: ")
-	secure1psid, _ := reader.ReadString('\n')
-	secure1psid = strings.TrimSpace(secure1psid)
+	fmt.Print("Paste your full Google Cookie and press Enter: ")
+	rawCookie, _ := reader.ReadString('\n')
+	rawCookie = strings.TrimSpace(rawCookie)
+	if rawCookie == "" {
+		log.Fatal("Cookie cannot be empty")
+		return
+	}
 
+	// Parse K=V cookie pairs separated by ';'
+	cookieMap := make(map[string]string)
+	parts := strings.Split(rawCookie, ";")
+	for _, p := range parts {
+		p = strings.TrimSpace(p)
+		if p == "" {
+			continue
+		}
+		if eq := strings.Index(p, "="); eq > 0 {
+			k := strings.TrimSpace(p[:eq])
+			v := strings.TrimSpace(p[eq+1:])
+			if k != "" {
+				cookieMap[k] = v
+			}
+		}
+	}
+
+	secure1psid := strings.TrimSpace(cookieMap["__Secure-1PSID"])
+	secure1psidts := strings.TrimSpace(cookieMap["__Secure-1PSIDTS"])
+	// Fallback: prompt user to input missing values
 	if secure1psid == "" {
-		log.Fatal("The __Secure-1PSID value cannot be empty.")
-		return
+		fmt.Print("Cookie missing __Secure-1PSID. Enter __Secure-1PSID: ")
+		v, _ := reader.ReadString('\n')
+		secure1psid = strings.TrimSpace(v)
 	}
-
-	fmt.Print("Enter your __Secure-1PSIDTS cookie value: ")
-	secure1psidts, _ := reader.ReadString('\n')
-	secure1psidts = strings.TrimSpace(secure1psidts)
-
 	if secure1psidts == "" {
-		fmt.Println("The __Secure-1PSIDTS value cannot be empty.")
+		fmt.Print("Cookie missing __Secure-1PSIDTS. Enter __Secure-1PSIDTS: ")
+		v, _ := reader.ReadString('\n')
+		secure1psidts = strings.TrimSpace(v)
+	}
+	if secure1psid == "" || secure1psidts == "" {
+		log.Fatal("__Secure-1PSID and __Secure-1PSIDTS cannot be empty")
 		return
 	}
 
-	tokenStorage := &gemini.GeminiWebTokenStorage{
-		Secure1PSID:   secure1psid,
-		Secure1PSIDTS: secure1psidts,
+	// Build HTTP client with proxy settings respected.
+	httpClient := &http.Client{Timeout: 15 * time.Second}
+	httpClient = util.SetProxy(cfg, httpClient)
+
+	// Request ListAccounts to extract email as label (use POST per upstream behavior).
+	req, err := http.NewRequest(http.MethodPost, "https://accounts.google.com/ListAccounts", nil)
+	if err != nil {
+		fmt.Printf("Failed to create request: %v\n", err)
+		return
+	}
+	req.Header.Set("Cookie", rawCookie)
+	req.Header.Set("Accept", "application/json, text/plain, */*")
+	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36")
+	req.Header.Set("Origin", "https://accounts.google.com")
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded;charset=UTF-8")
+
+	resp, err := httpClient.Do(req)
+	email := ""
+	if err != nil {
+		fmt.Printf("Request to ListAccounts failed: %v\n", err)
+	} else {
+		defer resp.Body.Close()
+		if resp.StatusCode != http.StatusOK {
+			fmt.Printf("ListAccounts returned status code: %d\n", resp.StatusCode)
+		} else {
+			var payload []any
+			if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
+				fmt.Printf("Failed to parse ListAccounts response: %v\n", err)
+			} else {
+				// Expected structure like: ["gaia.l.a.r", [["gaia.l.a",1,"Name","email@example.com", ... ]]]
+				if len(payload) >= 2 {
+					if accounts, ok := payload[1].([]any); ok && len(accounts) >= 1 {
+						if first, ok := accounts[0].([]any); ok && len(first) >= 4 {
+							if em, ok := first[3].(string); ok {
+								email = strings.TrimSpace(em)
+							}
+						}
+					}
+				}
+				if email == "" {
+					fmt.Println("Failed to parse email from ListAccounts response")
+				}
+			}
+		}
 	}
 
 	// Generate a filename based on the SHA256 hash of the PSID
@@ -49,9 +123,25 @@ func DoGeminiWebAuth(cfg *config.Config) {
 	hasher.Write([]byte(secure1psid))
 	hash := hex.EncodeToString(hasher.Sum(nil))
 	fileName := fmt.Sprintf("gemini-web-%s.json", hash[:16])
-	// Set a stable label for logging, e.g. gemini-web-<hash>
-	if tokenStorage != nil {
-		tokenStorage.Label = strings.TrimSuffix(fileName, ".json")
+
+	// Decide label: prefer email; fallback prompt then file name without .json
+	defaultLabel := strings.TrimSuffix(fileName, ".json")
+	label := email
+	if label == "" {
+		fmt.Printf("Enter label for this auth (default: %s): ", defaultLabel)
+		v, _ := reader.ReadString('\n')
+		v = strings.TrimSpace(v)
+		if v != "" {
+			label = v
+		} else {
+			label = defaultLabel
+		}
+	}
+
+	tokenStorage := &gemini.GeminiWebTokenStorage{
+		Secure1PSID:   secure1psid,
+		Secure1PSIDTS: secure1psidts,
+		Label:         label,
 	}
 	record := &sdkAuth.TokenRecord{
 		Provider: "gemini-web",

internal/provider/gemini-web/client.go

@@ -9,8 +9,6 @@ import (
 	"net/http"
 	"net/http/cookiejar"
 	"net/url"
-	"os"
-	"path/filepath"
 	"regexp"
 	"strings"
 	"time"
@@ -126,19 +124,6 @@ func getAccessToken(baseCookies map[string]string, proxy string, verbose bool, i
 		}
 	}
 
-	cacheDir := "temp"
-	_ = os.MkdirAll(cacheDir, 0o755)
-	if v1, ok1 := baseCookies["__Secure-1PSID"]; ok1 {
-		cacheFile := filepath.Join(cacheDir, ".cached_1psidts_"+v1+".txt")
-		if b, err := os.ReadFile(cacheFile); err == nil {
-			cv := strings.TrimSpace(string(b))
-			if cv != "" {
-				merged := map[string]string{"__Secure-1PSID": v1, "__Secure-1PSIDTS": cv}
-				trySets = append(trySets, merged)
-			}
-		}
-	}
-
 	if len(extraCookies) > 0 {
 		trySets = append(trySets, extraCookies)
 	}
@@ -307,7 +292,7 @@ func (c *GeminiClient) Close(delaySec float64) {
 	c.Running = false
 }
 
-// ensureRunning mirrors the Python decorator behavior and retries on APIError.
+// ensureRunning mirrors the decorator behavior and retries on APIError.
 func (c *GeminiClient) ensureRunning() error {
 	if c.Running {
 		return nil
@@ -434,7 +419,7 @@ func (c *GeminiClient) generateOnce(prompt string, files []string, model Model,
 	}()
 
 	if resp.StatusCode == 429 {
-		// Surface 429 as TemporarilyBlocked to match Python behavior
+		// Surface 429 as TemporarilyBlocked to match reference behavior
 		c.Close(0)
 		return empty, &TemporarilyBlocked{GeminiError{Msg: "Too many requests. IP temporarily blocked."}}
 	}
@@ -514,7 +499,7 @@ func (c *GeminiClient) generateOnce(prompt string, files []string, model Model,
 				}
 			}
 		}
-		// Parse nested error code to align with Python mapping
+		// Parse nested error code to align with error mapping
 		var top []any
 		// Prefer lastTop from fallback scan; otherwise try parts[2]
 		if len(lastTop) > 0 {
@@ -537,7 +522,7 @@ func (c *GeminiClient) generateOnce(prompt string, files []string, model Model,
 			}
 		}
 		// Debug("Invalid response: control frames only; no body found")
-		// Close the client to force re-initialization on next request (parity with Python client behavior)
+		// Close the client to force re-initialization on next request (parity with reference client behavior)
 		c.Close(0)
 		return empty, &APIError{Msg: "Failed to generate contents. Invalid response data received."}
 	}
@@ -760,7 +745,7 @@ func (c *GeminiClient) generateOnce(prompt string, files []string, model Model,
 }
 
 // extractErrorCode attempts to navigate the known nested error structure and fetch the integer code.
-// Mirrors Python path: response_json[0][5][2][0][1][0]
+// Mirrors reference path: response_json[0][5][2][0][1][0]
 func extractErrorCode(top []any) (int, bool) {
 	if len(top) == 0 {
 		return 0, false

internal/provider/gemini-web/media.go

@@ -52,7 +52,7 @@ func (i Image) Save(path string, filename string, cookies map[string]string, ver
 			filename = q[0]
 		}
 	}
-	// Regex validation (align with Python: ^(.*\.\w+)) to extract name with extension.
+	// Regex validation (pattern: ^(.*\.\w+)) to extract name with extension.
 	if filename != "" {
 		re := regexp.MustCompile(`^(.*\.\w+)`)
 		if m := re.FindStringSubmatch(filename); len(m) >= 2 {
@@ -70,7 +70,7 @@ func (i Image) Save(path string, filename string, cookies map[string]string, ver
 	client := newHTTPClient(httpOptions{ProxyURL: i.Proxy, Insecure: insecure, FollowRedirects: true})
 	client.Timeout = 120 * time.Second
 
-	// Helper to set raw Cookie header using provided cookies (to mirror Python client behavior).
+	// Helper to set raw Cookie header using provided cookies (parity with the reference client behavior).
 	buildCookieHeader := func(m map[string]string) string {
 		if len(m) == 0 {
 			return ""

internal/provider/gemini-web/state.go

@@ -98,16 +98,16 @@ func (s *GeminiWebState) Label() string {
 }
 
 func (s *GeminiWebState) loadConversationCaches() {
-	if path := s.convPath(); path != "" {
-		if store, err := LoadConvStore(path); err == nil {
-			s.convStore = store
-		}
+	path := s.convPath()
+	if path == "" {
+		return
 	}
-	if path := s.convPath(); path != "" {
-		if items, index, err := LoadConvData(path); err == nil {
-			s.convData = items
-			s.convIndex = index
-		}
+	if store, err := LoadConvStore(path); err == nil {
+		s.convStore = store
+	}
+	if items, index, err := LoadConvData(path); err == nil {
+		s.convData = items
+		s.convIndex = index
 	}
 }