feat: integrate creator hub douyin workflows

cmd/docker_gateway/douyin.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import base64
+import binascii
 import http.client
 import json
 import logging
@@ -31,6 +33,7 @@ IDENTITY_URL = (
 WORKS_PATH = "/aweme/v1/web/aweme/post/"
 COMMENTS_PATH = "/aweme/v1/web/comment/list/"
 RESPONSE_LIMIT = 1 << 20
+MEDIA_RESPONSE_LIMIT = 32 << 20
 CONTROL_TIMEOUT = 15.0
 UID_RE = re.compile(r"^[1-9][0-9]{0,19}$")
 ID_RE = re.compile(r"^[1-9][0-9]{0,63}$")
@@ -62,6 +65,13 @@ class BrowserResponse:
     challenge: str = ""
 
 
+@dataclass(frozen=True)
+class BrowserMediaResponse:
+    status: int
+    content_type: str
+    body_base64: str
+
+
 class CDPConnection:
     def __init__(self, socket: websocket.WebSocket) -> None:
         self.socket = socket
@@ -212,10 +222,20 @@ class CDPConnection:
 
 
 class DouyinBrowser:
-    def __init__(self, endpoint: Callable[[str], str] | None = None) -> None:
+    def __init__(
+        self,
+        endpoint: Callable[[str], str] | None = None,
+        *,
+        origin: str = ORIGIN,
+        url_validator: Callable[[object], bool] | None = None,
+        media_validator: Callable[[object], bool] | None = None,
+    ) -> None:
         self.endpoint = endpoint or (
             lambda alias: f"http://creatorhub-browser-{alias}:9222"
         )
+        self.origin = origin
+        self.url_validator = url_validator or is_douyin_url
+        self.media_validator = media_validator or is_douyin_media_url
 
     @contextmanager
     def connection(self, alias: str):
@@ -272,14 +292,18 @@ class DouyinBrowser:
             for target in targets
             if isinstance(target, dict) and target.get("type") == "page"
         ]
