mijia-go-api/auth.go

package mijia

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/mdp/qrterminal/v3"
	"rsc.io/qr"
)

const qrLoginTimeout = 120 * time.Second

type AuthData struct {
	UA                     string            `json:"ua"`
	DeviceID               string            `json:"deviceId"`
	PassO                  string            `json:"pass_o"`
	Psecurity              string            `json:"psecurity"`
	Nonce                  string            `json:"nonce"`
	Ssecurity              string            `json:"ssecurity"`
	PassToken              string            `json:"passToken"`
	UserID                 string            `json:"userId"`
	CUserID                string            `json:"cUserId"`
	ServiceToken           string            `json:"serviceToken"`
	YetAnotherServiceToken string            `json:"yetAnotherServiceToken"`
	ExpireTime             int64             `json:"expireTime"`
	SaveTime               int64             `json:"saveTime"`
	Extra                  map[string]string `json:"-"`
}

func (data AuthData) MarshalJSON() ([]byte, error) {
	fields := map[string]any{
		"ua": data.UA, "deviceId": data.DeviceID, "pass_o": data.PassO,
		"psecurity": data.Psecurity, "nonce": data.Nonce, "ssecurity": data.Ssecurity,
		"passToken": data.PassToken, "userId": data.UserID, "cUserId": data.CUserID,
		"serviceToken": data.ServiceToken, "yetAnotherServiceToken": data.YetAnotherServiceToken,
		"expireTime": data.ExpireTime, "saveTime": data.SaveTime,
	}
	for key, value := range data.Extra {
		if _, stable := fields[key]; !stable {
			fields[key] = value
		}
	}
	return json.Marshal(fields)
}

func (data *AuthData) UnmarshalJSON(payload []byte) error {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(payload, &fields); err != nil {
		return err
	}
	stable := map[string]any{
		"ua": &data.UA, "deviceId": &data.DeviceID, "pass_o": &data.PassO,
		"psecurity": &data.Psecurity, "nonce": &data.Nonce, "ssecurity": &data.Ssecurity,
		"passToken": &data.PassToken, "userId": &data.UserID, "cUserId": &data.CUserID,
		"serviceToken": &data.ServiceToken, "yetAnotherServiceToken": &data.YetAnotherServiceToken,
		"expireTime": &data.ExpireTime, "saveTime": &data.SaveTime,
	}
	data.Extra = make(map[string]string)
	for key, raw := range fields {
		if target, ok := stable[key]; ok {
			if err := json.Unmarshal(raw, target); err != nil {
				return fmt.Errorf("decode auth field %s: %w", key, err)
			}
			continue
		}
		var value string
		if json.Unmarshal(raw, &value) == nil {
			data.Extra[key] = value
		}
	}
	return nil
}

func (data AuthData) complete() bool {
	return data.UA != "" && data.Ssecurity != "" && data.UserID != "" && data.CUserID != "" && data.ServiceToken != ""
}

func (data AuthData) zero() bool {
	return data.UA == "" && data.DeviceID == "" && data.PassO == "" && data.Psecurity == "" &&
		data.Nonce == "" && data.Ssecurity == "" && data.PassToken == "" && data.UserID == "" &&
		data.CUserID == "" && data.ServiceToken == "" && data.YetAnotherServiceToken == "" &&
		data.ExpireTime == 0 && data.SaveTime == 0 && len(data.Extra) == 0
}

func (data AuthData) yetAnotherServiceToken() string {
	if data.YetAnotherServiceToken != "" {
		return data.YetAnotherServiceToken
	}
	if data.Extra != nil && data.Extra["yetAnotherServiceToken"] != "" {
		return data.Extra["yetAnotherServiceToken"]
	}
	return data.ServiceToken
}

func (data AuthData) clone() AuthData {
	clone := data
	if data.Extra != nil {
		clone.Extra = make(map[string]string, len(data.Extra))
		for key, value := range data.Extra {
			clone.Extra[key] = value
		}
	}
	return clone
}

// AuthData returns a snapshot of the client's current authentication data.
func (client *Client) AuthData() AuthData {
	client.authMu.RLock()
	defer client.authMu.RUnlock()
	return client.authData.clone()
}

