1package auth
2
3import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "fmt"
8 "io"
9 "net/http"
10 "time"
11)
12
13// KimiDeviceAuthorization is the OAuth 2 device-code response used by
14// the official Kimi Code CLI.
15type KimiDeviceAuthorization struct {
16 UserCode string `json:"user_code"`
17 DeviceCode string `json:"device_code"`
18 VerificationURI string `json:"verification_uri"`
19 VerificationURIComplete string `json:"verification_uri_complete"`
20 ExpiresIn int `json:"expires_in"`
21 Interval int `json:"interval"`
22}
23
24// RequestKimiDeviceAuthorization starts Kimi Code's device-code login.
25func RequestKimiDeviceAuthorization(ctx context.Context) (KimiDeviceAuthorization, error) {
26 form := bytes.NewBufferString("client_id=" + KimiOAuth.ClientID)
27 req, err := http.NewRequestWithContext(ctx, "POST", "https://auth.kimi.com/api/oauth/device_authorization", form)
28 if err != nil {
29 return KimiDeviceAuthorization{}, err
30 }
31 req.Header.Set("content-type", "application/x-www-form-urlencoded")
32 req.Header.Set("accept", "application/json")
33 req.Header.Set("user-agent", "zot")
34
35 resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req)
36 if err != nil {
37 return KimiDeviceAuthorization{}, fmt.Errorf("kimi device authorization: %w", err)
38 }
39 defer resp.Body.Close()
40 body, _ := io.ReadAll(resp.Body)
41 if resp.StatusCode < 200 || resp.StatusCode >= 300 {
42 return KimiDeviceAuthorization{}, fmt.Errorf("kimi device authorization http %d: %s", resp.StatusCode, string(body))
43 }
44 var out KimiDeviceAuthorization
45 if err := json.Unmarshal(body, &out); err != nil {
46 return out, fmt.Errorf("parse kimi device authorization: %w", err)
47 }
48 if out.Interval <= 0 {
49 out.Interval = 5
50 }
51 return out, nil
52}
53
54// PollKimiDeviceToken polls until the browser/device-code login completes.
55func PollKimiDeviceToken(ctx context.Context, auth KimiDeviceAuthorization) (*OAuthToken, error) {
56 interval := time.Duration(auth.Interval) * time.Second
57 if interval <= 0 {
58 interval = 5 * time.Second
59 }
60 for {
61 tok, retry, err := pollKimiDeviceTokenOnce(ctx, auth.DeviceCode)
62 if err != nil {
63 return nil, err
64 }
65 if tok != nil {
66 return tok, nil
67 }
68 select {
69 case <-ctx.Done():
70 return nil, ctx.Err()
71 case <-time.After(retry):
72 }
73 }
74}
75
76func pollKimiDeviceTokenOnce(ctx context.Context, deviceCode string) (*OAuthToken, time.Duration, error) {
77 form := bytes.NewBufferString("client_id=" + KimiOAuth.ClientID + "&device_code=" + deviceCode + "&grant_type=urn:ietf:params:oauth:grant-type:device_code")
78 req, err := http.NewRequestWithContext(ctx, "POST", KimiOAuth.TokenURL, form)
79 if err != nil {
80 return nil, 0, err
81 }
82 req.Header.Set("content-type", "application/x-www-form-urlencoded")
83 req.Header.Set("accept", "application/json")
84 req.Header.Set("user-agent", "zot")
85 resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req)
86 if err != nil {
87 return nil, 0, fmt.Errorf("kimi token poll: %w", err)
88 }
89 defer resp.Body.Close()
90 body, _ := io.ReadAll(resp.Body)
91 var raw struct {
92 AccessToken string `json:"access_token"`
93 RefreshToken string `json:"refresh_token"`
94 TokenType string `json:"token_type"`
95 ExpiresIn float64 `json:"expires_in"`
96 Scope string `json:"scope"`
97 Error string `json:"error"`
98 ErrorDescription string `json:"error_description"`
99 }
100 _ = json.Unmarshal(body, &raw)
101 if resp.StatusCode == http.StatusOK && raw.AccessToken != "" {
102 return &OAuthToken{
103 AccessToken: raw.AccessToken,
104 RefreshToken: raw.RefreshToken,
105 TokenType: raw.TokenType,
106 Scope: raw.Scope,
107 ClientID: KimiOAuth.ClientID,
108 Expiry: time.Now().Add(time.Duration(raw.ExpiresIn) * time.Second),
109 }, 0, nil
110 }
111 if raw.Error == "authorization_pending" || raw.Error == "slow_down" || resp.StatusCode == http.StatusBadRequest {
112 return nil, 5 * time.Second, nil
113 }
114 if raw.Error != "" {
115 return nil, 0, fmt.Errorf("kimi token poll: %s: %s", raw.Error, raw.ErrorDescription)
116 }
117 return nil, 0, fmt.Errorf("kimi token poll http %d: %s", resp.StatusCode, string(body))
118}