-        douyin_targets = [
-            target for target in page_targets if is_douyin_url(target.get("url", ""))
+        matching_targets = [
+            target
+            for target in page_targets
+            if self.url_validator(target.get("url", ""))
         ]
-        if len(douyin_targets) > 1 or (not douyin_targets and len(page_targets) > 1):
+        if len(matching_targets) > 1 or (
+            not matching_targets and len(page_targets) > 1
+        ):
             raise DouyinError("browser has more than one page target")
         target = (
-            douyin_targets[0]
-            if douyin_targets
+            matching_targets[0]
+            if matching_targets
             else page_targets[0]
             if page_targets
             else None
@@ -319,63 +343,9 @@ class DouyinBrowser:
             raise DouyinError("browser CDP connection failed") from exc
         return CDPConnection(socket)
 
-    def set_cookies(self, alias: str, cookies: list[dict]) -> None:
-        with self.connection(alias) as cdp:
-            cdp.command("Network.enable")
-            cdp.command("Network.clearBrowserCookies")
-            cdp.command("Page.enable")
-            navigation = cdp.command("Page.navigate", {"url": ORIGIN_URL})
-            frame_id = navigation.get("frameId")
-            if not isinstance(frame_id, str) or navigation.get("errorText"):
-                raise DouyinError("browser navigation failed")
-            cdp.wait_event(
-                "Page.frameNavigated",
-                lambda params: (
-                    params.get("frame", {}).get("id") == frame_id
-                    and is_douyin_url(params.get("frame", {}).get("url"))
-                ),
-            )
-            deadline = time.monotonic() + CONTROL_TIMEOUT
-            last_error: DouyinError | None = None
-            while time.monotonic() < deadline:
-                try:
-                    state = cdp.evaluate(
-                        "({origin:location.origin,state:document.readyState})"
-                    )
-                    if (
-                        isinstance(state, dict)
-                        and state.get("origin") == ORIGIN
-                        and state.get("state") in {"interactive", "complete"}
-                    ):
-                        break
-                except DouyinError as exc:
-                    last_error = exc
-                time.sleep(0.1)
-            else:
-                if last_error:
-                    raise DouyinError("browser did not reach Douyin") from last_error
-                raise DouyinError("browser did not reach Douyin")
-            values = []
-            for cookie in cookies:
-                value = {
-                    "name": cookie["name"],
-                    "value": cookie["value"],
-                    "url": ORIGIN_URL,
-                    "domain": cookie["domain"],
-                    "path": cookie.get("path") or "/",
-                    "secure": bool(cookie.get("secure")),
-                    "httpOnly": bool(cookie.get("http_only")),
-                }
-                if cookie.get("same_site"):
-                    value["sameSite"] = cookie["same_site"]
-                if cookie.get("expires"):
-                    value["expires"] = cookie["expires"]
-                values.append(value)
-            cdp.command("Network.setCookies", {"cookies": values})
-
     def get(self, alias: str, target: str) -> BrowserResponse:
         with self.connection(alias) as cdp:
-            if cdp.evaluate("location.origin") != ORIGIN:
+            if cdp.evaluate("location.origin") != self.origin:
                 raise DouyinError("restricted browser origin changed")
             expression = f"""(async()=>{{
               const r=await fetch({json.dumps(target)},{{credentials:'include',redirect:'error'}});
@@ -402,6 +372,72 @@ class DouyinBrowser:
             raise DouyinError("restricted browser fetch returned invalid body")
         return BrowserResponse(status, body, detect_challenge(status, body))
 
+    def get_media(self, alias: str, target: str) -> BrowserMediaResponse:
+        if not self.media_validator(target):
+            raise DouyinError("restricted browser media target is invalid")
+        with self.connection(alias) as cdp:
+            navigation = cdp.command("Page.navigate", {"url": target})
+            frame_id = navigation.get("frameId")
+            if not isinstance(frame_id, str) or navigation.get("errorText"):
+                raise DouyinError("Douyin media page navigation failed")
+            cdp.wait_event(
+                "Page.frameNavigated",
+                lambda params: (
+                    params.get("frame", {}).get("id") == frame_id
+                    and self.url_validator(params.get("frame", {}).get("url"))
+                ),
+            )
+            deadline = time.monotonic() + CONTROL_TIMEOUT
+            while time.monotonic() < deadline:
+                if cdp.evaluate("document.readyState") in {"interactive", "complete"}:
+                    break
+                time.sleep(0.1)
+            else:
+                raise DouyinError("Douyin media page did not load")
+            result = cdp.evaluate(
+                f"""(async()=>{{
+                  const video=document.querySelector('video');
+                  const source=video?.currentSrc||video?.src||'';
+                  if(!source)return {{error:'media_source_unavailable'}};
+                  const r=await fetch(source,{{credentials:'include',redirect:'error'}});
+                  if(!r.body)return {{status:r.status,content_type:r.headers.get('content-type')||'',body:''}};
+                  const reader=r.body.getReader(), chunks=[]; let size=0;
+                  for(;;){{const item=await reader.read();if(item.done)break;
+                    if(size+item.value.byteLength>{MEDIA_RESPONSE_LIMIT}){{await reader.cancel();return {{too_large:true}};}}
+                    size+=item.value.byteLength;chunks.push(item.value);
+                  }}
+                  const bytes=new Uint8Array(size); let offset=0;
+                  for(const chunk of chunks){{bytes.set(chunk,offset);offset+=chunk.length;}}
+                  let binary='';
+                  for(let offset=0;offset<bytes.length;offset+=0x8000)binary+=String.fromCharCode(...bytes.subarray(offset,Math.min(offset+0x8000,bytes.length)));
+                  return {{status:r.status,content_type:r.headers.get('content-type')||'',body:btoa(binary)}};
+                }})()"""
+            )
+        if (
+            not isinstance(result, dict)
+            or result.get("too_large")
+            or result.get("error")
+        ):
+            raise DouyinError("Douyin media download failed")
+        status = result.get("status")
+        content_type = result.get("content_type")
+        body = result.get("body")
+        if (
+            not isinstance(status, int)
+            or not isinstance(content_type, str)
+            or not isinstance(body, str)
+        ):
+            raise DouyinError("Douyin media response is invalid")
+        if status < 200 or status >= 300:
+            raise DouyinError("Douyin media response was not successful")
+        try:
+            decoded_size = len(base64.b64decode(body, validate=True))
+        except (ValueError, binascii.Error) as exc:
+            raise DouyinError("Douyin media response is invalid") from exc
+        if decoded_size > MEDIA_RESPONSE_LIMIT:
+            raise DouyinError("Douyin media response is too large")
+        return BrowserMediaResponse(status, content_type, body)
+
     def identity(self, alias: str, expected_uid: str | None = None) -> dict:
         response = self.get(alias, IDENTITY_URL)
         try:
@@ -521,7 +557,7 @@ class DouyinBrowser:
 
     def _evaluate(self, alias: str, expression: str) -> object:
         with self.connection(alias) as cdp:
-            if cdp.evaluate("location.origin") != ORIGIN:
+            if cdp.evaluate("location.origin") != self.origin:
                 raise DouyinError("restricted browser origin changed")
             return cdp.evaluate(expression)
 
@@ -1020,6 +1056,43 @@ def _notice_id(value: object) -> str:
     return ""
 
 
+def _direct_message_type(detail: dict, message: object) -> str:
+    raw_type = (
+        detail.get("message_type") or detail.get("msg_type") or detail.get("type")
+    )
+    if isinstance(message, dict):
+        raw_type = raw_type or message.get("type") or message.get("message_type")
+    if isinstance(raw_type, int) and not isinstance(raw_type, bool):
+        raw_type = {1: "text", 2: "image", 3: "voice", 4: "video", 5: "sticker"}.get(
+            raw_type
+        )
+    if isinstance(raw_type, str):
+        normalized = raw_type.strip().lower()
+        aliases = {
+            "text": "text",
+            "txt": "text",
+            "image": "image",
+            "img": "image",
+            "picture": "image",
+            "voice": "voice",
+            "audio": "voice",
+            "video": "video",
+            "sticker": "sticker",
+            "emoji": "sticker",
+        }
+        if normalized in aliases:
+            return aliases[normalized]
+    if any(detail.get(key) for key in ("voice", "audio", "audio_url")):
+        return "voice"
+    if any(detail.get(key) for key in ("image", "picture", "image_url")):
+        return "image"
+    if any(detail.get(key) for key in ("video", "video_url")):
+        return "video"
+    if isinstance(message, str) and message:
+        return "text"
+    return "non_text"
+
+
 def normalize_notice(notice: object) -> dict | None:
     if not isinstance(notice, dict):
         raise DouyinError("notification detail is invalid")
@@ -1062,6 +1135,7 @@ def normalize_notice(notice: object) -> dict | None:
             or notice.get("text")
             or ""
         )
+        message_type = _direct_message_type(detail, message)
         if isinstance(message, dict):
             message = message.get("text") or message.get("content") or ""
         if not isinstance(message, str) or len(message) > 100000:
@@ -1082,6 +1156,7 @@ def normalize_notice(notice: object) -> dict | None:
             "interactor_uid": _notice_id(sender.get("uid") or sender.get("user_id")),
             "comment_id": "",
             "work_id": "",
+            "message_type": message_type,
             "message_text": message,
         }
         create_time = notice.get("create_time")
@@ -1180,6 +1255,23 @@ def is_douyin_url(value: object) -> bool:
         return False
 
 
+def is_douyin_media_url(value: object) -> bool:
+    if not isinstance(value, str):
+        return False
+    try:
+        parsed = urlsplit(value)
+        return (
+            parsed.scheme == "https"
+            and parsed.hostname == "www.douyin.com"
+            and parsed.port is None
+            and parsed.username is None
+            and parsed.password is None
+            and parsed.path.startswith("/video/")
+        )
+    except (TypeError, ValueError):
+        return False
+
+
 def detect_challenge(status: int, body: str) -> str:
     if status == 429:
         return ""