claude_usage_service.go 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. package repository
  2. import (
  3. "context"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "net/http"
  8. "time"
  9. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
  11. "github.com/Wei-Shaw/sub2api/internal/service"
  12. )
  13. const defaultClaudeUsageURL = "https://api.anthropic.com/api/oauth/usage"
  14. // 默认 User-Agent,与用户抓包的请求一致
  15. const defaultUsageUserAgent = "claude-code/2.1.7"
  16. type claudeUsageService struct {
  17. usageURL string
  18. allowPrivateHosts bool
  19. httpUpstream service.HTTPUpstream
  20. }
  21. // NewClaudeUsageFetcher 创建 Claude 用量获取服务
  22. // httpUpstream: 可选,如果提供则支持 TLS 指纹伪装
  23. func NewClaudeUsageFetcher(httpUpstream service.HTTPUpstream) service.ClaudeUsageFetcher {
  24. return &claudeUsageService{
  25. usageURL: defaultClaudeUsageURL,
  26. httpUpstream: httpUpstream,
  27. }
  28. }
  29. // FetchUsage 简单版本,不支持 TLS 指纹(向后兼容)
  30. func (s *claudeUsageService) FetchUsage(ctx context.Context, accessToken, proxyURL string) (*service.ClaudeUsageResponse, error) {
  31. return s.FetchUsageWithOptions(ctx, &service.ClaudeUsageFetchOptions{
  32. AccessToken: accessToken,
  33. ProxyURL: proxyURL,
  34. })
  35. }
  36. // FetchUsageWithOptions 完整版本,支持 TLS 指纹和自定义 User-Agent
  37. func (s *claudeUsageService) FetchUsageWithOptions(ctx context.Context, opts *service.ClaudeUsageFetchOptions) (*service.ClaudeUsageResponse, error) {
  38. if opts == nil {
  39. return nil, fmt.Errorf("options is nil")
  40. }
  41. // 创建请求
  42. req, err := http.NewRequestWithContext(ctx, "GET", s.usageURL, nil)
  43. if err != nil {
  44. return nil, fmt.Errorf("create request failed: %w", err)
  45. }
  46. // 设置请求头(与抓包一致,但不设置 Accept-Encoding,让 Go 自动处理压缩)
  47. req.Header.Set("Accept", "application/json, text/plain, */*")
  48. req.Header.Set("Content-Type", "application/json")
  49. req.Header.Set("Authorization", "Bearer "+opts.AccessToken)
  50. req.Header.Set("anthropic-beta", "oauth-2025-04-20")
  51. // 设置 User-Agent(优先使用缓存的 Fingerprint,否则使用默认值)
  52. userAgent := defaultUsageUserAgent
  53. if opts.Fingerprint != nil && opts.Fingerprint.UserAgent != "" {
  54. userAgent = opts.Fingerprint.UserAgent
  55. }
  56. req.Header.Set("User-Agent", userAgent)
  57. var resp *http.Response
  58. // 如果有 TLS Profile 且有 HTTPUpstream,使用 DoWithTLS
  59. if opts.TLSProfile != nil && s.httpUpstream != nil {
  60. resp, err = s.httpUpstream.DoWithTLS(req, opts.ProxyURL, opts.AccountID, 0, opts.TLSProfile)
  61. if err != nil {
  62. return nil, fmt.Errorf("request with TLS fingerprint failed: %w", err)
  63. }
  64. } else {
  65. // 不启用 TLS 指纹,使用普通 HTTP 客户端
  66. client, err := httpclient.GetClient(httpclient.Options{
  67. ProxyURL: opts.ProxyURL,
  68. Timeout: 30 * time.Second,
  69. ValidateResolvedIP: true,
  70. AllowPrivateHosts: s.allowPrivateHosts,
  71. })
  72. if err != nil {
  73. return nil, fmt.Errorf("create http client failed: %w", err)
  74. }
  75. resp, err = client.Do(req)
  76. if err != nil {
  77. return nil, fmt.Errorf("request failed: %w", err)
  78. }
  79. }
  80. defer func() { _ = resp.Body.Close() }()
  81. if resp.StatusCode != http.StatusOK {
  82. body, _ := io.ReadAll(resp.Body)
  83. msg := fmt.Sprintf("API returned status %d: %s", resp.StatusCode, string(body))
  84. return nil, infraerrors.New(http.StatusInternalServerError, "UPSTREAM_ERROR", msg)
  85. }
  86. var usageResp service.ClaudeUsageResponse
  87. if err := json.NewDecoder(resp.Body).Decode(&usageResp); err != nil {
  88. return nil, fmt.Errorf("decode response failed: %w", err)
  89. }
  90. return &usageResp, nil
  91. }