Merge pull request #2 from code-yeongyu/feature/replace-yt-dlp

Replace `yt-dlp` to `browser-cookie3` for cookie parsing

Author: code-yeongyu

poetry.lock

@@ -42,8 +42,8 @@ select = ["PLE", "PLW", "E", "W", "F", "I", "Q", "B"]
 [tool.poetry.dependencies]
 python = "^3.9"
 poetry = "^1.4.2"
-yt-dlp = "^2023.3.4"
 httpx = "^0.24.0"
+browser-cookie3 = "^0.17.1"
 
 [tool.pyright.defineConstant]
 DEBUG = true

pyproject.toml

@@ -1,3 +1,6 @@
 from .chatgpt_access_token_parser import ChatGPTAccessTokenParser as ChatGPTAccessTokenParser
 from .chatgpt_access_token_parser import get_access_token as get_access_token
 from .openai_cookie_parser import OpenAICookieParser as OpenAICookieParser
+from .supported_browser import SupportedBrowser as SupportedBrowser
+from .utils import get_cookie_loader as get_cookie_loader
+from .utils import load_cookies as load_cookies

revChatGPTAuth/__init__.py

@@ -1,14 +1,18 @@
+from typing import Union
+
 import httpx
 
+from revChatGPTAuth.supported_browser import SupportedBrowser
+
 from .openai_cookie_parser import OpenAICookieParser
 
 
 class ChatGPTAccessTokenParser:
 
-    def __init__(self, browser_name: str):
+    def __init__(self, browser_name: SupportedBrowser):
         self.openai_cookie_parser = OpenAICookieParser(browser_name)
 
-    def get_openai_chatgpt_access_token(self) -> str:
+    def get_access_token(self) -> str:
         cookie = self._get_stringified_cookies()
         response = self._auth_openai(cookie)
         response.raise_for_status()
@@ -25,6 +29,8 @@ def _auth_openai(self, cookie: str) -> httpx.Response:
         return response
 
 
-def get_access_token(browser_name: str) -> str:
+def get_access_token(browser_name: Union[str, SupportedBrowser]) -> str:
+    if isinstance(browser_name, str):
+        browser_name = SupportedBrowser(browser_name)
     chatgpt_access_token_parser = ChatGPTAccessTokenParser(browser_name)
-    return chatgpt_access_token_parser.get_openai_chatgpt_access_token()
+    return chatgpt_access_token_parser.get_access_token()

revChatGPTAuth/chatgpt_access_token_parser.py

@@ -1,19 +1,19 @@
 from http.cookiejar import Cookie
 from typing import Any, Optional
 
-from yt_dlp.cookies import SUPPORTED_BROWSERS, extract_cookies_from_browser
-from yt_dlp.utils import YoutubeDLCookieJar
+from revChatGPTAuth.supported_browser import SupportedBrowser
+from revChatGPTAuth.utils import load_cookies
 
 
 class OpenAICookieParser:
 
-    def __init__(self, browser_name: str):
-        if browser_name not in SUPPORTED_BROWSERS:
-            raise ValueError(f'Browser {browser_name} is not supported. Supported browsers are: {SUPPORTED_BROWSERS}')
+    SUPPORTED_BROWSERS = [browser.value for browser in SupportedBrowser]
+
+    def __init__(self, browser_name: SupportedBrowser):
         self.BROWSER_NAME = browser_name
 
     def parse_cookie(self):
-        all_cookies: YoutubeDLCookieJar = extract_cookies_from_browser(self.BROWSER_NAME)
+        all_cookies = load_cookies(self.BROWSER_NAME)
         openai_cookies = self._get_openai_cookies(all_cookies.__dict__)
         openai_cookies_dict: dict[str, Optional[str]] = {key: cookie.value for key, cookie in openai_cookies.items()}
         return openai_cookies_dict

revChatGPTAuth/openai_cookie_parser.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class SupportedBrowser(Enum):
+    '''Supported browsers for cookie extraction.'''
+    CHROME = 'chrome'
+    FIREFOX = 'firefox'
+    OPERA = 'opera'
+    OPERA_GX = 'opera gx'
+    EDGE = 'edge'
+    CHROMIUM = 'chromium'
+    BRAVE = 'brave'
+    VIVALDI = 'vivaldi'
+    SAFARI = 'safari'

revChatGPTAuth/supported_browser.py

@@ -27,7 +27,7 @@ def test_get_openai_chatgpt_access_token(
         chatgpt_access_token_parser = ChatGPTAccessTokenParser('chrome')
 
         # when
-        access_token = chatgpt_access_token_parser.get_openai_chatgpt_access_token()
+        access_token = chatgpt_access_token_parser.get_access_token()
 
         # then
         mock_cookie_parser.parse_cookie.assert_called_once()

revChatGPTAuth/test/test_chatgpt_access_token_parser.py

@@ -0,0 +1,35 @@
+from http.cookiejar import CookieJar
+from typing import Callable, Optional
+
+import browser_cookie3
+
+from revChatGPTAuth.supported_browser import SupportedBrowser
+
+
+def get_cookie_loader(browser_name: Optional[SupportedBrowser]) -> Callable[[str], CookieJar]:
+    if not browser_name:
+        return browser_cookie3.load
+    if browser_name == SupportedBrowser.CHROME:
+        return browser_cookie3.chrome
+    if browser_name == SupportedBrowser.FIREFOX:
+        return browser_cookie3.firefox
+    if browser_name == SupportedBrowser.OPERA:
+        return browser_cookie3.opera
+    if browser_name == SupportedBrowser.OPERA_GX:
+        return browser_cookie3.opera_gx
+    if browser_name == SupportedBrowser.EDGE:
+        return browser_cookie3.edge
+    if browser_name == SupportedBrowser.CHROMIUM:
+        return browser_cookie3.chromium
+    if browser_name == SupportedBrowser.BRAVE:
+        return browser_cookie3.brave
+    if browser_name == SupportedBrowser.VIVALDI:
+        return browser_cookie3.vivaldi
+    if browser_name == SupportedBrowser.SAFARI:
+        return browser_cookie3.safari
+    raise ValueError(f'Unsupported browser: {browser_name}')
+
+
+def load_cookies(browser_name: Optional[SupportedBrowser] = None, domain: str = '') -> CookieJar:
+    cookie_loader = get_cookie_loader(browser_name=browser_name)
+    return cookie_loader(domain)