From 9bc381e63b2f3e5692e039599f4717c8b50b6266 Mon Sep 17 00:00:00 2001
From: Ben Steel <bendavidsteel@gmail.com>
Date: Mon, 29 Jul 2024 20:17:40 +0100
Subject: [PATCH] Restored functionality for download bytes method (#1174)

---
 TikTokApi/api/video.py    | 38 +++++++++++++-------------------------
 TikTokApi/helpers.py      | 12 ++++++++++++
 TikTokApi/tiktok.py       | 10 ++++++++++
 examples/video_example.py |  3 +++
 4 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/TikTokApi/api/video.py b/TikTokApi/api/video.py
index bc8e9ca1..16271c4b 100644
--- a/TikTokApi/api/video.py
+++ b/TikTokApi/api/video.py
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from ..helpers import extract_video_id_from_url
+from ..helpers import extract_video_id_from_url, requests_cookie_to_playwright_cookie
 from typing import TYPE_CHECKING, ClassVar, Iterator, Optional
 from datetime import datetime
 import requests
@@ -153,6 +153,13 @@ async def info(self, **kwargs) -> dict:
             
         self.as_dict = video_info
         self.__extract_from_data()
+
+        cookies = [requests_cookie_to_playwright_cookie(c) for c in r.cookies]
+
+        await self.parent.set_session_cookies(
+            session, 
+            cookies
+        )
         return video_info
 
     async def bytes(self, **kwargs) -> bytes:
@@ -172,37 +179,18 @@ async def bytes(self, **kwargs) -> bytes:
                     output.write(video_bytes)
         """
 
-        raise NotImplementedError
         i, session = self.parent._get_session(**kwargs)
         downloadAddr = self.as_dict["video"]["downloadAddr"]
 
         cookies = await self.parent.get_session_cookies(session)
-        cookie_str = "; ".join([f"{k}={v}" for k, v in cookies.items()])
 
         h = session.headers
-        h["cookie"] = cookie_str
-
-        # Fetching the video bytes using a browser fetch within the page context
-        file_bytes = await session.page.evaluate(
-            """
-        async (url, headers) => {
-            const response = await fetch(url, { headers });
-            if (response.ok) {
-                const buffer = await response.arrayBuffer();
-                return new Uint8Array(buffer);
-            } else {
-                return `Error: ${response.statusText}`;  // Return an error message if the fetch fails
-            }
-        }
-        """,
-            (downloadAddr, h),
-        )
+        h["range"] = 'bytes=0-'
+        h["accept-encoding"] = 'identity;q=1, *;q=0'
+        h["referer"] = 'https://www.tiktok.com/'
 
-        byte_values = [
-            value
-            for key, value in sorted(file_bytes.items(), key=lambda item: int(item[0]))
-        ]
-        return bytes(byte_values)
+        resp = requests.get(downloadAddr, headers=h, cookies=cookies)
+        return resp.content
 
     def __extract_from_data(self) -> None:
         data = self.as_dict
diff --git a/TikTokApi/helpers.py b/TikTokApi/helpers.py
index 5050b560..cbca75c3 100644
--- a/TikTokApi/helpers.py
+++ b/TikTokApi/helpers.py
@@ -22,3 +22,15 @@ def random_choice(choices: list):
     if choices is None or len(choices) == 0:
         return None
     return random.choice(choices)
+
+def requests_cookie_to_playwright_cookie(req_c):
+    c = {
+        'name': req_c.name,
+        'value': req_c.value,
+        'domain': req_c.domain,
+        'path': req_c.path,
+        'secure': req_c.secure
+    }
+    if req_c.expires:
+        c['expires'] = req_c.expires
+    return c
diff --git a/TikTokApi/tiktok.py b/TikTokApi/tiktok.py
index 86d29b41..50167eee 100644
--- a/TikTokApi/tiktok.py
+++ b/TikTokApi/tiktok.py
@@ -317,6 +317,16 @@ def _get_session(self, **kwargs):
             i = random.randint(0, self.num_sessions - 1)
         return i, self.sessions[i]
 
+    async def set_session_cookies(self, session, cookies):
+        """
+        Set the cookies for a session
+
+        Args:
+            session (TikTokPlaywrightSession): The session to set the cookies for.
+            cookies (dict): The cookies to set for the session.
+        """
+        await session.context.add_cookies(cookies)
+
     async def get_session_cookies(self, session):
         """
         Get the cookies for a session
diff --git a/examples/video_example.py b/examples/video_example.py
index 3d35a72b..90f28606 100644
--- a/examples/video_example.py
+++ b/examples/video_example.py
@@ -20,6 +20,9 @@ async def get_video_example():
 
         video_info = await video.info()  # is HTML request, so avoid using this too much
         print(video_info)
+        video_bytes = await video.bytes()
+        with open("video.mp4", "wb") as f:
+            f.write(video_bytes)
 
 
 if __name__ == "__main__":