OpenList/drivers/139/util_test.go

package _139

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/OpenListTeam/OpenList/v4/drivers/base"
	"github.com/go-resty/resty/v2"
)

func useRetryingTestClient(t *testing.T) {
	t.Helper()
	oldClient := base.RestyClient
	base.RestyClient = resty.New().SetRetryCount(3)
	t.Cleanup(func() {
		base.RestyClient = oldClient
	})
}

func TestMergeMailCookiesPreservesExistingOrderAndAppendsNewNames(t *testing.T) {
	got := mergeMailCookieHeader("z=zv; behaviorid=b; Os_SSo_Sid=old", []*http.Cookie{
		{Name: "RMKEY", Value: "rm"},
		{Name: "Os_SSo_Sid", Value: "sid"},
		{Name: "a", Value: "av"},
	})
	want := "z=zv;behaviorid=b;Os_SSo_Sid=sid;RMKEY=rm;a=av"
	if got != want {
		t.Fatalf("mergeMailCookieHeader() = %q, want %q", got, want)
	}
}

func TestExtractFastLoginCookies(t *testing.T) {
	sid, rmkey := extractFastLoginCookies("RMKEY=rm; Os_SSo_Sid=sid")
	if sid != "sid" || rmkey != "rm" {
		t.Fatalf("extractFastLoginCookies() = %q, %q; want sid, rm", sid, rmkey)
	}
}

