From b01b371c66d7d87a30606fa32544552a6df58d02 Mon Sep 17 00:00:00 2001
From: Mads Marquart <madsmtm@gmail.com>
Date: Wed, 26 Jun 2019 20:03:38 +0200
Subject: [PATCH] Refactor session cookie handling into State

---
 fbchat/_client.py | 32 ++++++++++++++------------------
 fbchat/_state.py  | 25 +++++++++++++++++++++----
 2 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/fbchat/_client.py b/fbchat/_client.py
index c859489..cc31737 100644
--- a/fbchat/_client.py
+++ b/fbchat/_client.py
@@ -115,8 +115,8 @@ class Client(object):
         It may be a bad idea to do this in an exception handler, if you have a better method, please suggest it!
         """
         if error_code == "1357004":
-            log.warning("Got error #1357004. Doing a _postLogin, and resending request")
-            self._postLogin()
+            log.warning("Got error #1357004. Refreshing state and resending request")
+            self._state = State.from_session(session=self._state._session)
             return True
         return False
 
@@ -263,15 +263,6 @@ class Client(object):
         self._uid = None
         self._client_id = hex(int(random() * 2147483648))[2:]
 
-    def _postLogin(self):
-        self._uid = self._state._session.cookies.get_dict().get("c_user")
-        if self._uid is None:
-            raise FBchatException("Could not find c_user cookie")
-        self._uid = str(self._uid)
-
-        r = self._get("/")
-        self._state = State.from_base_request(self._state._session, r.text)
-
     def _login(self, email, password):
         soup = bs(self._get("https://m.facebook.com/").text, "html.parser")
         data = dict(
@@ -294,7 +285,10 @@ class Client(object):
             r = self._get("https://m.facebook.com/login/save-device/cancel/")
 
         if "home" in r.url:
-            self._postLogin()
+            self._state = State.from_session(session=self._state._session)
+            self._uid = self._state.get_user_id()
+            if self._uid is None:
+                raise FBchatException("Could not find c_user cookie")
             return True, r.url
         else:
             return False, r.url
@@ -378,7 +372,7 @@ class Client(object):
         :return: A dictionay containing session cookies
         :rtype: dict
         """
-        return self._state._session.cookies.get_dict()
+        return self._state.get_cookies()
 
     def setSession(self, session_cookies):
         """Loads session cookies
@@ -394,14 +388,16 @@ class Client(object):
 
         try:
             # Load cookies into current session
-            self._state._session.cookies = requests.cookies.merge_cookies(
-                self._state._session.cookies, session_cookies
-            )
-            self._postLogin()
+            self._state = State.from_cookies(session_cookies)
         except Exception as e:
             log.exception("Failed loading session")
-            self._resetValues()
+            self._state = State()
             return False
+        uid = self._state.get_user_id()
+        if uid is None:
+            log.warning("Could not find c_user cookie")
+            return False
+        self._uid = uid
         return True
 
     def login(self, email, password, max_tries=5):
diff --git a/fbchat/_state.py b/fbchat/_state.py
index 6822f95..a6f321d 100644
--- a/fbchat/_state.py
+++ b/fbchat/_state.py
@@ -21,6 +21,12 @@ class State(object):
     _counter = attr.ib(0)
     _logout_h = attr.ib(None)
 
+    def get_user_id(self):
+        rtn = self.get_cookies().get("c_user")
+        if rtn is None:
+            return None
+        return str(rtn)
+
     @property
     def logout_h(self):
         return self._logout_h
@@ -37,17 +43,19 @@ class State(object):
         }
 
     @classmethod
-    def from_base_request(cls, session, content):
-        soup = bs4.BeautifulSoup(content, "html.parser")
+    def from_session(cls, session):
+        r = session.get(_util.prefix_url("/"))
+
+        soup = bs4.BeautifulSoup(r.text, "html.parser")
 
         fb_dtsg_element = soup.find("input", {"name": "fb_dtsg"})
         if fb_dtsg_element:
             fb_dtsg = fb_dtsg_element["value"]
         else:
             # Fall back to searching with a regex
-            fb_dtsg = FB_DTSG_REGEX.search(content).group(1)
+            fb_dtsg = FB_DTSG_REGEX.search(r.text).group(1)
 
-        revision = int(content.split('"client_revision":', 1)[1].split(",", 1)[0])
+        revision = int(r.text.split('"client_revision":', 1)[1].split(",", 1)[0])
 
         logout_h_element = soup.find("input", {"name": "h"})
         logout_h = logout_h_element["value"] if logout_h_element else None
@@ -55,3 +63,12 @@ class State(object):
         return cls(
             session=session, fb_dtsg=fb_dtsg, revision=revision, logout_h=logout_h
         )
+
+    def get_cookies(self):
+        return self._session.cookies.get_dict()
+
+    @classmethod
+    def from_cookies(cls, cookies):
+        session = requests.session()
+        session.cookies = requests.cookies.merge_cookies(session.cookies, cookies)
+        return cls.from_session(session=session)