func (client *Client) setAuthData(authData AuthData) {
	client.authMu.Lock()
	client.authData = authData.clone()
	client.authMu.Unlock()
}

func (client *Client) updateAuthData(update func(*AuthData)) AuthData {
	client.authMu.Lock()
	defer client.authMu.Unlock()
	update(&client.authData)
	return client.authData.clone()
}

func (client *Client) loadAuthData() error {
	payload, err := os.ReadFile(client.authPath)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("read auth data: %w", err)
	}
	var authData AuthData
	if err := json.Unmarshal(payload, &authData); err != nil {
		return fmt.Errorf("decode auth data: %w", err)
	}
	client.setAuthData(authData)
	return nil
}

func (client *Client) ensureIdentity() {
	client.updateAuthData(func(authData *AuthData) {
		if authData.PassO == "" {
			authData.PassO = randomString(16, "0123456789abcdef")
		}
		if authData.DeviceID == "" {
			authData.DeviceID = randomString(16, "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ_-")
		}
		if authData.UA != "" {
			return
		}
		countryCode := "CN"
		if parts := strings.Split(client.locale, "_"); len(parts) == 2 {
			countryCode = parts[1]
		}
		id1 := randomString(40, "0123456789ABCDEF")
		id2 := randomString(32, "0123456789ABCDEF")
		id3 := randomString(32, "0123456789ABCDEF")
		id4 := randomString(40, "0123456789ABCDEF")
		authData.UA = fmt.Sprintf("Android-15-11.0.701-Xiaomi-23046RP50C-OS2.0.212.0.VMYCNXM-%s-%s-%s-%s-SmartHome-MI_APP_STORE-%s|%s|%s-64", id1, countryCode, id3, id2, id1, id4, authData.PassO)
	})
}

func randomString(length int, alphabet string) string {
	random := make([]byte, length)
	if _, err := rand.Read(random); err != nil {
		panic(fmt.Sprintf("generate random identity: %v", err))
	}
	for index := range random {
		random[index] = alphabet[int(random[index])%len(alphabet)]
	}
	return string(random)
}

func (client *Client) saveAuthData() error {
	authData := client.AuthData()
	authData.SaveTime = time.Now().UnixMilli()
	if err := client.writeAuthData(authData); err != nil {
		return err
	}
	client.setAuthData(authData)
	return nil
}

func (client *Client) commitAuthData(authData AuthData) error {
	if !authData.complete() {
		return fmt.Errorf("incomplete auth data")
	}
	authData.SaveTime = time.Now().UnixMilli()
	if client.authPath != "" {
		if err := client.writeAuthData(authData); err != nil {
			return err
		}
	} else if client.authDataChanged != nil {
		if err := client.authDataChanged(authData.clone()); err != nil {
			return fmt.Errorf("persist changed auth data: %w", err)
		}
	}
	client.setAuthData(authData)
	return nil
}

func (client *Client) writeAuthData(authData AuthData) error {
	if client.authPath == "" {
		return nil
	}
	payload, err := json.MarshalIndent(authData, "", "  ")
	if err != nil {
		return fmt.Errorf("encode auth data: %w", err)
	}
	directory := filepath.Dir(client.authPath)
	if err := os.MkdirAll(directory, 0o700); err != nil {
		return fmt.Errorf("create auth directory: %w", err)
	}
	temporary, err := os.CreateTemp(directory, ".auth-*.json")
	if err != nil {
		return fmt.Errorf("create temporary auth file: %w", err)
	}
	temporaryPath := temporary.Name()
	defer os.Remove(temporaryPath)
	if err := temporary.Chmod(0o600); err != nil {
		temporary.Close()
		return fmt.Errorf("secure temporary auth file: %w", err)
	}
	if _, err := temporary.Write(payload); err != nil {
		temporary.Close()
		return fmt.Errorf("write auth data: %w", err)
	}
	if err := temporary.Sync(); err != nil {
		temporary.Close()
		return fmt.Errorf("sync auth data: %w", err)
	}
	if err := temporary.Close(); err != nil {
		return fmt.Errorf("close auth data: %w", err)
	}
	if err := os.Rename(temporaryPath, client.authPath); err != nil {
		return fmt.Errorf("replace auth data: %w", err)
	}
	return nil
}

