From ef8e7d4251cbd8a4b33d27052a2c91b54795f5a0 Mon Sep 17 00:00:00 2001
From: Mads Marquart <madsmtm@gmail.com>
Date: Thu, 25 Jul 2019 23:54:17 +0200
Subject: [PATCH] Move user id handling to State

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

diff --git a/fbchat/_client.py b/fbchat/_client.py
index f19f620..8d598e5 100644
--- a/fbchat/_client.py
+++ b/fbchat/_client.py
@@ -174,16 +174,11 @@ class Client(object):
         """
         try:
             # Load cookies into current session
-            state = State.from_cookies(session_cookies, user_agent=user_agent)
+            self._state = State.from_cookies(session_cookies, user_agent=user_agent)
+            self._uid = self._state.user_id
         except Exception as e:
             log.exception("Failed loading session")
             return False
-        uid = state.get_user_id()
-        if uid is None:
-            log.warning("Could not find c_user cookie")
-            return False
-        self._state = state
-        self._uid = uid
         return True
 
     def login(self, email, password, max_tries=5, user_agent=None):
@@ -209,23 +204,19 @@ class Client(object):
 
         for i in range(1, max_tries + 1):
             try:
-                state = State.login(
+                self._state = State.login(
                     email,
                     password,
                     on_2fa_callback=self.on2FACode,
                     user_agent=user_agent,
                 )
-                uid = state.get_user_id()
-                if uid is None:
-                    raise FBchatException("Could not find user id")
+                self._uid = self._state.user_id
             except Exception:
                 if i >= max_tries:
                     raise
                 log.exception("Attempt #{} failed, retrying".format(i))
                 time.sleep(1)
             else:
-                self._state = state
-                self._uid = uid
                 self.onLoggedIn(email=email)
                 break
 
diff --git a/fbchat/_state.py b/fbchat/_state.py
index fb7e0d7..1dbe020 100644
--- a/fbchat/_state.py
+++ b/fbchat/_state.py
@@ -12,6 +12,14 @@ from . import _graphql, _util, _exception
 FB_DTSG_REGEX = re.compile(r'name="fb_dtsg" value="(.*?)"')
 
 
+def get_user_id(session):
+    # TODO: Optimize this `.get_dict()` call!
+    rtn = session.cookies.get_dict().get("c_user")
+    if rtn is None:
+        raise _exception.FBchatException("Could not find user id")
+    return str(rtn)
+
+
 def find_input_fields(html):
     return bs4.BeautifulSoup(html, "html.parser", parse_only=bs4.SoupStrainer("input"))
 
@@ -91,18 +99,13 @@ def _2fa_helper(session, code, r):
 class State(object):
     """Stores and manages state required for most Facebook requests."""
 
+    user_id = attr.ib()
     fb_dtsg = attr.ib()
     _revision = attr.ib()
     _session = attr.ib(factory=session_factory)
     _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)
-
     def get_params(self):
         self._counter += 1  # TODO: Make this operation atomic / thread-safe
         return {
@@ -163,6 +166,9 @@ class State(object):
 
     @classmethod
     def from_session(cls, session):
+        # TODO: Automatically set user_id when the cookie changes in the session
+        user_id = get_user_id(session)
+
         r = session.get(_util.prefix_url("/"))
 
         soup = find_input_fields(r.text)
@@ -180,7 +186,11 @@ class State(object):
         logout_h = logout_h_element["value"] if logout_h_element else None
 
         return cls(
-            fb_dtsg=fb_dtsg, revision=revision, session=session, logout_h=logout_h
+            user_id=user_id,
+            fb_dtsg=fb_dtsg,
+            revision=revision,
+            session=session,
+            logout_h=logout_h,
         )
 
     def get_cookies(self):
@@ -197,6 +207,7 @@ class State(object):
         # It may be a bad idea to do this in an exception handler, if you have a better method, please suggest it!
         _util.log.warning("Refreshing state and resending request")
         new = State.from_session(session=self._session)
+        self.user_id = new.user_id
         self.fb_dtsg = new.fb_dtsg
         self._revision = new._revision
         self._counter = new._counter