From cedfb5670e48f41d635317f7f2bed5587a656009 Mon Sep 17 00:00:00 2001
From: Romain Bignon <romain@symlink.me>
Date: Sat, 21 Feb 2015 14:21:54 +0100
Subject: [PATCH] save and load state of browser into storage

---
 modules/creditmutuel/browser.py        |  6 +++-
 weboob/applications/boobank/boobank.py |  3 ++
 weboob/browser/browsers.py             | 43 ++++++++++++++++++++++++--
 weboob/tools/backend.py                | 14 +++++++--
 4 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/modules/creditmutuel/browser.py b/modules/creditmutuel/browser.py
index 1a118d02..7c7be832 100644
--- a/modules/creditmutuel/browser.py
+++ b/modules/creditmutuel/browser.py
@@ -45,7 +45,9 @@ class CreditMutuelBrowser(LoginBrowser):
     PROFILE = Wget()
     BASEURL = 'https://www.creditmutuel.fr'
 
-    login =       URL('/groupe/fr/index.html',                               LoginPage)
+    login =       URL('/groupe/fr/index.html',
+                      '/(?P<subbank>.*)/fr/banques/accueil.html',
+                      LoginPage)
     login_error = URL('/(?P<subbank>.*)/fr/identification/default.cgi',      LoginErrorPage)
     accounts =    URL('/(?P<subbank>.*)/fr/banque/situation_financiere.cgi', AccountsPage)
     user_space =  URL('/(?P<subbank>.*)/fr/banque/espace_personnel.aspx',    UserSpacePage)
@@ -70,6 +72,8 @@ class CreditMutuelBrowser(LoginBrowser):
 
     currentSubBank = None
 
+    __states__ = ['currentSubBank']
+
     def do_login(self):
         self.login.stay_or_go()
 
diff --git a/weboob/applications/boobank/boobank.py b/weboob/applications/boobank/boobank.py
index 6d7960c3..e0cef259 100644
--- a/weboob/applications/boobank/boobank.py
+++ b/weboob/applications/boobank/boobank.py
@@ -320,6 +320,9 @@ class Boobank(ReplApplication):
                            }
     COLLECTION_OBJECTS = (Account, Transaction, )
 
+    def load_default_backends(self):
+        self.load_backends(CapBank, storage=self.create_storage())
+
     def _complete_account(self, exclude=None):
         if exclude:
             exclude = '%s@%s' % self.parse_id(exclude)
diff --git a/weboob/browser/browsers.py b/weboob/browser/browsers.py
index 41602244..43ddfa96 100644
--- a/weboob/browser/browsers.py
+++ b/weboob/browser/browsers.py
@@ -20,6 +20,9 @@
 from __future__ import absolute_import, print_function
 
 import re
+import pickle
+import base64
+import zlib
 try:
     import urllib3
 except ImportError:
@@ -88,6 +91,11 @@ class Browser(object):
     Maximum of threads for asynchronous requests.
     """
 
+    __states__ = []
+    """
+    Saved state variables.
+    """
+
     @classmethod
     def asset(cls, localfile):
         """
@@ -110,7 +118,27 @@ class Browser(object):
         if isinstance(self.VERIFY, basestring):
             self.VERIFY = self.asset(self.VERIFY)
 
-    def _save(self, response, warning=False, **kwargs):
+    def load_state(self, state):
+        if 'cookies' in state:
+            try:
+                self.session.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.logger.info('Reloaded cookies from storage')
+        for attrname in self.__states__:
+            if attrname in state:
+                setattr(self, attrname, state[attrname])
+
+    def dump_state(self):
+        state = {}
+        state['cookies'] = base64.b64encode(zlib.compress(pickle.dumps(self.session.cookies, -1)))
+        for attrname in self.__states__:
+            state[attrname] = getattr(self, attrname)
+        self.logger.info('Stored cookies into storage')
+        return state
+
+    def save_response(self, response, warning=False, **kwargs):
         if self.responses_dirname is None:
             import tempfile
             self.responses_dirname = tempfile.mkdtemp(prefix='weboob_session_')
@@ -194,7 +222,7 @@ class Browser(object):
         profile.setup_session(session)
 
         if self.logger.settings['save_responses']:
-            session.hooks['response'].append(self._save)
+            session.hooks['response'].append(self.save_response)
 
         self.session = session
 
@@ -606,6 +634,17 @@ class PagesBrowser(DomainBrowser):
         for url in self._urls.itervalues():
             url.browser = self
 
+    def load_state(self, state):
+        super(PagesBrowser, self).load_state(state)
+        if 'url' in state:
+            self.location(state['url'])
+
+    def dump_state(self):
+        state = super(PagesBrowser, self).dump_state()
+        if self.page:
+            state['url'] = self.page.url
+        return state
+
     def open(self, *args, **kwargs):
         """
         Same method than
diff --git a/weboob/tools/backend.py b/weboob/tools/backend.py
index 2802a7e7..3c4467c3 100644
--- a/weboob/tools/backend.py
+++ b/weboob/tools/backend.py
@@ -285,7 +285,12 @@ class Module(object):
         """
         This abstract method is called when the backend is unloaded.
         """
-        pass
+        if self._browser is None:
+            return
+
+        if hasattr(self.browser, 'dump_state'):
+            self.storage.set('browser_state', self.browser.dump_state())
+            self.storage.save()
 
     _browser = None
 
@@ -346,7 +351,12 @@ class Module(object):
             kwargs.setdefault('responses_dirname', os.path.join(self.logger.settings['responses_dirname'],
                                                                 self._private_config.get('_debug_dir', self.name)))
 
-        return self.BROWSER(*args, **kwargs)
+        browser = self.BROWSER(*args, **kwargs)
+
+        if hasattr(browser, 'load_state'):
+            browser.load_state(self.storage.get('browser_state', default={}))
+
+        return browser
 
     @classmethod
     def iter_caps(klass):