func TestCredentialState(t *testing.T) {
	tests := []struct {
		name string
		d    Yun139
		want credentialState
		err  bool
	}{
		{
			name: "authorization",
			d:    Yun139{Addition: Addition{Authorization: " auth "}},
			want: credentialStateAuthorization,
		},
		{
			name: "full login",
			d: Yun139{Addition: Addition{
				MailCookies: "RMKEY=rm; Os_SSo_Sid=sid",
				Username:    "user",
				Password:    "password",
			}},
			want: credentialStateFullLogin,
		},
		{
			name: "full login without initial cookies",
			d: Yun139{Addition: Addition{
				Username: "user",
				Password: "password",
			}},
			want: credentialStateFullLogin,
		},
		{
			name: "cookies only",
			d:    Yun139{Addition: Addition{MailCookies: "RMKEY=rm; Os_SSo_Sid=sid"}},
			want: credentialStateCookiesOnly,
		},
		{
			name: "partial password login",
			d:    Yun139{Addition: Addition{Username: "user"}},
			err:  true,
		},
		{
			name: "missing credentials",
			d:    Yun139{},
			err:  true,
		},
		{
			name: "invalid cookie",
			d:    Yun139{Addition: Addition{MailCookies: "invalid-cookie"}},
			err:  true,
		},
		{
			name: "authorization with basic prefix",
			d:    Yun139{Addition: Addition{Authorization: "Basic abc"}},
			err:  true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := tt.d.credentialState()
			if tt.err {
				if err == nil {
					t.Fatal("credentialState() expected error")
				}
				return
			}
			if err != nil {
				t.Fatalf("credentialState() unexpected error: %v", err)
			}
			if got != tt.want {
				t.Fatalf("credentialState() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestIntegrationLoginObtainsAuthorization(t *testing.T) {
	if os.Getenv("OPENLIST_139_INTEGRATION") != "1" {
		t.Skip("set OPENLIST_139_INTEGRATION=1 to run live 139Yun login checks")
	}
	base.RestyClient = resty.New().
		SetHeader("user-agent", base.UserAgent).
		SetRetryCount(3).
		SetRetryResetReaders(true).
		SetTimeout(30 * time.Second)

	username := os.Getenv("OPENLIST_139_USERNAME")
	password := os.Getenv("OPENLIST_139_PASSWORD")
	mailCookies := os.Getenv("OPENLIST_139_MAIL_COOKIES")
	authorization := strings.TrimSpace(os.Getenv("OPENLIST_139_AUTHORIZATION"))

	if authorization != "" {
		t.Run("authorization", func(t *testing.T) {
			d := Yun139{Addition: Addition{Authorization: authorization}}
			state, err := d.credentialState()
			if err != nil {
				t.Fatalf("credentialState() unexpected error: %v", err)
			}
			if state != credentialStateAuthorization {
				t.Fatalf("credentialState() = %v, want authorization", state)
			}
			if d.Authorization == "" || strings.HasPrefix(strings.ToLower(d.Authorization), "basic ") {
				t.Fatal("authorization should be present without Basic prefix")
			}
		})
	}

	if mailCookies == "" {
		t.Fatal("OPENLIST_139_MAIL_COOKIES is required")
	}

	runFastLogin := func(mailCookies string) error {
		d := Yun139{Addition: Addition{MailCookies: mailCookies}}
		state, err := d.credentialState()
		if err != nil {
			return fmt.Errorf("credentialState() error: %w", err)
		}
		if state != credentialStateCookiesOnly {
			return fmt.Errorf("credentialState() = %v, want cookies only", state)
		}
		sid, rmkey := extractFastLoginCookies(d.MailCookies)
		if sid == "" || rmkey == "" {
			return fmt.Errorf("mail cookies are missing Os_SSo_Sid or RMKEY")
		}
		token, err := d.step2_get_single_token(sid)
		if err != nil {
			return fmt.Errorf("step2_get_single_token() error: %w", err)
		}
		auth, err := d.step3_third_party_login(token)
		if err != nil {
			return fmt.Errorf("step3_third_party_login() error: %w", err)
		}
		if auth == "" {
			return fmt.Errorf("authorization is empty after fast login")
		}
		return nil
	}

	if username == "" || password == "" {
		t.Fatal("OPENLIST_139_USERNAME and OPENLIST_139_PASSWORD are required for password fallback")
	}

	var refreshedMailCookies string
	var generatedAuthorization string
	t.Run("password login fallback", func(t *testing.T) {
		d := Yun139{Addition: Addition{
			MailCookies: mailCookies,
			Username:    username,
			Password:    password,
		}}
		state, err := d.credentialState()
		if err != nil {
			t.Fatalf("credentialState() unexpected error: %v", err)
		}
		if state != credentialStateFullLogin {
			t.Fatalf("credentialState() = %v, want full login", state)
		}
		passId, err := d.step1_password_login()
		if err != nil {
			t.Fatalf("step1_password_login() error: %v", err)
		}
		token, err := d.step2_get_single_token(passId)
		if err != nil {
			t.Fatalf("step2_get_single_token() error: %v", err)
		}
		auth, err := d.step3_third_party_login(token)
		if err != nil {
			t.Fatalf("step3_third_party_login() error: %v", err)
		}
		d.Authorization = auth
		if auth == "" || d.Authorization == "" {
			t.Fatal("authorization is empty after password login")
		}
		generatedAuthorization = auth
		refreshedMailCookies = d.MailCookies
	})

	t.Run("authorization generated by password login", func(t *testing.T) {
		if generatedAuthorization == "" {
			t.Fatal("password login did not generate authorization")
		}
		d := Yun139{Addition: Addition{Authorization: generatedAuthorization}}
		state, err := d.credentialState()
		if err != nil {
			t.Fatalf("credentialState() unexpected error: %v", err)
		}
		if state != credentialStateAuthorization {
			t.Fatalf("credentialState() = %v, want authorization", state)
		}
		if strings.HasPrefix(strings.ToLower(d.Authorization), "basic ") {
			t.Fatal("authorization should not include Basic prefix")
		}
	})

	t.Run("mail cookies fast login from input", func(t *testing.T) {
		sid, rmkey := extractFastLoginCookies(mailCookies)
		if sid == "" || rmkey == "" {
			t.Skip("input mail cookies are missing Os_SSo_Sid or RMKEY")
		}
		if err := runFastLogin(mailCookies); err != nil {
			t.Skipf("input mail cookies are stale or unusable: %v", err)
		}
	})

	t.Run("mail cookies fast login after password login", func(t *testing.T) {
		if refreshedMailCookies == "" {
			t.Fatal("password login did not refresh mail cookies")
		}
		if err := runFastLogin(refreshedMailCookies); err != nil {
			t.Fatalf("refreshed mail cookies fast login failed: %v", err)
		}
	})
}

func TestInvalidAuthorizationDoesNotUseCookieFastLogin(t *testing.T) {
	d := Yun139{Addition: Addition{
		Authorization: "not-base64",
		MailCookies:   "Os_SSo_Sid=sid; RMKEY=rmkey",
	}}

	err := d.refreshToken()
	if err == nil || !strings.Contains(err.Error(), "password login failed") {
		t.Fatalf("refreshToken() error = %v, want password login fallback error", err)
	}
	if d.Authorization != "not-base64" {
		t.Fatalf("Authorization = %q, want original invalid value retained", d.Authorization)
	}
}

func TestSMSSceneForRisk(t *testing.T) {
	tests := map[string]int{
		"S025":         1,
		"S035":         1,
		"PML401010062": 2,
		"MW0016":       4,
	}
	for riskCode, want := range tests {
		got, ok := smsSceneForRisk(riskCode)
		if !ok || got != want {
			t.Fatalf("smsSceneForRisk(%q) = %d, %t; want %d, true", riskCode, got, ok, want)
		}
	}
	if _, ok := smsSceneForRisk("PICTURE_ONLY"); ok {
		t.Fatal("smsSceneForRisk() accepted unsupported risk code")
	}
}

func TestSendSMSVerificationCode(t *testing.T) {
	useRetryingTestClient(t)

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := io.ReadAll(r.Body)
		if err != nil {
			t.Fatalf("read request body: %v", err)
		}
		if r.URL.Query().Get("func") != "login:sendSmsCodeByScene" {
			t.Errorf("func = %q", r.URL.Query().Get("func"))
		}
		if !strings.Contains(string(body), `<string name="scene">1</string>`) {
			t.Errorf("request body = %s", body)
		}
		if !strings.Contains(r.Header.Get("Cookie"), "device=fingerprint") {
			t.Errorf("Cookie = %q", r.Header.Get("Cookie"))
		}
		http.SetCookie(w, &http.Cookie{Name: "challenge", Value: "sms"})
		_, _ = io.WriteString(w, `{"code":"S_OK"}`)
	}))
	defer server.Close()

	oldURL := mailSMSURL
	mailSMSURL = server.URL
	defer func() { mailSMSURL = oldURL }()

	d := Yun139{Addition: Addition{Username: "18800000000", MailCookies: "device=fingerprint"}}
	if err := d.sendSMSVerificationCode("S025"); err != nil {
		t.Fatalf("sendSMSVerificationCode() error = %v", err)
	}
	if !strings.Contains(d.MailCookies, "challenge=sms") {
		t.Fatalf("MailCookies = %q", d.MailCookies)
	}
}

func TestSendSMSVerificationCodeStopsAtPictureChallenge(t *testing.T) {
	useRetryingTestClient(t)

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = io.WriteString(w, `{"code":"PML401010021"}`)
	}))
	defer server.Close()

	oldURL := mailSMSURL
	mailSMSURL = server.URL
	defer func() { mailSMSURL = oldURL }()

	d := Yun139{Addition: Addition{Username: "18800000000"}}
	err := d.sendSMSVerificationCode("S025")
	if err == nil || !strings.Contains(err.Error(), "requires picture verification") {
		t.Fatalf("sendSMSVerificationCode() error = %v", err)
	}
}