func parseServiceResponse(payload []byte, target any) error {
	payload = []byte(strings.TrimPrefix(string(payload), "&&&START&&&"))
	if err := json.Unmarshal(payload, target); err != nil {
		return fmt.Errorf("decode login response: %w", err)
	}
	return nil
}

type serviceLoginData struct {
	Code      int    `json:"code"`
	Desc      string `json:"desc"`
	Location  string `json:"location"`
	Ssecurity string `json:"ssecurity"`
}

type qrLoginData struct {
	Code     int    `json:"code"`
	Desc     string `json:"desc"`
	LoginURL string `json:"loginUrl"`
	QR       string `json:"qr"`
	LP       string `json:"lp"`
}

type longPollData struct {
	Code      int            `json:"code"`
	Desc      string         `json:"desc"`
	Location  string         `json:"location"`
	Psecurity string         `json:"psecurity"`
	Nonce     stringOrNumber `json:"nonce"`
	Ssecurity string         `json:"ssecurity"`
	PassToken string         `json:"passToken"`
	UserID    stringOrNumber `json:"userId"`
	CUserID   stringOrNumber `json:"cUserId"`
}

type stringOrNumber string

type recordingWriter struct {
	w   io.Writer
	err error
}

func (writer *recordingWriter) Write(payload []byte) (int, error) {
	if writer.err != nil {
		return 0, writer.err
	}
	written, err := writer.w.Write(payload)
	if err == nil && written < len(payload) {
		err = io.ErrShortWrite
	}
	if err != nil {
		writer.err = err
	}
	return written, err
}

func (value *stringOrNumber) UnmarshalJSON(payload []byte) error {
	decoded, err := decodeStringOrNumber(payload)
	if err != nil {
		return err
	}
	*value = stringOrNumber(decoded)
	return nil
}

func (client *Client) Login(ctx context.Context) (AuthData, error) {
	client.loginMu.Lock()
	defer client.loginMu.Unlock()
	client.ensureIdentity()

	location, refreshedAuthData, err := client.getLocation(ctx)
	if err != nil {
		return AuthData{}, err
	}
	if refreshedAuthData != nil {
		if err := client.commitAuthData(*refreshedAuthData); err != nil {
			return AuthData{}, err
		}
		return client.AuthData(), nil
	}
	loginData, err := client.getQRLoginData(ctx, location)
	if err != nil {
		return AuthData{}, err
	}
	if client.qrWriter != nil {
		if _, err := qr.Encode(loginData.LoginURL, qr.L); err != nil {
			return AuthData{}, fmt.Errorf("encode login QR code: %w", err)
		}
		writer := &recordingWriter{w: client.qrWriter}
		if _, err := fmt.Fprintf(writer, "请使用米家APP扫描下方二维码\n%s\n", loginData.LoginURL); err != nil {
			return AuthData{}, fmt.Errorf("write QR login output: %w", err)
		}
		qrterminal.GenerateHalfBlock(loginData.LoginURL, qrterminal.L, writer)
		if writer.err != nil {
			return AuthData{}, fmt.Errorf("write QR login output: %w", writer.err)
		}
		if loginData.QR != "" {
			if _, err := fmt.Fprintf(writer, "二维码图片: %s\n", loginData.QR); err != nil {
				return AuthData{}, fmt.Errorf("write QR login output: %w", err)
			}
		}
	}
	return client.completeQRLogin(ctx, loginData)
}

