OpenAI Key Checker

main.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/fatih/color"
	"github.com/valyala/fasthttp"
)

// Configuration
const (
	defaultWorkers = 50
	openAIEndpoint = "https://api.openai.com/v1/chat/completions"
)

// Result represents the outcome of checking a key
type Result struct {
	Key      string
	IsValid  bool
	HasQuota bool
	Message  string
}

func main() {
	// Parse command line flags
	filePath := flag.String("file", "", "Path to file containing OpenAI keys (one per line)")
	workers := flag.Int("workers", defaultWorkers, "Number of concurrent workers")
	output := flag.String("output", "", "Optional output file for results")
	flag.Parse()

	if *filePath == "" {
		fmt.Println("Please specify a file path using -file flag")
		flag.Usage()
		os.Exit(1)
	}

	// Read keys from file
	keys, err := readKeysFromFile(*filePath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error reading file: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Found %d keys in file\n", len(keys))
	fmt.Printf("Starting quota check with %d workers\n", *workers)

	// Create channels for communication
	keysChan := make(chan string, len(keys))
	resultsChan := make(chan Result, len(keys))

	// Start worker pool
	var wg sync.WaitGroup
	for i := 0; i < *workers; i++ {
		wg.Add(1)
		go worker(i, keysChan, resultsChan, &wg)
	}

	// Send keys to workers
	for _, key := range keys {
		keysChan <- key
	}
	close(keysChan)

	// Wait for all workers to finish in a separate goroutine
	go func() {
		wg.Wait()
		close(resultsChan)
	}()

	// Prepare for results processing
	successColor := color.New(color.FgGreen).SprintFunc()
	errorColor := color.New(color.FgRed).SprintFunc()
	warnColor := color.New(color.FgYellow).SprintFunc()

	// Collect results
	var validKeys, invalidKeys, quotaExceededKeys int
	var results []string

	// Process results
	for result := range resultsChan {
		keyMasked := maskAPIKey(result.Key)
		var line string

		if result.IsValid {
			if result.HasQuota {
				line = fmt.Sprintf("%s: %s", successColor("✓ VALID (HAS QUOTA)"), keyMasked)
				validKeys++
			} else {
				line = fmt.Sprintf("%s: %s - %s", warnColor("! QUOTA EXCEEDED"), keyMasked, result.Message)
				quotaExceededKeys++
			}
		} else {
			line = fmt.Sprintf("%s: %s - %s", errorColor("✗ ERROR"), keyMasked, result.Message)
			invalidKeys++
		}
		
		fmt.Println(line)
		results = append(results, line)
	}

	// Summary
	fmt.Println("\nSummary:")
	fmt.Printf("Valid keys with quota: %s\n", successColor(validKeys))
	fmt.Printf("Valid keys with exceeded quota: %s\n", warnColor(quotaExceededKeys))
	fmt.Printf("Invalid keys: %s\n", errorColor(invalidKeys))

	// Write results to file if requested
	if *output != "" {
		writeResultsToFile(*output, results)
	}
}

// worker handles checking API keys
func worker(id int, keys <-chan string, results chan<- Result, wg *sync.WaitGroup) {
	defer wg.Done()

	for key := range keys {
		result := checkAPIKey(key)
		results <- result
		// Small delay to prevent overwhelming the API
		time.Sleep(100 * time.Millisecond)
	}
}

// checkAPIKey makes an HTTP request to OpenAI API to check if the key is valid and has quota
func checkAPIKey(apiKey string) Result {
	req := fasthttp.AcquireRequest()
	resp := fasthttp.AcquireResponse()
	defer fasthttp.ReleaseRequest(req)
	defer fasthttp.ReleaseResponse(resp)

	// Set up a minimal chat completion request
	payload := `{
		"model": "gpt-3.5-turbo",
		"messages": [
			{"role": "user", "content": "Hello"}
		],
		"max_tokens": 1
	}`

	req.SetRequestURI(openAIEndpoint)
	req.Header.SetMethod("POST")
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", "application/json")
	req.SetBodyString(payload)

	err := fasthttp.Do(req, resp)
	if err != nil {
		return Result{
			Key:      apiKey,
			IsValid:  false,
			HasQuota: false,
			Message:  fmt.Sprintf("HTTP request failed: %v", err),
		}
	}

	statusCode := resp.StatusCode()
	body := string(resp.Body())

	switch {
	case statusCode == 200:
		return Result{
			Key:      apiKey,
			IsValid:  true,
			HasQuota: true,
			Message:  "Key is valid and has quota",
		}
	case statusCode == 401:
		return Result{
			Key:      apiKey,
			IsValid:  false,
			HasQuota: false,
			Message:  "Authentication failed",
		}
	case statusCode == 429:
		// The key is valid (authenticated) but has no quota or rate limited
		return Result{
			Key:      apiKey,
			IsValid:  true,
			HasQuota: false,
			Message:  "Rate limit or quota exceeded",
		}
	default:
		// Check if the error message indicates quota issues
		if strings.Contains(strings.ToLower(body), "quota") ||
			strings.Contains(strings.ToLower(body), "rate limit") {
			return Result{
				Key:      apiKey,
				IsValid:  true,
				HasQuota: false,
				Message:  "Quota issue detected",
			}
		}

		return Result{
			Key:      apiKey,
			IsValid:  false,
			HasQuota: false,
			Message:  fmt.Sprintf("Error: %d - %s", statusCode, body),
		}
	}
}

// readKeysFromFile reads API keys from a file (one key per line)
func readKeysFromFile(filePath string) ([]string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var keys []string
	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && strings.HasPrefix(line, "sk-") {
			keys = append(keys, line)
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return keys, nil
}

// maskAPIKey masks most of the API key for display
func maskAPIKey(key string) string {
	if len(key) <= 10 {
		return key
	}
	return key[:7] + "..." + key[len(key)-5:]
}

// writeResultsToFile writes the results to a file
func writeResultsToFile(filename string, results []string) {
	file, err := os.Create(filename)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating output file: %v\n", err)
		return
	}
	defer file.Close()

	for _, line := range results {
		// Remove color codes for file output
		cleanLine := stripANSI(line)
		file.WriteString(cleanLine + "\n")
	}

	fmt.Printf("Results written to %s\n", filename)
}

// stripANSI removes ANSI color codes from a string
func stripANSI(str string) string {
	// Simple regex replacement to remove ANSI color codes
	for i := 0; i < 100; i++ { // Simple brute-force approach
		str = strings.ReplaceAll(str, fmt.Sprintf("\x1b[%dm", i), "")
	}
	return str
}