func TestVerifySMSCode(t *testing.T) {
	useRetryingTestClient(t)

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := io.ReadAll(r.Body)
		if err != nil {
			t.Fatalf("read request body: %v", err)
		}
		if r.URL.Query().Get("func") != "/login/inlogin.action" {
			t.Errorf("func = %q", r.URL.Query().Get("func"))
		}
		wantHash := sha1Hash("fetion.com.cn:123456")
		if !strings.Contains(string(body), `<string name="loginPassword">`+wantHash+`</string>`) {
			t.Errorf("request body does not contain SMS code hash: %s", body)
		}
		http.SetCookie(w, &http.Cookie{Name: "RMKEY", Value: "new-rmkey"})
		_, _ = io.WriteString(w, `{"code":"S_OK","var":{"loginSuccessUrl":"https://mail.10086.cn/?sid=sms-sid"}}`)
	}))
	defer server.Close()

	oldURL := mailSMSURL
	mailSMSURL = server.URL
	defer func() { mailSMSURL = oldURL }()

	d := Yun139{Addition: Addition{
		Username:    "18800000000",
		SmsCode:     "123456",
		MailCookies: "challenge=sms",
	}}
	sid, err := d.verifySMSCode("S025")
	if err != nil {
		t.Fatalf("verifySMSCode() error = %v", err)
	}
	if sid != "sms-sid" {
		t.Fatalf("sid = %q", sid)
	}
	if d.SmsCode != "" {
		t.Fatalf("SmsCode = %q, want cleared", d.SmsCode)
	}
	if !strings.Contains(d.MailCookies, "RMKEY=new-rmkey") {
		t.Fatalf("MailCookies = %q", d.MailCookies)
	}
}

