✨ feat: use protobuf danmaku source when login (#340)

Author: SigureMo

README.md

@@ -274,7 +274,11 @@ https://github.com/orgs/community/discussions/16925#discussioncomment-7571187
 -  可选值 `"ass" | "xml" | "protobuf"`
 -  默认值 `"ass"`
 
-B 站提供了 `xml` 与 `protobuf` 两种弹幕数据接口，yutto 会自动下载 `xml` 格式弹幕并转换为 `ass` 格式，如果你不喜欢 yutto 自动转换的效果，可以选择输出格式为 `xml` 或 `protobuf`，手动通过一些工具进行转换，比如 yutto 和 bilili 所使用的 [biliass](https://github.com/yutto-dev/yutto/tree/main/packages/biliass)，或者使用 [us-danmaku](https://tiansh.github.io/us-danmaku/bilibili/) 进行在线转换。
+B 站提供了 `xml` 与 `protobuf` 两种弹幕数据接口，`xml` 接口为旧接口，弹幕数上限较低，`protobuf` 接口相对较高，但不登录情况下只能获取很少的弹幕
+
+为了确保无论是否登录都能获取最多的弹幕，yutto 在登录时会下载 `protobuf` 源数据，在未登录时会下载 `xml` 源数据，并将其转换为主流播放器支持的 `ass` 格式
+
+如果你不喜欢 yutto 自动转换的效果，可以选择输出格式为 `xml` 或 `protobuf`，手动通过一些工具进行转换，比如 yutto 和 bilili 所使用的 [biliass](https://github.com/yutto-dev/yutto/tree/main/packages/biliass)，或者使用 [us-danmaku](https://tiansh.github.io/us-danmaku/bilibili/) 进行在线转换。
 
 如果你不想下载弹幕，只需要使用参数 `--no-danmaku` 即可。
 

packages/biliass/rust/src/lib.rs

@@ -24,6 +24,7 @@ fn biliass_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<python::PyDanmakuElem>()?;
     m.add_class::<python::PyComment>()?;
     m.add_class::<python::PyCommentPosition>()?;
+    m.add_function(wrap_pyfunction!(python::py_get_danmaku_meta_size, m)?)?;
     m.add_function(wrap_pyfunction!(python::py_read_comments_from_xml, m)?)?;
     m.add_function(wrap_pyfunction!(python::py_read_comments_from_protobuf, m)?)?;
     m.add_function(wrap_pyfunction!(python::py_parse_special_comment, m)?)?;

packages/biliass/rust/src/python/convert.rs

@@ -38,7 +38,6 @@ pub fn py_xml_to_ass(
 #[allow(clippy::too_many_arguments)]
 #[pyfunction(name = "protobuf_to_ass")]
 pub fn py_protobuf_to_ass(
-    // inputs: Vec<Py<PyAny>>,
     inputs: Vec<PyBackedBytes>,
     stage_width: u32,
     stage_height: u32,

packages/biliass/rust/src/python/mod.rs

@@ -5,7 +5,7 @@ mod reader;
 
 pub use comment::{PyComment, PyCommentPosition};
 pub use convert::{py_protobuf_to_ass, py_xml_to_ass};
-pub use proto::{PyDanmakuElem, PyDmSegMobileReply};
+pub use proto::{py_get_danmaku_meta_size, PyDanmakuElem, PyDmSegMobileReply};
 pub use reader::{
     py_parse_special_comment, py_read_comments_from_protobuf, py_read_comments_from_xml,
 };

packages/biliass/rust/src/python/proto.rs

@@ -130,3 +130,13 @@ impl PyDmSegMobileReply {
         ))
     }
 }
+
+#[pyfunction(name = "get_danmaku_meta_size")]
+pub fn py_get_danmaku_meta_size(buffer: &[u8]) -> PyResult<usize> {
+    let dm_sge_opt = proto::DmWebViewReply::decode(&mut Cursor::new(buffer))
+        .map(|reply| reply.dm_sge)
+        .map_err(error::DecodeError::from)
+        .map_err(error::BiliassError::from)?;
+
+    Ok(dm_sge_opt.map_or(0, |dm_sge| dm_sge.total as usize))
+}

packages/biliass/src/biliass/__init__.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from biliass._core import get_danmaku_meta_size as get_danmaku_meta_size
+
 from .biliass import (
     Danmaku2ASS as Danmaku2ASS,
     read_comments_bilibili_protobuf as read_comments_bilibili_protobuf,

packages/biliass/src/biliass/_core.pyi

@@ -86,3 +86,4 @@ def protobuf_to_ass(
     comment_filter: list[str],
     is_reduce_comments: bool,
 ) -> str: ...
+def get_danmaku_meta_size(buffer: bytes) -> int: ...

src/yutto/_typing.py

@@ -56,12 +56,57 @@ class AvId(BilibiliId):
     ```
     """
 
+    # id conversion based on
+    # https://github.com/SocialSisterYi/bilibili-API-collect/blob/master/docs/misc/bvid_desc.md
+
+    XOR_CODE = 23442827791579
+    MASK_CODE = 2251799813685247
+    MAX_AID = 1 << 51
+    ALPHABET = "FcwAPNKTMug3GV5Lj7EJnHpWsx4tb8haYeviqBz6rkCy12mUSDQX9RdoZf"
+    ENCODE_MAP = 8, 7, 0, 5, 1, 3, 2, 4, 6
+    DECODE_MAP = tuple(reversed(ENCODE_MAP))
+
+    BASE = len(ALPHABET)
+    PREFIX = "BV1"
+    PREFIX_LEN = len(PREFIX)
+    CODE_LEN = len(ENCODE_MAP)
+
+    @staticmethod
+    def av2bv(aid: int) -> str:
+        bvid = [""] * 9
+        tmp = (AvId.MAX_AID | aid) ^ AvId.XOR_CODE
+        for i in range(AvId.CODE_LEN):
+            bvid[AvId.ENCODE_MAP[i]] = AvId.ALPHABET[tmp % AvId.BASE]
+            tmp //= AvId.BASE
+        return AvId.PREFIX + "".join(bvid)
+
+    @staticmethod
+    def bv2av(bvid: str) -> int:
+        assert bvid[:3] == AvId.PREFIX
+
+        bvid = bvid[3:]
+        tmp = 0
+        for i in range(AvId.CODE_LEN):
+            idx = AvId.ALPHABET.index(bvid[AvId.DECODE_MAP[i]])
+            tmp = tmp * AvId.BASE + idx
+        return (tmp & AvId.MASK_CODE) ^ AvId.XOR_CODE
+
     def to_dict(self) -> dict[str, str]:
         raise NotImplementedError("请不要直接使用 AvId")
 
     def to_url(self) -> str:
         raise NotImplementedError("请不要直接使用 AvId")
 
+    def as_aid(self) -> AId:
+        if isinstance(self, AId):
+            return self
+        return AId(str(self.bv2av(self.value)))
+
+    def as_bvid(self) -> BvId:
+        if isinstance(self, BvId):
+            return self
+        return BvId(self.av2bv(int(self.value)))
+
 
 class AId(AvId):
     """AID"""

src/yutto/api/danmaku.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
+import asyncio
 from typing import TYPE_CHECKING
 
+from biliass import get_danmaku_meta_size
+
+from yutto.api.user_info import get_user_info
 from yutto.utils.fetcher import Fetcher
 
 if TYPE_CHECKING:
     import httpx
 
-    from yutto._typing import CId
+    from yutto._typing import AvId, CId
     from yutto.utils.danmaku import DanmakuData, DanmakuSaveType
 
 
@@ -18,21 +22,34 @@ async def get_xml_danmaku(client: httpx.AsyncClient, cid: CId) -> str:
     return results
 
 
-async def get_protobuf_danmaku(client: httpx.AsyncClient, cid: CId, segment_id: int = 1) -> bytes:
+async def get_protobuf_danmaku_segment(client: httpx.AsyncClient, cid: CId, segment_id: int = 1) -> bytes:
     danmaku_api = "http://api.bilibili.com/x/v2/dm/web/seg.so?type=1&oid={cid}&segment_index={segment_id}"
     results = await Fetcher.fetch_bin(client, danmaku_api.format(cid=cid, segment_id=segment_id))
     assert results is not None
     return results
 
 
+async def get_protobuf_danmaku(client: httpx.AsyncClient, avid: AvId, cid: CId) -> list[bytes]:
+    danmaku_meta_api = "https://api.bilibili.com/x/v2/dm/web/view?type=1&oid={cid}&pid={aid}"
+    aid = avid.as_aid()
+    meta_results = await Fetcher.fetch_bin(client, danmaku_meta_api.format(cid=cid, aid=aid.value))
+    assert meta_results is not None
+    size = get_danmaku_meta_size(meta_results)
+
+    results = await asyncio.gather(
+        *[get_protobuf_danmaku_segment(client, cid, segment_id) for segment_id in range(1, size + 1)]
+    )
+    return results
+
+
 async def get_danmaku(
     client: httpx.AsyncClient,
     cid: CId,
+    avid: AvId,
     save_type: DanmakuSaveType,
-    last_n_segments: int = 2,
 ) -> DanmakuData:
-    # 暂时默认使用 XML 源
-    source_type = "xml" if save_type == "xml" or save_type == "ass" else "protobuf"
+    # 在已经登录的情况下，使用 protobuf，因为未登录时 protobuf 弹幕会少非常多
+    source_type = "xml" if save_type == "xml" or not (await get_user_info(client))["is_login"] else "protobuf"
     danmaku_data: DanmakuData = {
         "source_type": source_type,
         "save_type": save_type,
@@ -42,6 +59,5 @@ async def get_danmaku(
     if source_type == "xml":
         danmaku_data["data"].append(await get_xml_danmaku(client, cid))
     else:
-        for i in range(1, last_n_segments + 1):
-            danmaku_data["data"].append(await get_protobuf_danmaku(client, cid, i))
+        danmaku_data["data"].extend(await get_protobuf_danmaku(client, avid, cid))
     return danmaku_data

src/yutto/api/user_info.py

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING, Any, TypedDict
 
 from yutto._typing import UserInfo
+from yutto.utils.asynclib import async_cache
 from yutto.utils.fetcher import Fetcher
 
 if TYPE_CHECKING:
@@ -26,6 +27,7 @@ class WbiImg(TypedDict):
 dm_cover_img_str_cache: str = base64.b64encode("".join(random.choices(string.printable, k=random.randint(32, 128))).encode())[:-2].decode()  # fmt: skip
 
 
+@async_cache(lambda _: "user_info")
 async def get_user_info(client: AsyncClient) -> UserInfo:
     info_api = "https://api.bilibili.com/x/web-interface/nav"
     res_json = await Fetcher.fetch_json(client, info_api)