From 389d2d7e381b49c50dca49728ace499e33711266 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mike=20F=C3=A4hrmann?= <mike_faehrmann@web.de>
Date: Fri, 18 Oct 2019 21:31:33 +0200
Subject: [PATCH] implement 'cookies-update' option (#445)

---
 docs/configuration.rst         | 10 ++++++++++
 docs/gallery-dl.conf           |  1 +
 gallery_dl/extractor/common.py | 24 +++++++++++++++++++++---
 gallery_dl/job.py              |  2 ++
 4 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/docs/configuration.rst b/docs/configuration.rst
index 39a3b9e1..ac1935bd 100644
--- 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
index ebf47ffe..7fd6f9ec 100644
--- 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
index 0bb12777..8cb2abe6 100644
--- a/gallery_dl/extractor/common.py
+++ b/gallery_dl/extractor/common.py
@@ -18,7 +18,7 @@ import requests
 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 @@ class Extractor():
         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 @@ class Extractor():
         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
index 8b610248..b8eef639 100644
--- a/gallery_dl/job.py
+++ b/gallery_dl/job.py
@@ -271,6 +271,8 @@ class DownloadJob(Job):
                 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)