From 248830752dea79cd9af875ee4b4e43da7540fc5c Mon Sep 17 00:00:00 2001
From: Romain Bignon <romain@symlink.me>
Date: Sat, 7 Mar 2015 15:43:35 +0100
Subject: [PATCH] add (deprecated) StateBrowser which saves cookies into
 storage

---
 weboob/deprecated/browser/__init__.py |  5 +--
 weboob/deprecated/browser/browser.py  | 47 ++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/weboob/deprecated/browser/__init__.py b/weboob/deprecated/browser/__init__.py
index 66a26125..e46bbda7 100644
--- a/weboob/deprecated/browser/__init__.py
+++ b/weboob/deprecated/browser/__init__.py
@@ -23,9 +23,10 @@ from weboob.deprecated.browser.browser import BrowserIncorrectPassword, BrowserB
                                          BrowserHTTPNotFound, BrowserHTTPError, \
                                          Page, Browser, BrokenPageError, \
                                          StandardBrowser, BrowserPasswordExpired, \
-                                         BrowserForbidden
+                                         BrowserForbidden, StateBrowser
 
 
 __all__ = ['BrowserIncorrectPassword', 'BrowserPasswordExpired', 'BrowserBanned',
            'BrowserUnavailable', 'BrowserRetry', 'BrowserHTTPNotFound', 'BrowserHTTPError',
-           'Page', 'Browser', 'BrokenPageError', 'StandardBrowser', 'BrowserForbidden']
+           'Page', 'Browser', 'BrokenPageError', 'StandardBrowser', 'BrowserForbidden',
+           'StateBrowser']
diff --git a/weboob/deprecated/browser/browser.py b/weboob/deprecated/browser/browser.py
index 3803771e..a213a2c2 100644
--- a/weboob/deprecated/browser/browser.py
+++ b/weboob/deprecated/browser/browser.py
@@ -33,6 +33,9 @@ try:
 except ImportError:
     raise ImportError('Please install python-mechanize')
 
+import pickle
+import base64
+import zlib
 import os
 import re
 from threading import RLock
@@ -58,7 +61,7 @@ from weboob.deprecated.browser.parsers import get_parser
 
 __all__ = ['BrowserIncorrectPassword', 'BrowserForbidden', 'BrowserBanned', 'BrowserUnavailable', 'BrowserRetry',
            'BrowserPasswordExpired', 'BrowserHTTPNotFound', 'BrowserHTTPError', 'BrokenPageError', 'Page',
-           'StandardBrowser', 'Browser']
+           'StandardBrowser', 'Browser', 'StateBrowser']
 
 
 class BrowserRetry(Exception):
@@ -692,6 +695,48 @@ class Browser(StandardBrowser):
             self._cookie.save()
 
 
+class StateBrowser(Browser):
+    """
+    This browser aims to store state (cookies, location and attributes).
+    """
+
+    __states__ = []
+    """
+    Saved state variables.
+    """
+
+    def __init__(self, *args, **kwargs):
+        kwargs['get_home'] = False
+        Browser.__init__(self, *args, **kwargs)
+
+    def load_state(self, state):
+        if 'cookies' in state:
+            try:
+                cookies = pickle.loads(zlib.decompress(base64.b64decode(state['cookies'])))
+            except (TypeError, zlib.error, EOFError, ValueError):
+                self.logger.error('Unable to reload cookies from storage')
+            else:
+                self.set_cookiejar(cookies)
+                self.logger.info('Reloaded cookies from storage')
+        for attrname in self.__states__:
+            if attrname in state:
+                setattr(self, attrname, state[attrname])
+        if 'url' in state:
+            self.location(state['url'], nologin=True)
+
+
+    def dump_state(self):
+        state = {}
+        cookiejar = self._ua_handlers["_cookies"].cookiejar
+        state['cookies'] = base64.b64encode(zlib.compress(pickle.dumps(cookiejar, -1)))
+        for attrname in self.__states__:
+            state[attrname] = getattr(self, attrname)
+        if self.page:
+            state['url'] = self.page.url
+        self.logger.info('Stored cookies into storage')
+        return state
+
+
 def mywrap_socket(sock, *args, **kwargs):
     kwargs['do_handshake_on_connect'] = False
     sock = ssl.wrap_socketold(sock, *args, **kwargs)