implement 'cookies-update' option (#445)

mikf · Oct 19, 2019 · 389d2d7 · 389d2d7
1 parent fbc0a6a
commit 389d2d7
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 3 deletions.
diff --git a/docs/configuration.rst b/docs/configuration.rst
@@ -223,6 +223,16 @@ Description Source to read additional cookies from.
 =========== =====
 
 
+extractor.*.cookies-update
+--------------------------
+=========== =====
+Type        ``bool``
+Default     ``false``
+Description If `extractor.*.cookies`_ specifies a cookies.txt file, update its
+            contents with cookies received during data extraction.
+=========== =====
+
+
 extractor.*.proxy
 -----------------
 =========== =====

diff --git a/docs/gallery-dl.conf b/docs/gallery-dl.conf
@@ -5,6 +5,7 @@
         "postprocessors": null,
         "archive": null,
         "cookies": null,
+        "cookies-update": false,
         "proxy": null,
         "skip": true,
         "sleep": 0,

diff --git a/gallery_dl/extractor/common.py b/gallery_dl/extractor/common.py
@@ -18,7 +18,7 @@
 import threading
 import http.cookiejar
 from .message import Message
-from .. import config, text, exception, cloudflare
+from .. import config, text, util, exception, cloudflare
 
 
 class Extractor():
@@ -43,6 +43,7 @@ def __init__(self, match):
         self._retries = self.config("retries", 4)
         self._timeout = self.config("timeout", 30)
         self._verify = self.config("verify", True)
+        self._cookiefile = None
 
         if self._retries < 0:
             self._retries = float("inf")
@@ -174,20 +175,37 @@ def _init_cookies(self):
         if cookies:
             if isinstance(cookies, dict):
                 self._update_cookies_dict(cookies, self.cookiedomain)
-            else:
+            elif isinstance(cookies, str):
+                cookiefile = util.expand_path(cookies)
                 cookiejar = http.cookiejar.MozillaCookieJar()
                 try:
-                    cookiejar.load(cookies)
+                    cookiejar.load(cookiefile)
                 except OSError as exc:
                     self.log.warning("cookies: %s", exc)
                 else:
                     self.session.cookies.update(cookiejar)
+                    self._cookiefile = cookiefile
+            else:
+                self.log.warning(
+                    "expected 'dict' or 'str' value for 'cookies' option, "
+                    "got '%s' (%s)", cookies.__class__.__name__, cookies)
 
         cookies = cloudflare.cookies(self.category)
         if cookies:
             domain, cookies = cookies
             self._update_cookies_dict(cookies, domain)
 
+    def _store_cookies(self):
+        """Store the session's cookiejar in a cookies.txt file"""
+        if self._cookiefile and self.config("cookies-update", False):
+            cookiejar = http.cookiejar.MozillaCookieJar()
+            for cookie in self.session.cookies:
+                cookiejar.set_cookie(cookie)
+            try:
+                cookiejar.save(self.cookiefile)
+            except OSError as exc:
+                self.log.warning("cookies: %s", exc)
+
     def _update_cookies(self, cookies, *, domain=""):
         """Update the session's cookiejar with 'cookies'"""
         if isinstance(cookies, dict):

diff --git a/gallery_dl/job.py b/gallery_dl/job.py
@@ -271,6 +271,8 @@ def handle_finalize(self):
                 pp.finalize()
         if self.archive:
             self.archive.close()
+        if self.pathfmt:
+            self.extractor._store_cookies()
 
     def handle_skip(self):
         self.out.skip(self.pathfmt.path)