Refactor MQTT connecting

fbchat/_mqtt.py

@@ -22,46 +22,70 @@ def fetch_sequence_id(state):
         raise
 
 
+HOST = "edge-chat.facebook.com"
+
+
 @attr.s(slots=True)
 class Mqtt:
     _state = attr.ib()
     _mqtt = attr.ib()
+    _session_id = attr.ib()
 
-    @classmethod
-    def connect(cls, state, foreground):
-        mqtt = paho.mqtt.client.Client(
+    def __init__(self, state):
+        self._state = state
+
+        # Generate a random session ID between 1 and 9007199254740991
+        self._session_id = random.randint(1, 2 ** 53)
+
+        self._mqtt = paho.mqtt.client.Client(
             client_id="mqttwsclient",
             clean_session=True,
             protocol=paho.mqtt.client.MQTTv31,
             transport="websockets",
         )
-        mqtt.enable_logger()
+        self._mqtt.enable_logger()
+        # self._mqtt.max_inflight_messages_set(20)
+        # self._mqtt.max_queued_messages_set(0) # unlimited
+        # self._mqtt.message_retry_set(5)
+        # self._mqtt.reconnect_delay_set(min_delay=1, max_delay=1)
 
-        # Generate a random session ID between 1 and 9007199254740991
-        session_id = random.randint(1, 2 ** 53)
-        last_seq_id = fetch_sequence_id(state)
+        # TODO: Is region (lla | atn | odn | others?) important?
+        self._mqtt.ws_set_options(
+            path="/chat?sid={}".format(self._session_id), headers=self._create_headers
+        )
+        self._mqtt.tls_set()
+
+    def _create_headers(self, headers):
+        headers["Cookie"] = _util.get_cookie_header(self._state._session, HOST)
+        headers["User-Agent"] = self._state._session.headers["User-Agent"]
+        headers["Origin"] = "https://www.facebook.com"
+        headers["Host"] = HOST
+        return headers
+
+    def connect(self, foreground):
+        last_seq_id = fetch_sequence_id(self._state)
 
         messenger_sync_create_queue_payload = {
             "sync_api_version": 10,
             "max_deltas_able_to_process": 1000,
             "delta_batch_size": 500,
             "encoding": "JSON",
-            "entity_fbid": state.user_id,
+            "entity_fbid": self._state.user_id,
             "initial_titan_sequence_id": str(last_seq_id),
             "device_params": None,
         }
 
         username = {
             # The user ID
-            "u": state.user_id,
+            "u": self._state.user_id,
             # Session ID
-            "s": session_id,
+            "s": self._session_id,
             # Active status setting
             "chat_on": True,
             # foreground_state - Whether the window is focused
             "fg": foreground,
             # Can be any random ID
-            "d": state._client_id,
+            "d": self._state._client_id,
             # Application ID, taken from facebook.com
             "aid": 219994525426954,
             # MQTT extension by FB, allows making a SUBSCRIBE while CONNECTing
@@ -115,23 +139,11 @@ class Mqtt:
             "gas": None,
             "pack": [],
         }
-        mqtt.username_pw_set(_util.json_minimal(username))
+        self._mqtt.username_pw_set(_util.json_minimal(username))
 
-        headers = {
-            "Cookie": _util.get_cookie_header(state._session, "edge-chat.facebook.com"),
-            "User-Agent": state._session.headers["User-Agent"],
-            "Origin": "https://www.facebook.com",
-            "Host": "edge-chat.facebook.com",
-        }
-
-        # TODO: Is region (lla | atn | odn | others?) important?
-        mqtt.ws_set_options(path="/chat?sid={}".format(session_id), headers=headers)
-        mqtt.tls_set()
-        response_code = mqtt.connect("edge-chat.facebook.com", 443, keepalive=10)
+        response_code = self._mqtt.connect(HOST, 443, keepalive=10)
         # TODO: Handle response code
 
-        return cls(state=state, mqtt=mqtt)
-
     def listen(self, on_message):
         def real_on_message(client, userdata, message):
             on_message(message.topic, message.payload)