From d52c9c5304e83274c69f295299afc7c725299798 Mon Sep 17 00:00:00 2001
From: Nick Sweeting <git@nicksweeting.com>
Date: Thu, 21 Feb 2019 12:57:16 -0500
Subject: [PATCH] allow passing COOKIES_FILE to wget

---
 archivebox/archive_methods.py | 13 +++++++++++--
 archivebox/config.py          |  9 ++++++++-
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/archivebox/archive_methods.py b/archivebox/archive_methods.py
index 26530d22..c7924318 100644
--- a/archivebox/archive_methods.py
+++ b/archivebox/archive_methods.py
@@ -25,6 +25,7 @@ from config import (
     RESOLUTION,
     CHECK_SSL_VALIDITY,
     SUBMIT_ARCHIVE_DOT_ORG,
+    COOKIES_FILE,
     WGET_USER_AGENT,
     CHROME_USER_DATA_DIR,
     CHROME_SANDBOX,
@@ -229,6 +230,7 @@ def fetch_wget(link_dir, link, requisites=FETCH_WGET_REQUISITES, warc=FETCH_WARC
         *(('--warc-file={}'.format(warc_path),) if warc else ()),
         *(('--page-requisites',) if FETCH_WGET_REQUISITES else ()),
         *(('--user-agent={}'.format(WGET_USER_AGENT),) if WGET_USER_AGENT else ()),
+        *(('--load-cookies', COOKIES_FILE) if COOKIES_FILE else ()),
         *((() if CHECK_SSL_VALIDITY else ('--no-check-certificate', '--no-hsts'))),
         link['url'],
     ]
@@ -260,12 +262,19 @@ def fetch_wget(link_dir, link, requisites=FETCH_WGET_REQUISITES, warc=FETCH_WARC
             raise Exception('Got an error from the server')
     except Exception as e:
         end()
+
+        # to let the user copy-paste the command and run it safely we have
+        # to quote some of the arguments that could have spaces in them
+        quoted_cmd = ' '.join(CMD)
+        quoted_cmd = quoted_cmd.replace(WGET_USER_AGENT, '"{}"'.format(WGET_USER_AGENT))
+        if COOKIES_FILE:
+            quoted_cmd = quoted_cmd.replace(COOKIES_FILE, '"{}"'.format(COOKIES_FILE))
+
         print('        {}Some resources were skipped: {}{}'.format(ANSI['lightyellow'], e, ANSI['reset']))
         print('        Run to see full output:')
         print('            cd {};'.format(link_dir))
-        print('            {}'.format(' '.join(CMD).replace(WGET_USER_AGENT, '"{}"'.format(WGET_USER_AGENT))))
+        print('            {}'.format(quoted_cmd))
         output = e
-
     return {
         'cmd': CMD,
         'output': output,
diff --git a/archivebox/config.py b/archivebox/config.py
index 1202fd3c..fec3411e 100644
--- a/archivebox/config.py
+++ b/archivebox/config.py
@@ -33,8 +33,9 @@ SUBMIT_ARCHIVE_DOT_ORG = os.getenv('SUBMIT_ARCHIVE_DOT_ORG', 'True'
 CHECK_SSL_VALIDITY =     os.getenv('CHECK_SSL_VALIDITY',     'True'             ).lower() == 'true'
 RESOLUTION =             os.getenv('RESOLUTION',             '1440,2000'        )
 GIT_DOMAINS =            os.getenv('GIT_DOMAINS',            'github.com,bitbucket.org,gitlab.com').split(',')
+COOKIES_FILE =           os.getenv('COOKIES_FILE',           None)
 WGET_USER_AGENT =        os.getenv('WGET_USER_AGENT',        'ArchiveBox/{GIT_SHA} (+https://github.com/pirate/ArchiveBox/) wget/{WGET_VERSION}')
-CHROME_USER_DATA_DIR =   os.getenv('CHROME_USER_DATA_DIR',    None)
+CHROME_USER_DATA_DIR =   os.getenv('CHROME_USER_DATA_DIR',   None)
 
 CHROME_BINARY =          os.getenv('CHROME_BINARY',          None)  # change to google-chrome browser if using google-chrome
 WGET_BINARY =            os.getenv('WGET_BINARY',            'wget'             )
@@ -122,6 +123,12 @@ except Exception:
 
 WGET_USER_AGENT = WGET_USER_AGENT.format(GIT_SHA=GIT_SHA[:9], WGET_VERSION=WGET_VERSION)
 
+try:
+    COOKIES_FILE = os.path.abspath(COOKIES_FILE) if COOKIES_FILE else None
+except Exception:
+    print('[!] Warning: unable to get full path to COOKIES_FILE, are you sure you specified it correctly?')
+    raise
+
 if sys.stdout.encoding.upper() not in ('UTF-8', 'UTF8'):
     print('[X] Your system is running python3 scripts with a bad locale setting: {} (it should be UTF-8).'.format(sys.stdout.encoding))
     print('    To fix it, add the line "export PYTHONIOENCODING=UTF-8" to your ~/.bashrc file (without quotes)')