Use MozillaCookieJar instead of json to save cookies

wkentaro · Jan 20, 2024 · 30fa8e7 · 30fa8e7
1 parent f68cdfe
commit 30fa8e7
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 18 deletions.
diff --git a/gdown/cli.py b/gdown/cli.py
@@ -87,7 +87,7 @@ def main():
     parser.add_argument(
         "--no-cookies",
         action="store_true",
-        help="don't use cookies in ~/.cache/gdown/cookies.json",
+        help="don't use cookies in ~/.cache/gdown/cookies.txt",
     )
     parser.add_argument(
         "--no-check-certificate",

diff --git a/gdown/download.py b/gdown/download.py
@@ -1,4 +1,3 @@
-import json
 import os
 import os.path as osp
 import re
@@ -8,6 +7,7 @@
 import textwrap
 import time
 import urllib
+from http.cookiejar import MozillaCookieJar
 
 import requests
 import tqdm
@@ -62,12 +62,11 @@ def _get_session(proxy, use_cookies, user_agent, return_cookies_file=False):
         print("Using proxy:", proxy, file=sys.stderr)
 
     # Load cookies if exists
-    cookies_file = osp.join(home, ".cache/gdown/cookies.json")
-    if osp.exists(cookies_file) and use_cookies:
-        with open(cookies_file) as f:
-            cookies = json.load(f)
-        for k, v in cookies:
-            sess.cookies[k] = v
+    cookies_file = osp.join(home, ".cache/gdown/cookies.txt")
+    if use_cookies and osp.exists(cookies_file):
+        cookie_jar = MozillaCookieJar(cookies_file)
+        cookie_jar.load()
+        sess.cookies.update(cookie_jar)
 
     if return_cookies_file:
         return sess, cookies_file
@@ -206,16 +205,10 @@ def download(
             continue
 
         if use_cookies:
-            if not osp.exists(osp.dirname(cookies_file)):
-                os.makedirs(osp.dirname(cookies_file))
-            # Save cookies
-            with open(cookies_file, "w") as f:
-                cookies = [
-                    (k, v)
-                    for k, v in sess.cookies.items()
-                    if not k.startswith("download_warning_")
-                ]
-                json.dump(cookies, f, indent=2)
+            cookie_jar = MozillaCookieJar(cookies_file)
+            for cookie in sess.cookies:
+                cookie_jar.set_cookie(cookie)
+            cookie_jar.save()
 
         if "Content-Disposition" in res.headers:
             # This is the file