func TestPasswordLoginWithoutInitialCookiesStopsAtRedirectAndPersistsCookies(t *testing.T) {
	useRetryingTestClient(t)

	var loginRequests, redirectRequests int
	var server *httptest.Server
	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/Login/Login.ashx":
			loginRequests++
			http.SetCookie(w, &http.Cookie{Name: "RMKEY", Value: "new-rm"})
			http.SetCookie(w, &http.Cookie{Name: "Os_SSo_Sid", Value: "new-sid"})
			w.Header().Set("Location", server.URL+"/appmail?sid=single-sid")
			w.WriteHeader(http.StatusFound)
		case "/appmail":
			redirectRequests++
			w.WriteHeader(http.StatusOK)
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	oldURL := mailPasswordURL
	mailPasswordURL = server.URL + "/Login/Login.ashx"
	defer func() { mailPasswordURL = oldURL }()

	d := Yun139{Addition: Addition{Username: "18800000000", Password: "password"}}
	sid, err := d.step1_password_login()
	if err != nil || sid != "single-sid" {
		t.Fatalf("sid=%q err=%v", sid, err)
	}
	if loginRequests != 1 || redirectRequests != 0 {
		t.Fatalf("login/redirect requests=%d/%d, want 1/0", loginRequests, redirectRequests)
	}
	if !strings.Contains(d.MailCookies, "RMKEY=new-rm") || !strings.Contains(d.MailCookies, "Os_SSo_Sid=new-sid") {
		t.Fatalf("response cookies were not persisted: %q", d.MailCookies)
	}
}

func TestPasswordLoginReusesRawRefreshedCookies(t *testing.T) {
	useRetryingTestClient(t)

	var n int
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n++
		cookie := r.Header.Get("Cookie")
		if !strings.Contains(cookie, "JSESSIONID=stale") || !strings.Contains(cookie, "behaviorid=device") {
			t.Errorf("login %d lost raw cookie context: %q", n, cookie)
		}
		if n == 1 {
			http.SetCookie(w, &http.Cookie{Name: "RMKEY", Value: "rm-1"})
			http.SetCookie(w, &http.Cookie{Name: "Os_SSo_Sid", Value: "sid-1"})
			w.Header().Set("Location", "https://mail.10086.cn/?sid=first")
		} else {
			if !strings.Contains(cookie, "RMKEY=rm-1") || !strings.Contains(cookie, "Os_SSo_Sid=sid-1") {
				t.Errorf("second login did not reuse first refreshed cookies: %q", cookie)
			}
			http.SetCookie(w, &http.Cookie{Name: "RMKEY", Value: "rm-2"})
			http.SetCookie(w, &http.Cookie{Name: "Os_SSo_Sid", Value: "sid-2"})
			w.Header().Set("Location", "https://mail.10086.cn/?sid=second")
		}
		w.WriteHeader(http.StatusFound)
	}))
	defer server.Close()

	oldURL := mailPasswordURL
	mailPasswordURL = server.URL
	defer func() { mailPasswordURL = oldURL }()

	d := Yun139{Addition: Addition{Username: "18800000000", Password: "password", MailCookies: "Os_SSo_Sid=old; RMKEY=old; JSESSIONID=stale; behaviorid=device"}}
	if sid, err := d.step1_password_login(); err != nil || sid != "first" {
		t.Fatalf("first login sid=%q err=%v", sid, err)
	}
	refreshed := d.MailCookies
	if !strings.Contains(refreshed, "RMKEY=rm-1") || !strings.Contains(refreshed, "Os_SSo_Sid=sid-1") {
		t.Fatalf("first login did not persist refreshed cookies: %q", refreshed)
	}
	d.Authorization = ""
	d.MailCookies = refreshed
	if sid, err := d.step1_password_login(); err != nil || sid != "second" {
		t.Fatalf("second login sid=%q err=%v", sid, err)
	}
	if n != 2 {
		t.Fatalf("login count=%d, want 2", n)
	}
}