From 8d2fe3d7d55bfa8a3064ec40433dda5ce97f20d7 Mon Sep 17 00:00:00 2001
From: Joakim Holm <joakimholm@protonmail.com>
Date: Wed, 19 Apr 2023 22:53:33 +0200
Subject: [PATCH] Add authentication with netscape cookie file

---
 grawlix/__main__.py       | 21 ++++++++++++++++++++-
 grawlix/arguments.py      |  6 ++++++
 grawlix/sources/source.py | 36 +++++++++++++++++++++++++++++++++++-
 3 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/grawlix/__main__.py b/grawlix/__main__.py
index 6052ad1..55b33ac 100644
--- a/grawlix/__main__.py
+++ b/grawlix/__main__.py
@@ -5,9 +5,10 @@ from .sources import load_source, Source
 from .output import download_book
 from . import  arguments, logging
 
-from typing import Tuple
+from typing import Tuple, Optional
 from rich.progress import Progress
 from functools import partial
+import os
 
 
 def get_login(source: Source, config: Config, options) -> Tuple[str, str]:
@@ -45,6 +46,20 @@ def get_urls(options) -> list[str]:
     return urls
 
 
+def get_cookie_file(options) -> Optional[str]:
+    """
+    Get path to cookie file
+
+    :param options: Cli arguments
+    :returns: Path to cookie file
+    """
+    if options.cookie_file is not None and os.path.exists(options.cookie_file):
+        return options.cookie_file
+    if os.path.exists("./cookies.txt"):
+        return "./cookies.txt"
+    return None
+
+
 def authenticate(source: Source, config: Config, options):
     """
     Authenticate with source
@@ -58,6 +73,10 @@ def authenticate(source: Source, config: Config, options):
         username, password = get_login(source, config, options)
         source.login(username, password)
         source.authenticated = True
+    if source.supports_cookies:
+        cookie_file = get_cookie_file(options)
+        if cookie_file:
+            source.load_cookies(cookie_file)
     else:
         raise SourceNotAuthenticated
 
diff --git a/grawlix/arguments.py b/grawlix/arguments.py
index f4cdbf7..2ec087a 100644
--- a/grawlix/arguments.py
+++ b/grawlix/arguments.py
@@ -39,6 +39,12 @@ def parse_arguments():
         help = "Password for login",
         dest = "password",
     )
+    parser.add_argument(
+        '-c',
+        '--cookies',
+        help = "Path to netscape cookie file",
+        dest = "cookie_file"
+    )
     # Outputs
     parser.add_argument(
         '-o',
diff --git a/grawlix/sources/source.py b/grawlix/sources/source.py
index ac32744..1ddd9bc 100644
--- a/grawlix/sources/source.py
+++ b/grawlix/sources/source.py
@@ -1,7 +1,9 @@
 from grawlix.book import Book, Series, Result
 
-from typing import Generic, TypeVar, Tuple
+from typing import Generic, TypeVar, Tuple, Optional
+from http.cookiejar import MozillaCookieJar
 import requests
+import re
 
 T = TypeVar("T")
 
@@ -41,6 +43,25 @@ class Source(Generic[T]):
         raise NotImplementedError
 
 
+    @property
+    def supports_cookies(self) -> bool:
+        """Does the source support authentication with cookie file"""
+        return "cookies" in self._authentication_methods
+
+
+    def load_cookies(self, cookie_file: str):
+        """
+        Authenticate with source with netscape cookie file
+
+        :param cookie_file: Path to netscape cookie file
+        """
+        if self.supports_cookies:
+            cookie_jar = MozillaCookieJar()
+            cookie_jar.load(cookie_file, ignore_expires=True)
+            self._session.cookies.update(cookie_jar)
+            self.authenticated = True
+
+
     def download(self, url: str) -> Result[T]:
         """
         Download book metadata from source
@@ -59,3 +80,16 @@ class Source(Generic[T]):
         :returns: Downloaded book metadata
         """
         raise NotImplementedError
+
+
+    def get_match_index(self, url: str) -> Optional[int]:
+        """
+        Find the first regex in `self.match` that matches url
+
+        :param url: Url to match
+        :returns: Index of regex
+        """
+        for index, match in enumerate(self.match):
+            if re.match(match, url):
+                return index
+        return None