fanqie-go/internal/fqapi/videodecrypt_test.go

package fqapi

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/binary"
	"bytes"
	"encoding/base64"
	"encoding/hex"
	"io"
	"os"
	"strings"
	"testing"
)

// 番茄短剧真实样本（video_model video_1..4 共用同一 spade_a）
const testSpadeA = "kbwfwFiKI8RuvCnHV4ss9lK6HMFmvS/2f7wv7H+mAfRRvxmXlw=="

func TestDecodeSpadeKey(t *testing.T) {
	key, err := DecodeSpadeKey(testSpadeA)
	if err != nil {
		t.Fatal(err)
	}
	want := "4f04d75b0df5cee67dbe13738872e965"
	if got := hex.EncodeToString(key); got != want {
		t.Fatalf("key mismatch: got %s want %s", got, want)
	}
}

// 端到端：下载真实 360p 视频 → DecryptCENCMP4 → 输出文件（ffprobe 验证）
func TestDecryptRealVideo(t *testing.T) {
	mainURLB64 := os.Getenv("FQ_VIDEO_URL")
	spadeA := strings.TrimSpace(os.Getenv("FQ_VIDEO_SPADE"))
	if mainURLB64 == "" || spadeA == "" {
		t.Skip("need FQ_VIDEO_URL (main_url base64) and FQ_VIDEO_SPADE")
	}
	url, err := base64.StdEncoding.DecodeString(mainURLB64)
	if err != nil {
		t.Fatal(err)
	}
	body, _, err := downloadMedia(string(bytes.TrimSpace(url)), "", map[string]string{})
	if err != nil {
		t.Fatal(err)
	}
	defer body.Close()
	data, err := io.ReadAll(body)
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("downloaded %d bytes", len(data))

	key, err := DecodeSpadeKey(spadeA)
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("key=%x", key)
	if err := DecryptCENCMP4(data, key); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile("/tmp/video_decrypted.mp4", data, 0644); err != nil {
		t.Fatal(err)
	}
	t.Log("written /tmp/video_decrypted.mp4")
}

// ---- 合成 CENC MP4 单元测试：加密 → 解密 → 比对明文 ----

func box(typ string, payload ...[]byte) []byte {
	var body []byte
	for _, p := range payload {
		body = append(body, p...)
	}
	b := make([]byte, 8+len(body))
	binary.BigEndian.PutUint32(b[:4], uint32(len(b)))
	copy(b[4:8], typ)
	copy(b[8:], body)
	return b
}

func u32(v uint32) []byte { b := make([]byte, 4); binary.BigEndian.PutUint32(b, v); return b }

