From 903f009726c9a669bf7ecf624ec67a68a4e8dd82 Mon Sep 17 00:00:00 2001
From: madwind <madwind.cn@gmail.com>
Date: Wed, 24 Aug 2022 12:00:31 +0800
Subject: [PATCH] cookie_backup

---
 auto_sign_in.py         |  1 +
 ptsites/base/request.py |  1 +
 ptsites/executor.py     | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/auto_sign_in.py b/auto_sign_in.py
index 162d5b27..8176f7c3 100644
--- a/auto_sign_in.py
+++ b/auto_sign_in.py
@@ -21,6 +21,7 @@ class PluginAutoSignIn:
             'max_workers': {'type': 'integer'},
             'get_messages': {'type': 'boolean', 'default': True},
             'get_details': {'type': 'boolean', 'default': True},
+            'cookie_backup': {'type': 'boolean', 'default': True},
             'aipocr': {
                 'type': 'object',
                 'properties': {
diff --git a/ptsites/base/request.py b/ptsites/base/request.py
index 8838ea74..edafbb90 100644
--- a/ptsites/base/request.py
+++ b/ptsites/base/request.py
@@ -69,6 +69,7 @@ def request(self,
                 entry.fail_with_prefix('Detected CloudFlare DDoS-GUARD')
             elif response is not None and response.status_code != 200:
                 entry.fail_with_prefix(f'response.status_code={response.status_code}')
+            entry['session_cookie'] = (' '.join(list(map(lambda x: f'{x[0]}={x[1]};', self.session.cookies.items()))))
             return response
         except Exception as e:
             entry.fail_with_prefix(NetworkState.NETWORK_ERROR.value.format(url=url, error=e))
diff --git a/ptsites/executor.py b/ptsites/executor.py
index e627ca02..a53947e5 100644
--- a/ptsites/executor.py
+++ b/ptsites/executor.py
@@ -3,6 +3,8 @@
 import importlib
 import pathlib
 import pkgutil
+import threading
+import os
 
 from flexget import plugin
 from flexget.entry import Entry
@@ -14,6 +16,8 @@
 from .base.reseed import Reseed
 from .base.sign_in import SignIn
 
+lock = threading.Semaphore(1)
+
 
 def build_sign_in_schema() -> dict:
     module = None
@@ -37,6 +41,33 @@ def build_sign_in_entry(entry: SignInEntry, config: dict) -> None:
         raise plugin.PluginError(f"site: {entry['site_name']}, error: {e}")
 
 
+def save_cookie(entry):
+    file_name = 'cookies.bak'
+    site_name = entry['site_name']
+    session_cookie = entry.get('session_cookie')
+    if not session_cookie:
+        return
+    with lock:
+        mode = 'r+'
+        if not os.path.exists(file_name):
+            mode = 'w+'
+
+        with open(file_name, mode, encoding='utf-8') as f:
+            lines = f.readlines()
+            updated = False
+            cookie = f"{site_name}: '{session_cookie}'\n"
+            for i in range(len(lines)):
+                if lines[i].startswith(f'{site_name}:'):
+                    lines[i] = cookie
+                    updated = True
+                    break
+            if not updated:
+                lines.append(cookie)
+            lines.sort()
+            f.seek(0)
+            f.writelines(lines)
+
+
 def sign_in(entry: SignInEntry, config: dict) -> None:
     try:
         site_class = get_site_class(entry['class_name'])
@@ -68,6 +99,11 @@ def sign_in(entry: SignInEntry, config: dict) -> None:
             return
         if entry['details']:
             logger.info(f"site_name: {entry['site_name']}, details: {entry['details']}")
+
+    if config.get('cookie_backup', True):
+        if entry.failed:
+            return
+        save_cookie(entry)
     clean_entry_attr(entry)