func (client *Client) getLocation(ctx context.Context) (url.Values, *AuthData, error) {
	httpClient := client.newSession()
	serviceURL, err := url.Parse(client.serviceLoginURL)
	if err != nil {
		return nil, nil, fmt.Errorf("parse service login URL: %w", err)
	}
	query := serviceURL.Query()
	query.Set("_json", "true")
	query.Set("sid", "mijia")
	query.Set("_locale", client.locale)
	serviceURL.RawQuery = query.Encode()
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, serviceURL.String(), nil)
	if err != nil {
		return nil, nil, err
	}
	client.setLoginHeaders(request, true)
	var data serviceLoginData
	if err := client.doLoginRequestWithClient(httpClient, request, false, &data); err != nil {
		return nil, nil, err
	}
	if data.Location == "" {
		return nil, nil, &LoginError{Code: data.Code, Message: "登录响应缺少 location"}
	}
	if data.Code == 0 {
		refreshRequest, err := http.NewRequestWithContext(ctx, http.MethodGet, data.Location, nil)
		if err != nil {
			return nil, nil, err
		}
		client.setLoginHeaders(refreshRequest, false)
		response, err := httpClient.Do(refreshRequest)
		if err != nil {
			return nil, nil, fmt.Errorf("refresh login token: %w", err)
		}
		body, readErr := readHTTPResponse(response)
		response.Body.Close()
		if readErr != nil {
			return nil, nil, fmt.Errorf("read token refresh response: %w", readErr)
		}
		if response.StatusCode != http.StatusOK {
			return nil, nil, &LoginError{Code: response.StatusCode, Message: string(body)}
		}
		if string(body) != "ok" {
			return nil, nil, &LoginError{Code: -1, Message: string(body)}
		}
		candidate := client.AuthData()
		serviceTokenReceived := updateAuthDataFromCookies(&candidate, httpClient, response.Request.URL)
		candidate.Ssecurity = data.Ssecurity
		if !serviceTokenReceived || !candidate.complete() {
			return nil, nil, fmt.Errorf("%w: %w", ErrReauthenticationRequired, &LoginError{Code: -1, Message: "刷新Token响应认证信息不完整"})
		}
		candidate.ExpireTime = time.Now().Add(30 * 24 * time.Hour).UnixMilli()
		return nil, &candidate, nil
	}
	locationURL, err := url.Parse(data.Location)
	if err != nil {
		return nil, nil, fmt.Errorf("parse login location: %w", err)
	}
	return locationURL.Query(), nil, nil
}

func (client *Client) getQRLoginData(ctx context.Context, location url.Values) (qrLoginData, error) {
	location.Set("theme", "")
	location.Set("bizDeviceType", "")
	location.Set("_hasLogo", "false")
	location.Set("_qrsize", "240")
	location.Set("_dc", fmt.Sprintf("%d", time.Now().UnixMilli()))
	loginURL, err := url.Parse(client.loginURL)
	if err != nil {
		return qrLoginData{}, err
	}
	loginURL.RawQuery = location.Encode()
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, loginURL.String(), nil)
	if err != nil {
		return qrLoginData{}, err
	}
	client.setLoginHeaders(request, false)
	var data qrLoginData
	if err := client.doLoginRequest(request, true, &data); err != nil {
		return qrLoginData{}, err
	}
	if data.LoginURL == "" || data.LP == "" {
		return qrLoginData{}, &LoginError{Code: data.Code, Message: "二维码登录响应不完整"}
	}
	return data, nil
}

func (client *Client) completeQRLogin(ctx context.Context, loginData qrLoginData) (AuthData, error) {
	pollContext, cancel := context.WithTimeout(ctx, qrLoginTimeout)
	defer cancel()
	httpClient := client.newSession()
	request, err := http.NewRequestWithContext(pollContext, http.MethodGet, loginData.LP, nil)
	if err != nil {
		return AuthData{}, err
	}
	client.setLoginHeaders(request, false)
	var data longPollData
	if err := client.doLoginRequestWithClient(httpClient, request, true, &data); err != nil {
		if errors.Is(err, context.DeadlineExceeded) {
			return AuthData{}, &LoginError{Code: -1, Message: "超时，请重试"}
		}
		return AuthData{}, err
	}
	callback, err := http.NewRequestWithContext(ctx, http.MethodGet, data.Location, nil)
	if err != nil {
		return AuthData{}, err
	}
	client.setLoginHeaders(callback, false)
	response, err := httpClient.Do(callback)
	if err != nil {
		return AuthData{}, fmt.Errorf("complete login callback: %w", err)
	}
	_, readErr := readHTTPResponse(response)
	response.Body.Close()
	if readErr != nil {
		return AuthData{}, fmt.Errorf("read login callback response: %w", readErr)
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return AuthData{}, &LoginError{Code: response.StatusCode, Message: "登录回调失败"}
	}
	candidate := client.AuthData()
	candidate.Ssecurity = ""
	candidate.UserID = ""
	candidate.CUserID = ""
	candidate.ServiceToken = ""
	serviceTokenReceived := updateAuthDataFromCookies(&candidate, httpClient, response.Request.URL)
	candidate.Psecurity = data.Psecurity
	candidate.Nonce = string(data.Nonce)
	candidate.Ssecurity = data.Ssecurity
	candidate.PassToken = data.PassToken
	candidate.UserID = string(data.UserID)
	candidate.CUserID = string(data.CUserID)
	if !serviceTokenReceived || !candidate.complete() {
		return AuthData{}, &LoginError{Code: -1, Message: "登录回调认证信息不完整"}
	}
	candidate.ExpireTime = time.Now().Add(30 * 24 * time.Hour).UnixMilli()
	if err := client.commitAuthData(candidate); err != nil {
		return AuthData{}, err
	}
	return client.AuthData(), nil
}

