Handle headers in State

fbchat/_client.py

@@ -78,24 +78,16 @@ class Client(object):
         :raises: FBchatException on failed login
         """
         self._sticky, self._pool = (None, None)
-        self._resetValues()
+        self._state = State.with_user_agent(user_agent=user_agent)
+        self._seq = "0"
+        self._uid = None
+        self._client_id = hex(int(random() * 2 ** 31))[2:]
         self._default_thread_id = None
         self._default_thread_type = None
         self._pull_channel = 0
         self._markAlive = True
         self._buddylist = dict()
 
-        if not user_agent:
-            user_agent = choice(USER_AGENTS)
-
-        self._header = {
-            "Content-Type": "application/x-www-form-urlencoded",
-            "Referer": "https://www.facebook.com",
-            "Origin": "https://www.facebook.com",
-            "User-Agent": user_agent,
-            "Connection": "keep-alive",
-        }
-
         handler.setLevel(logging_level)
 
         # If session cookies aren't set, not properly loaded or gives us an invalid session, then do the login
@@ -139,10 +131,7 @@ class Client(object):
     ):
         payload = self._generatePayload(query)
         r = self._state._session.get(
-            prefix_url(url),
-            headers=self._header,
-            params=payload,
-            allow_redirects=allow_redirects,
+            prefix_url(url), params=payload, allow_redirects=allow_redirects
         )
         if not fix_request:
             return r
@@ -170,9 +159,7 @@ class Client(object):
         error_retries=3,
     ):
         payload = self._generatePayload(query)
-        r = self._state._session.post(
-            prefix_url(url), headers=self._header, data=payload
-        )
+        r = self._state._session.post(prefix_url(url), data=payload)
         if not fix_request:
             return r
         try:
@@ -203,13 +190,7 @@ class Client(object):
         error_retries=3,
     ):
         payload = self._generatePayload(query)
-        # Removes 'Content-Type' from the header
-        headers = dict(
-            (i, self._header[i]) for i in self._header if i != "Content-Type"
-        )
-        r = self._state._session.post(
-            prefix_url(url), headers=headers, data=payload, files=files
-        )
+        r = self._state._session.post(prefix_url(url), data=payload, files=files)
         if not fix_request:
             return r
         try:
@@ -260,12 +241,6 @@ class Client(object):
     LOGIN METHODS
     """
 
-    def _resetValues(self):
-        self._state = State()
-        self._seq = "0"
-        self._uid = None
-        self._client_id = hex(int(random() * 2147483648))[2:]
-
     def _login(self, email, password):
         soup = bs(self._get("https://m.facebook.com/").text, "html.parser")
         data = dict(
@@ -391,7 +366,9 @@ class Client(object):
 
         try:
             # Load cookies into current session
-            self._state = State.from_cookies(session_cookies)
+            self._state = State.from_cookies(
+                session_cookies, user_agent=self._state._session.headers["User-Agent"]
+            )
         except Exception as e:
             log.exception("Failed loading session")
             self._state = State()
@@ -455,7 +432,8 @@ class Client(object):
         data = {"ref": "mb", "h": logout_h}
 
         r = self._get("/logout.php", data)
-        self._resetValues()
+        self._state = None
+        self._uid = None
         return r.ok
 
     """

fbchat/_state.py

@@ -5,17 +5,26 @@ import attr
 import bs4
 import re
 import requests
+import random
 
 from . import _util
 
 FB_DTSG_REGEX = re.compile(r'name="fb_dtsg" value="(.*?)"')
 
 
+def session_factory(user_agent=None):
+    session = requests.session()
+    session.headers["Referer"] = "https://www.facebook.com"
+    # TODO: Deprecate setting the user agent manually
+    session.headers["User-Agent"] = user_agent or random.choice(_util.USER_AGENTS)
+    return session
+
+
 @attr.s(slots=True, kw_only=True)
 class State(object):
     """Stores and manages state required for most Facebook requests."""
 
-    _session = attr.ib(factory=requests.session)
+    _session = attr.ib(factory=session_factory)
     fb_dtsg = attr.ib(None)
     _revision = attr.ib(None)
     _counter = attr.ib(0)
@@ -68,7 +77,11 @@ class State(object):
         return self._session.cookies.get_dict()
 
     @classmethod
-    def from_cookies(cls, cookies):
-        session = requests.session()
+    def from_cookies(cls, cookies, user_agent=None):
+        session = session_factory(user_agent=user_agent)
         session.cookies = requests.cookies.merge_cookies(session.cookies, cookies)
         return cls.from_session(session=session)
+
+    @classmethod
+    def with_user_agent(cls, user_agent=None, **kwargs):
+        return cls(session=session_factory(user_agent=user_agent), **kwargs)