func TestDecryptCENCMP4Synthetic(t *testing.T) {
	key := []byte("0123456789abcdef")
	plain1 := []byte("HELLO_CENC_SAMPLE_1_PADPADP")
	plain2 := []byte("SECOND_SAMPLE")
	plain3 := []byte("THIRD")

	enc := func(pt []byte, iv8 []byte) []byte {
		iv := make([]byte, 16)
		copy(iv, iv8)
		block, _ := aes.NewCipher(key)
		ct := make([]byte, len(pt))
		cipher.NewCTR(block, iv).XORKeyStream(ct, pt)
		return ct
	}
	iv1 := []byte{1, 2, 3, 4, 5, 6, 7, 8}
	iv2 := []byte{8, 7, 6, 5, 4, 3, 2, 1}
	iv3 := []byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x11, 0x22}
	ct1, ct2, ct3 := enc(plain1, iv1), enc(plain2, iv2), enc(plain3, iv3)

	// buildFile(off1, off2)：chunk 绝对偏移写入 stco；偏移依赖 moov 长度，先占位测长再重建
	buildFile := func(off1, off2 uint32) []byte {
		entry := box("enca", make([]byte, 8), box("sinf", box("frma", []byte("mp4a"))))
		stsd := box("stsd", u32(0), u32(1), entry)
		stsz := box("stsz", u32(0), u32(0), u32(3), u32(uint32(len(plain1))), u32(uint32(len(plain2))), u32(uint32(len(plain3))))
		stsc := box("stsc", u32(0), u32(2),
			u32(1), u32(2), u32(1), // entry1: chunk1 起，每 chunk 2 样本
			u32(2), u32(1), u32(1)) // entry2: chunk2 起，每 chunk 1 样本
		stco := box("stco", u32(0), u32(2), u32(off1), u32(off2))
		sencBody := append(append(append(append(u32(0), u32(3)...), iv1...), iv2...), iv3...)
		senc := box("senc", sencBody)
		moov := box("moov", box("trak", box("mdia", box("minf", box("stbl", stsd, stsz, stsc, stco, senc)))))
		ftyp := box("ftyp", []byte("isom"))
		return append(append(ftyp, moov...), box("mdat", ct1, ct2, ct3)...)
	}

	// 占位构造测 moov 总长，得到真实 chunk 偏移后重建
	probe := buildFile(0, 0)
	chunk1Off := uint32(len(probe) - len(ct1) - len(ct2) - len(ct3))
	chunk2Off := chunk1Off + uint32(len(ct1)+len(ct2))
	data := buildFile(chunk1Off, chunk2Off)

	if err := DecryptCENCMP4(data, key); err != nil {
		t.Fatal(err)
	}
	if got := data[chunk1Off : chunk1Off+uint32(len(plain1))]; string(got) != string(plain1) {
		t.Errorf("sample1 解密不符: %q", got)
	}
	s2 := chunk1Off + uint32(len(plain1))
	if got := data[s2 : s2+uint32(len(plain2))]; string(got) != string(plain2) {
		t.Errorf("sample2 解密不符: %q", got)
	}
	if got := data[chunk2Off : chunk2Off+uint32(len(plain3))]; string(got) != string(plain3) {
		t.Errorf("sample3 解密不符: %q", got)
	}
	// stsd entry 类型应被 patch 为 frma 指示的 mp4a（原 enca 四字已不存在）
	if bytes.Index(data, []byte("enca")) >= 0 || bytes.Index(data, []byte("mp4a")) < 0 {
		t.Error("stsd entry 类型未修复为 mp4a")
	}
}

// 清晰度过滤：按 definition 与档位 key 均可命中；无匹配返回空列表由上层报错
func TestDecodeVideoModelQualityFilter(t *testing.T) {
	vm := `{"video_id":"v123","video_duration":1.5,"video_list":{
		"video_1":{"definition":"360p","vtype":"mp4","main_url":"` + base64.StdEncoding.EncodeToString([]byte("https://a/360")) + `"},
		"video_4":{"definition":"720p","vtype":"mp4","main_url":"` + base64.StdEncoding.EncodeToString([]byte("https://a/720")) + `"}}}`
	got := decodeVideoModelAsJSON(vm, "url", "", "720p").(map[string]interface{})
	vl := got["video_list"].(map[string]interface{})
	if len(vl) != 1 || vl["video_4"] == nil {
		t.Fatalf("definition 过滤失败: %v", vl)
	}
	got = decodeVideoModelAsJSON(vm, "url", "", "VIDEO_1").(map[string]interface{})
	vl = got["video_list"].(map[string]interface{})
	if len(vl) != 1 || vl["video_1"] == nil {
		t.Fatalf("档位key过滤失败: %v", vl)
	}
	if u := vl["video_1"].(map[string]interface{})["main_url"].(string); u != "https://a/360" {
		t.Fatalf("main_url 未解码: %s", u)
	}
	got = decodeVideoModelAsJSON(vm, "url", "", "1080p").(map[string]interface{})
	if len(got["video_list"].(map[string]interface{})) != 0 {
		t.Fatal("无匹配应返回空列表")
	}
	// max 取最高档（vheight 缺失时按 definition 数字比较）
	got = decodeVideoModelAsJSON(vm, "url", "", "max").(map[string]interface{})
	vl = got["video_list"].(map[string]interface{})
	if len(vl) != 1 || vl["video_4"] == nil {
		t.Fatalf("max 应只保留最高档: %v", vl)
	}
	// vheight 存在时优先按 vheight
	vm2 := `{"video_list":{
		"video_1":{"definition":"360p","vwidth":360,"vheight":640},
		"video_2":{"definition":"540p","vwidth":540,"vheight":960}}}`
	got = decodeVideoModelAsJSON(vm2, "url", "", "max").(map[string]interface{})
	if len(got["video_list"].(map[string]interface{})) != 1 || got["video_list"].(map[string]interface{})["video_2"] == nil {
		t.Fatal("max 应按 vheight 取 video_2")
	}
}