func (client *Client) doLoginRequest(request *http.Request, verifyCode bool, target any) error {
	return client.doLoginRequestWithClient(client.session(), request, verifyCode, target)
}

func (client *Client) doLoginRequestWithClient(httpClient *http.Client, request *http.Request, verifyCode bool, target any) error {
	response, err := httpClient.Do(request)
	if err != nil {
		return fmt.Errorf("send login request: %w", err)
	}
	defer response.Body.Close()
	body, err := readHTTPResponse(response)
	if err != nil {
		return fmt.Errorf("read login response: %w", err)
	}
	if response.StatusCode != http.StatusOK {
		return &LoginError{Code: response.StatusCode, Message: string(body)}
	}
	if err := parseServiceResponse(body, target); err != nil {
		return err
	}
	if verifyCode {
		encoded, _ := json.Marshal(target)
		var status struct {
			Code int    `json:"code"`
			Desc string `json:"desc"`
		}
		_ = json.Unmarshal(encoded, &status)
		if status.Code != 0 {
			return &LoginError{Code: status.Code, Message: status.Desc}
		}
	}
	return nil
}

func (client *Client) setLoginHeaders(request *http.Request, withCookies bool) {
	authData := client.AuthData()
	request.Header.Set("User-Agent", authData.UA)
	request.Header.Set("Connection", "keep-alive")
	request.Header.Set("Accept-Encoding", "gzip")
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if withCookies {
		request.Header.Set("Cookie", strings.Join([]string{
			"deviceId=" + authData.DeviceID,
			"pass_o=" + authData.PassO,
			"passToken=" + authData.PassToken,
			"userId=" + authData.UserID,
			"cUserId=" + authData.CUserID,
			"uLocale=" + client.locale,
		}, ";"))
	}
}

func updateAuthDataFromCookies(authData *AuthData, httpClient *http.Client, target *url.URL) bool {
	if httpClient.Jar == nil || target == nil {
		return false
	}
	cookies := httpClient.Jar.Cookies(target)
	if authData.Extra == nil {
		authData.Extra = make(map[string]string)
	}
	serviceTokenReceived := false
	for _, cookie := range cookies {
		switch cookie.Name {
		case "serviceToken":
			if cookie.Value != "" {
				authData.ServiceToken = cookie.Value
				serviceTokenReceived = true
			}
		case "yetAnotherServiceToken":
			authData.YetAnotherServiceToken = cookie.Value
		case "cUserId":
			authData.CUserID = cookie.Value
		default:
			authData.Extra[cookie.Name] = cookie.Value
		}
	}
	return serviceTokenReceived
}

func (client *Client) refreshToken(ctx context.Context) error {
	client.loginMu.Lock()
	defer client.loginMu.Unlock()

	available, _ := client.Available(ctx)
	if available {
		return nil
	}
	_, refreshedAuthData, err := client.getLocation(ctx)
	if err != nil {
		return err
	}
	if refreshedAuthData == nil {
		return fmt.Errorf("%w: %w", ErrReauthenticationRequired, &LoginError{Code: -1, Message: "刷新Token失败，请重新登录"})
	}
	if err := client.commitAuthData(*refreshedAuthData); err != nil {
		return err
	}
	return nil
}