conn.py 38 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963
  1. # mautrix-instagram - A Matrix-Instagram puppeting bridge.
  2. # Copyright (C) 2022 Tulir Asokan
  3. #
  4. # This program is free software: you can redistribute it and/or modify
  5. # it under the terms of the GNU Affero General Public License as published by
  6. # the Free Software Foundation, either version 3 of the License, or
  7. # (at your option) any later version.
  8. #
  9. # This program is distributed in the hope that it will be useful,
  10. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  11. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  12. # GNU Affero General Public License for more details.
  13. #
  14. # You should have received a copy of the GNU Affero General Public License
  15. # along with this program. If not, see <https://www.gnu.org/licenses/>.
  16. from __future__ import annotations
  17. from typing import Any, Awaitable, Callable, Iterable, Type, TypeVar
  18. from collections import defaultdict
  19. from socket import error as SocketError, socket
  20. import asyncio
  21. import json
  22. import logging
  23. import re
  24. import time
  25. import zlib
  26. from yarl import URL
  27. import paho.mqtt.client as pmc
  28. from mautrix.util.logging import TraceLogger
  29. from ..errors import (
  30. IrisSubscribeError,
  31. MQTTConnectionUnauthorized,
  32. MQTTNotConnected,
  33. MQTTNotLoggedIn,
  34. )
  35. from ..proxy import ProxyHandler
  36. from ..state import AndroidState
  37. from ..types import (
  38. AppPresenceEventPayload,
  39. ClientConfigUpdatePayload,
  40. CommandResponse,
  41. IrisPayload,
  42. IrisPayloadData,
  43. LiveVideoCommentPayload,
  44. MessageSyncEvent,
  45. MessageSyncMessage,
  46. Operation,
  47. PubsubEvent,
  48. PubsubPayload,
  49. ReactionStatus,
  50. ReactionType,
  51. RealtimeDirectEvent,
  52. RealtimeZeroProvisionPayload,
  53. ThreadAction,
  54. ThreadItemType,
  55. ThreadRemoveEvent,
  56. ThreadSyncEvent,
  57. TypingStatus,
  58. )
  59. from .events import Connect, Disconnect, NewSequenceID, ProxyUpdate
  60. from .otclient import MQTToTClient
  61. from .subscription import GraphQLQueryID, RealtimeTopic, everclear_subscriptions
  62. from .thrift import ForegroundStateConfig, IncomingMessage, RealtimeClientInfo, RealtimeConfig
  63. try:
  64. import socks
  65. except ImportError:
  66. socks = None
  67. T = TypeVar("T")
  68. ACTIVITY_INDICATOR_REGEX = re.compile(
  69. r"/direct_v2/threads/([\w_]+)/activity_indicator_id/([\w_]+)"
  70. )
  71. INBOX_THREAD_REGEX = re.compile(r"/direct_v2/inbox/threads/([\w_]+)")
  72. class AndroidMQTT:
  73. _loop: asyncio.AbstractEventLoop
  74. _client: MQTToTClient
  75. log: TraceLogger
  76. state: AndroidState
  77. _graphql_subs: set[str]
  78. _skywalker_subs: set[str]
  79. _iris_seq_id: int | None
  80. _iris_snapshot_at_ms: int | None
  81. _publish_waiters: dict[int, asyncio.Future]
  82. _response_waiters: dict[RealtimeTopic, asyncio.Future]
  83. _response_waiter_locks: dict[RealtimeTopic, asyncio.Lock]
  84. _message_response_waiter_lock: asyncio.Lock
  85. _message_response_waiter_id: str | None
  86. _message_response_waiter: asyncio.Future | None
  87. _disconnect_error: Exception | None
  88. _event_handlers: dict[Type[T], list[Callable[[T], Awaitable[None]]]]
  89. _outgoing_events: asyncio.Queue
  90. _event_dispatcher_task: asyncio.Task | None
  91. # region Initialization
  92. def __init__(
  93. self,
  94. state: AndroidState,
  95. log: TraceLogger | None = None,
  96. proxy_handler: ProxyHandler | None = None,
  97. ) -> None:
  98. self._graphql_subs = set()
  99. self._skywalker_subs = set()
  100. self._iris_seq_id = None
  101. self._iris_snapshot_at_ms = None
  102. self._publish_waiters = {}
  103. self._response_waiters = {}
  104. self._message_response_waiter_lock = asyncio.Lock()
  105. self._message_response_waiter_id = None
  106. self._message_response_waiter = None
  107. self._disconnect_error = None
  108. self._response_waiter_locks = defaultdict(lambda: asyncio.Lock())
  109. self._event_handlers = defaultdict(lambda: [])
  110. self._event_dispatcher_task = None
  111. self._outgoing_events = asyncio.Queue()
  112. self.log = log or logging.getLogger("mauigpapi.mqtt")
  113. self._loop = asyncio.get_running_loop()
  114. self.state = state
  115. self._client = MQTToTClient(
  116. client_id=self._form_client_id(),
  117. clean_session=True,
  118. protocol=pmc.MQTTv31,
  119. transport="tcp",
  120. )
  121. self.proxy_handler = proxy_handler
  122. self.setup_proxy()
  123. self._client.enable_logger()
  124. self._client.tls_set()
  125. # mqtt.max_inflight_messages_set(20) # The rest will get queued
  126. # mqtt.max_queued_messages_set(0) # Unlimited messages can be queued
  127. # mqtt.message_retry_set(20) # Retry sending for at least 20 seconds
  128. # mqtt.reconnect_delay_set(min_delay=1, max_delay=120)
  129. self._client.connect_async("edge-mqtt.facebook.com", 443, keepalive=60)
  130. self._client.on_message = self._on_message_handler
  131. self._client.on_publish = self._on_publish_handler
  132. self._client.on_connect = self._on_connect_handler
  133. self._client.on_disconnect = self._on_disconnect_handler
  134. self._client.on_socket_open = self._on_socket_open
  135. self._client.on_socket_close = self._on_socket_close
  136. self._client.on_socket_register_write = self._on_socket_register_write
  137. self._client.on_socket_unregister_write = self._on_socket_unregister_write
  138. def setup_proxy(self):
  139. http_proxy = self.proxy_handler.get_proxy_url()
  140. if http_proxy:
  141. if not socks:
  142. self.log.warning("http_proxy is set, but pysocks is not installed")
  143. else:
  144. proxy_url = URL(http_proxy)
  145. proxy_type = {
  146. "http": socks.HTTP,
  147. "https": socks.HTTP,
  148. "socks": socks.SOCKS5,
  149. "socks5": socks.SOCKS5,
  150. "socks4": socks.SOCKS4,
  151. }[proxy_url.scheme]
  152. self._client.proxy_set(
  153. proxy_type=proxy_type,
  154. proxy_addr=proxy_url.host,
  155. proxy_port=proxy_url.port,
  156. proxy_username=proxy_url.user,
  157. proxy_password=proxy_url.password,
  158. )
  159. def _clear_response_waiters(self) -> None:
  160. for waiter in self._response_waiters.values():
  161. if not waiter.done():
  162. waiter.set_exception(
  163. MQTTNotConnected("MQTT disconnected before request returned response")
  164. )
  165. for waiter in self._publish_waiters.values():
  166. if not waiter.done():
  167. waiter.set_exception(
  168. MQTTNotConnected("MQTT disconnected before request was published")
  169. )
  170. if self._message_response_waiter and not self._message_response_waiter.done():
  171. self._message_response_waiter.set_exception(
  172. MQTTNotConnected("MQTT disconnected before message send returned response")
  173. )
  174. self._message_response_waiter = None
  175. self._message_response_waiter_id = None
  176. self._response_waiters = {}
  177. self._publish_waiters = {}
  178. def _form_client_id(self) -> bytes:
  179. subscribe_topics = [
  180. RealtimeTopic.PUBSUB, # 88
  181. RealtimeTopic.SUB_IRIS_RESPONSE, # 135
  182. RealtimeTopic.RS_REQ, # 244
  183. RealtimeTopic.REALTIME_SUB, # 149
  184. RealtimeTopic.REGION_HINT, # 150
  185. RealtimeTopic.RS_RESP, # 245
  186. RealtimeTopic.T_RTC_LOG, # 274
  187. RealtimeTopic.SEND_MESSAGE_RESPONSE, # 133
  188. RealtimeTopic.MESSAGE_SYNC, # 146
  189. RealtimeTopic.LIGHTSPEED_RESPONSE, # 179
  190. RealtimeTopic.UNKNOWN_PP, # 34
  191. ]
  192. subscribe_topic_ids = [int(topic.encoded) for topic in subscribe_topics]
  193. password = f"authorization={self.state.session.authorization}"
  194. cfg = RealtimeConfig(
  195. client_identifier=self.state.device.phone_id[:20],
  196. client_info=RealtimeClientInfo(
  197. user_id=int(self.state.user_id),
  198. user_agent=self.state.user_agent,
  199. client_capabilities=0b10110111,
  200. endpoint_capabilities=0,
  201. publish_format=1,
  202. no_automatic_foreground=True,
  203. make_user_available_in_foreground=False,
  204. device_id=self.state.device.phone_id,
  205. is_initially_foreground=False,
  206. network_type=1,
  207. network_subtype=-1,
  208. client_mqtt_session_id=int(time.time() * 1000) & 0xFFFFFFFF,
  209. subscribe_topics=subscribe_topic_ids,
  210. client_type="cookie_auth",
  211. app_id=567067343352427,
  212. # region_preference=self.state.session.region_hint or "LLA",
  213. device_secret="",
  214. client_stack=3,
  215. ),
  216. password=password,
  217. app_specific_info={
  218. "capabilities": self.state.application.CAPABILITIES,
  219. "app_version": self.state.application.APP_VERSION,
  220. "everclear_subscriptions": json.dumps(everclear_subscriptions),
  221. "User-Agent": self.state.user_agent,
  222. "Accept-Language": self.state.device.language.replace("_", "-"),
  223. "platform": "android",
  224. "ig_mqtt_route": "django",
  225. "pubsub_msg_type_blacklist": "direct, typing_type",
  226. "auth_cache_enabled": "1",
  227. },
  228. )
  229. return zlib.compress(cfg.to_thrift(), level=9)
  230. # endregion
  231. def _on_socket_open(self, client: MQTToTClient, _: Any, sock: socket) -> None:
  232. self._loop.add_reader(sock, client.loop_read)
  233. def _on_socket_close(self, client: MQTToTClient, _: Any, sock: socket) -> None:
  234. self._loop.remove_reader(sock)
  235. def _on_socket_register_write(self, client: MQTToTClient, _: Any, sock: socket) -> None:
  236. self._loop.add_writer(sock, client.loop_write)
  237. def _on_socket_unregister_write(self, client: MQTToTClient, _: Any, sock: socket) -> None:
  238. self._loop.remove_writer(sock)
  239. def _on_connect_handler(
  240. self, client: MQTToTClient, _: Any, flags: dict[str, Any], rc: int
  241. ) -> None:
  242. if rc != 0:
  243. err = pmc.connack_string(rc)
  244. self.log.error("MQTT Connection Error: %s (%d)", err, rc)
  245. if rc == pmc.CONNACK_REFUSED_NOT_AUTHORIZED:
  246. self._disconnect_error = MQTTConnectionUnauthorized()
  247. self.disconnect()
  248. return
  249. self._loop.create_task(self._post_connect())
  250. def _on_disconnect_handler(self, client: MQTToTClient, _: Any, rc: int) -> None:
  251. err_str = "Generic error." if rc == pmc.MQTT_ERR_NOMEM else pmc.error_string(rc)
  252. self.log.debug(f"MQTT disconnection code %d: %s", rc, err_str)
  253. self._clear_response_waiters()
  254. async def _post_connect(self) -> None:
  255. await self._dispatch(Connect())
  256. self.log.debug("Re-subscribing to things after connect")
  257. if self._graphql_subs:
  258. res = await self.graphql_subscribe(self._graphql_subs)
  259. self.log.trace("GraphQL subscribe response: %s", res)
  260. if self._skywalker_subs:
  261. res = await self.skywalker_subscribe(self._skywalker_subs)
  262. self.log.trace("Skywalker subscribe response: %s", res)
  263. if self._iris_seq_id:
  264. retry = 0
  265. while True:
  266. try:
  267. await self.iris_subscribe(self._iris_seq_id, self._iris_snapshot_at_ms)
  268. break
  269. except (asyncio.TimeoutError, IrisSubscribeError) as e:
  270. self.log.exception("Error requesting iris subscribe")
  271. retry += 1
  272. if retry >= 5 or isinstance(e, IrisSubscribeError):
  273. self._disconnect_error = e
  274. self.disconnect()
  275. break
  276. await asyncio.sleep(5)
  277. self.log.debug("Retrying iris subscribe")
  278. def _on_publish_handler(self, client: MQTToTClient, _: Any, mid: int) -> None:
  279. try:
  280. waiter = self._publish_waiters[mid]
  281. except KeyError:
  282. self.log.trace(f"Got publish confirmation for {mid}, but no waiters")
  283. return
  284. self.log.trace(f"Got publish confirmation for {mid}")
  285. waiter.set_result(None)
  286. # region Incoming event parsing
  287. def _parse_direct_thread_path(self, path: str) -> dict:
  288. try:
  289. blank, direct_v2, threads, thread_id, *rest = path.split("/")
  290. except (ValueError, IndexError) as e:
  291. self.log.debug(f"Got {e!r} while parsing path {path}")
  292. raise
  293. if (blank, direct_v2, threads) != ("", "direct_v2", "threads"):
  294. self.log.debug(f"Got unexpected first parts in direct thread path {path}")
  295. raise ValueError("unexpected first three parts in _parse_direct_thread_path")
  296. additional = {"thread_id": thread_id}
  297. if rest:
  298. subitem_key = rest[0]
  299. if subitem_key == "approval_required_for_new_members":
  300. additional["approval_required_for_new_members"] = True
  301. elif subitem_key == "participants" and len(rest) > 2 and rest[2] == "has_seen":
  302. additional["has_seen"] = int(rest[1])
  303. elif subitem_key == "items":
  304. additional["item_id"] = rest[1]
  305. if len(rest) > 4 and rest[2] == "reactions":
  306. additional["reaction_type"] = ReactionType(rest[3])
  307. additional["reaction_user_id"] = int(rest[4])
  308. elif subitem_key in "admin_user_ids":
  309. additional["admin_user_id"] = int(rest[1])
  310. elif subitem_key == "activity_indicator_id":
  311. additional["activity_indicator_id"] = rest[1]
  312. self.log.trace("Parsed path %s -> %s", path, additional)
  313. return additional
  314. def _on_messager_sync_item(self, part: IrisPayloadData, parsed_item: IrisPayload) -> bool:
  315. if part.path.startswith("/direct_v2/threads/"):
  316. raw_message = {
  317. "path": part.path,
  318. "op": part.op,
  319. **self._parse_direct_thread_path(part.path),
  320. }
  321. try:
  322. json_value = json.loads(part.value)
  323. if "reaction_type" in raw_message:
  324. self.log.trace("Treating %s as new reaction data", json_value)
  325. raw_message["new_reaction"] = json_value
  326. json_value["sender_id"] = raw_message.pop("reaction_user_id")
  327. json_value["type"] = raw_message.pop("reaction_type")
  328. json_value["client_context"] = parsed_item.mutation_token
  329. if part.op == Operation.REMOVE:
  330. json_value["emoji"] = None
  331. json_value["timestamp"] = None
  332. else:
  333. raw_message = {
  334. **raw_message,
  335. **json_value,
  336. }
  337. except (json.JSONDecodeError, TypeError):
  338. raw_message["value"] = part.value
  339. message = MessageSyncMessage.deserialize(raw_message)
  340. evt = MessageSyncEvent(iris=parsed_item, message=message)
  341. elif part.path.startswith("/direct_v2/inbox/threads/"):
  342. if part.op == Operation.REMOVE:
  343. blank, direct_v2, inbox, threads, thread_id, *_ = part.path.split("/")
  344. evt = ThreadRemoveEvent.deserialize(
  345. {
  346. "thread_id": thread_id,
  347. "path": part.path,
  348. "op": part.op,
  349. **json.loads(part.value),
  350. }
  351. )
  352. else:
  353. evt = ThreadSyncEvent.deserialize(
  354. {
  355. "path": part.path,
  356. "op": part.op,
  357. **json.loads(part.value),
  358. }
  359. )
  360. else:
  361. self.log.warning(f"Unsupported path {part.path}")
  362. return False
  363. self._outgoing_events.put_nowait(evt)
  364. return True
  365. def _on_message_sync(self, payload: bytes) -> None:
  366. parsed = json.loads(payload.decode("utf-8"))
  367. self.log.trace("Got message sync event: %s", parsed)
  368. has_items = False
  369. for sync_item in parsed:
  370. parsed_item = IrisPayload.deserialize(sync_item)
  371. if self._iris_seq_id < parsed_item.seq_id:
  372. self.log.trace(f"Got new seq_id: {parsed_item.seq_id}")
  373. self._iris_seq_id = parsed_item.seq_id
  374. self._iris_snapshot_at_ms = int(time.time() * 1000)
  375. asyncio.create_task(
  376. self._dispatch(NewSequenceID(self._iris_seq_id, self._iris_snapshot_at_ms))
  377. )
  378. for part in parsed_item.data:
  379. has_items = self._on_messager_sync_item(part, parsed_item) or has_items
  380. if has_items and not self._event_dispatcher_task:
  381. self._event_dispatcher_task = asyncio.create_task(self._dispatcher_loop())
  382. def _on_pubsub(self, payload: bytes) -> None:
  383. parsed_thrift = IncomingMessage.from_thrift(payload)
  384. self.log.trace(f"Got pubsub event {parsed_thrift.topic} / {parsed_thrift.payload}")
  385. message = PubsubPayload.parse_json(parsed_thrift.payload)
  386. for data in message.data:
  387. match = ACTIVITY_INDICATOR_REGEX.match(data.path)
  388. if match:
  389. evt = PubsubEvent(
  390. data=data,
  391. base=message,
  392. thread_id=match.group(1),
  393. activity_indicator_id=match.group(2),
  394. )
  395. self._loop.create_task(self._dispatch(evt))
  396. elif not data.double_publish:
  397. self.log.debug("Pubsub: no activity indicator on data: %s", data)
  398. else:
  399. self.log.debug("Pubsub: double publish: %s", data.path)
  400. def _parse_realtime_sub_item(self, topic: str | GraphQLQueryID, raw: dict) -> Iterable[Any]:
  401. if topic == GraphQLQueryID.APP_PRESENCE:
  402. yield AppPresenceEventPayload.deserialize(raw).presence_event
  403. elif topic == GraphQLQueryID.ZERO_PROVISION:
  404. yield RealtimeZeroProvisionPayload.deserialize(raw).zero_product_provisioning_event
  405. elif topic == GraphQLQueryID.CLIENT_CONFIG_UPDATE:
  406. yield ClientConfigUpdatePayload.deserialize(raw).client_config_update_event
  407. elif topic == GraphQLQueryID.LIVE_REALTIME_COMMENTS:
  408. yield LiveVideoCommentPayload.deserialize(raw).live_video_comment_event
  409. elif topic == "direct":
  410. event = raw["event"]
  411. for item in raw["data"]:
  412. yield RealtimeDirectEvent.deserialize(
  413. {
  414. "event": event,
  415. **self._parse_direct_thread_path(item["path"]),
  416. **item,
  417. }
  418. )
  419. def _on_realtime_sub(self, payload: bytes) -> None:
  420. parsed_thrift = IncomingMessage.from_thrift(payload)
  421. try:
  422. topic = GraphQLQueryID(parsed_thrift.topic)
  423. except ValueError:
  424. topic = parsed_thrift.topic
  425. self.log.trace(f"Got realtime sub event {topic} / {parsed_thrift.payload}")
  426. allowed = (
  427. "direct",
  428. GraphQLQueryID.APP_PRESENCE,
  429. GraphQLQueryID.ZERO_PROVISION,
  430. GraphQLQueryID.CLIENT_CONFIG_UPDATE,
  431. GraphQLQueryID.LIVE_REALTIME_COMMENTS,
  432. )
  433. if topic not in allowed:
  434. return
  435. parsed_json = json.loads(parsed_thrift.payload)
  436. for evt in self._parse_realtime_sub_item(topic, parsed_json):
  437. self._loop.create_task(self._dispatch(evt))
  438. def _handle_send_response(self, message: pmc.MQTTMessage) -> None:
  439. data = json.loads(message.payload.decode("utf-8"))
  440. try:
  441. ccid = data["payload"]["client_context"]
  442. except KeyError:
  443. self.log.warning(
  444. "Didn't find client_context in send message response: %s", message.payload
  445. )
  446. ccid = self._message_response_waiter_id
  447. else:
  448. if ccid != self._message_response_waiter_id:
  449. self.log.error(
  450. "Mismatching client_context in send message response (%s != %s)",
  451. ccid,
  452. self._message_response_waiter_id,
  453. )
  454. return
  455. if self._message_response_waiter and not self._message_response_waiter.done():
  456. self.log.debug("Got response to %s: %s", ccid, message.payload)
  457. self._message_response_waiter.set_result(message)
  458. self._message_response_waiter = None
  459. self._message_response_waiter_id = None
  460. else:
  461. self.log.warning("Didn't find task waiting for response %s", message.payload)
  462. def _on_message_handler(self, client: MQTToTClient, _: Any, message: pmc.MQTTMessage) -> None:
  463. try:
  464. topic = RealtimeTopic.decode(message.topic)
  465. # Instagram Android MQTT messages are always compressed
  466. message.payload = zlib.decompress(message.payload)
  467. if topic == RealtimeTopic.MESSAGE_SYNC:
  468. self._on_message_sync(message.payload)
  469. elif topic == RealtimeTopic.PUBSUB:
  470. self._on_pubsub(message.payload)
  471. elif topic == RealtimeTopic.REALTIME_SUB:
  472. self._on_realtime_sub(message.payload)
  473. elif topic == RealtimeTopic.SEND_MESSAGE_RESPONSE:
  474. self._handle_send_response(message)
  475. else:
  476. try:
  477. waiter = self._response_waiters.pop(topic)
  478. except KeyError:
  479. self.log.debug(
  480. "No handler for MQTT message in %s: %s", topic.value, message.payload
  481. )
  482. else:
  483. self.log.trace("Got response %s: %s", topic.value, message.payload)
  484. waiter.set_result(message)
  485. except Exception:
  486. self.log.exception("Error in incoming MQTT message handler")
  487. self.log.trace("Errored MQTT payload: %s", message.payload)
  488. # endregion
  489. async def _reconnect(self) -> None:
  490. try:
  491. self.log.trace("Trying to reconnect to MQTT")
  492. self._client.reconnect()
  493. except (SocketError, OSError, pmc.WebsocketConnectionError) as e:
  494. # TODO custom class
  495. raise MQTTNotLoggedIn("MQTT reconnection failed") from e
  496. def add_event_handler(
  497. self, evt_type: Type[T], handler: Callable[[T], Awaitable[None]]
  498. ) -> None:
  499. self._event_handlers[evt_type].append(handler)
  500. async def _dispatch(self, evt: T) -> None:
  501. for handler in self._event_handlers[type(evt)]:
  502. try:
  503. await handler(evt)
  504. except Exception:
  505. self.log.exception(f"Error in {type(evt).__name__} handler")
  506. def disconnect(self) -> None:
  507. self._client.disconnect()
  508. async def _dispatcher_loop(self) -> None:
  509. loop_id = f"{hex(id(self))}#{time.monotonic()}"
  510. self.log.debug(f"Dispatcher loop {loop_id} starting")
  511. try:
  512. while True:
  513. evt = await self._outgoing_events.get()
  514. await asyncio.shield(self._dispatch(evt))
  515. except asyncio.CancelledError:
  516. tasks = self._outgoing_events
  517. self._outgoing_events = asyncio.Queue()
  518. if not tasks.empty():
  519. self.log.debug(
  520. f"Dispatcher loop {loop_id} stopping after dispatching {tasks.qsize()} events"
  521. )
  522. while not tasks.empty():
  523. await self._dispatch(tasks.get_nowait())
  524. raise
  525. finally:
  526. self.log.debug(f"Dispatcher loop {loop_id} stopped")
  527. async def listen(
  528. self,
  529. graphql_subs: set[str] | None = None,
  530. skywalker_subs: set[str] | None = None,
  531. seq_id: int = None,
  532. snapshot_at_ms: int = None,
  533. retry_limit: int = 5,
  534. ) -> None:
  535. self._graphql_subs = graphql_subs or set()
  536. self._skywalker_subs = skywalker_subs or set()
  537. self._iris_seq_id = seq_id
  538. self._iris_snapshot_at_ms = snapshot_at_ms
  539. self.log.debug("Connecting to Instagram MQTT")
  540. await self._reconnect()
  541. connection_retries = 0
  542. while True:
  543. try:
  544. await asyncio.sleep(1)
  545. except asyncio.CancelledError:
  546. self.disconnect()
  547. # this might not be necessary
  548. self._client.loop_misc()
  549. break
  550. rc = self._client.loop_misc()
  551. # If disconnect() has been called
  552. # Beware, internal API, may have to change this to something more stable!
  553. if self._client._state == pmc.mqtt_cs_disconnecting:
  554. break # Stop listening
  555. if rc != pmc.MQTT_ERR_SUCCESS:
  556. # If known/expected error
  557. if rc == pmc.MQTT_ERR_CONN_LOST:
  558. await self._dispatch(Disconnect(reason="Connection lost, retrying"))
  559. elif rc == pmc.MQTT_ERR_NOMEM:
  560. # This error is wrongly classified
  561. # See https://github.com/eclipse/paho.mqtt.python/issues/340
  562. await self._dispatch(Disconnect(reason="Connection lost, retrying"))
  563. elif rc == pmc.MQTT_ERR_CONN_REFUSED:
  564. raise MQTTNotLoggedIn("MQTT connection refused")
  565. elif rc == pmc.MQTT_ERR_NO_CONN:
  566. if connection_retries > retry_limit:
  567. raise MQTTNotConnected(f"Connection failed {connection_retries} times")
  568. if self.proxy_handler.update_proxy_url():
  569. self.setup_proxy()
  570. await self._dispatch(ProxyUpdate())
  571. sleep = connection_retries * 2
  572. await self._dispatch(
  573. Disconnect(
  574. reason="MQTT Error: no connection, retrying "
  575. f"in {connection_retries} seconds"
  576. )
  577. )
  578. await asyncio.sleep(sleep)
  579. else:
  580. err = pmc.error_string(rc)
  581. self.log.error("MQTT Error: %s", err)
  582. await self._dispatch(Disconnect(reason=f"MQTT Error: {err}, retrying"))
  583. await self._reconnect()
  584. connection_retries += 1
  585. else:
  586. connection_retries = 0
  587. if self._event_dispatcher_task:
  588. self._event_dispatcher_task.cancel()
  589. self._event_dispatcher_task = None
  590. if self._disconnect_error:
  591. self.log.info("disconnect_error is set, raising and clearing variable")
  592. err = self._disconnect_error
  593. self._disconnect_error = None
  594. raise err
  595. # region Basic outgoing MQTT
  596. def publish(self, topic: RealtimeTopic, payload: str | bytes | dict) -> asyncio.Future:
  597. if isinstance(payload, dict):
  598. payload = json.dumps(payload)
  599. if isinstance(payload, str):
  600. payload = payload.encode("utf-8")
  601. self.log.trace(f"Publishing message in {topic.value} ({topic.encoded}): {payload}")
  602. payload = zlib.compress(payload, level=9)
  603. info = self._client.publish(topic.encoded, payload, qos=1)
  604. self.log.trace(f"Published message ID: {info.mid}")
  605. fut = asyncio.Future()
  606. self._publish_waiters[info.mid] = fut
  607. return fut
  608. async def request(
  609. self,
  610. topic: RealtimeTopic,
  611. response: RealtimeTopic,
  612. payload: str | bytes | dict,
  613. timeout: int | None = None,
  614. ) -> pmc.MQTTMessage:
  615. async with self._response_waiter_locks[response]:
  616. fut = asyncio.Future()
  617. self._response_waiters[response] = fut
  618. await self.publish(topic, payload)
  619. self.log.trace(
  620. f"Request published to {topic.value}, waiting for response {response.name}"
  621. )
  622. return await asyncio.wait_for(fut, timeout)
  623. async def iris_subscribe(self, seq_id: int, snapshot_at_ms: int) -> None:
  624. self.log.debug(f"Requesting iris subscribe {seq_id}/{snapshot_at_ms}")
  625. resp = await self.request(
  626. RealtimeTopic.SUB_IRIS,
  627. RealtimeTopic.SUB_IRIS_RESPONSE,
  628. {
  629. "seq_id": seq_id,
  630. "snapshot_at_ms": snapshot_at_ms,
  631. "snapshot_app_version": self.state.application.APP_VERSION,
  632. "timezone_offset": int(self.state.device.timezone_offset),
  633. "subscription_type": "message",
  634. },
  635. timeout=20,
  636. )
  637. self.log.debug("Iris subscribe response: %s", resp.payload.decode("utf-8"))
  638. resp_dict = json.loads(resp.payload.decode("utf-8"))
  639. if resp_dict["error_type"] and resp_dict["error_message"]:
  640. raise IrisSubscribeError(resp_dict["error_type"], resp_dict["error_message"])
  641. latest_seq_id = resp_dict.get("latest_seq_id")
  642. if latest_seq_id > self._iris_seq_id:
  643. self.log.info(f"Latest sequence ID is {latest_seq_id}, catching up from {seq_id}")
  644. self._iris_seq_id = latest_seq_id
  645. self._iris_snapshot_at_ms = resp_dict.get("subscribed_at_ms", int(time.time() * 1000))
  646. asyncio.create_task(
  647. self._dispatch(NewSequenceID(self._iris_seq_id, self._iris_snapshot_at_ms))
  648. )
  649. def graphql_subscribe(self, subs: set[str]) -> asyncio.Future:
  650. self._graphql_subs |= subs
  651. return self.publish(RealtimeTopic.REALTIME_SUB, {"sub": list(subs)})
  652. def graphql_unsubscribe(self, subs: set[str]) -> asyncio.Future:
  653. self._graphql_subs -= subs
  654. return self.publish(RealtimeTopic.REALTIME_SUB, {"unsub": list(subs)})
  655. def skywalker_subscribe(self, subs: set[str]) -> asyncio.Future:
  656. self._skywalker_subs |= subs
  657. return self.publish(RealtimeTopic.PUBSUB, {"sub": list(subs)})
  658. def skywalker_unsubscribe(self, subs: set[str]) -> asyncio.Future:
  659. self._skywalker_subs -= subs
  660. return self.publish(RealtimeTopic.PUBSUB, {"unsub": list(subs)})
  661. # endregion
  662. # region Actually sending messages and stuff
  663. async def send_foreground_state(self, state: ForegroundStateConfig) -> None:
  664. self.log.debug("Updating foreground state: %s", state)
  665. await self.publish(
  666. RealtimeTopic.FOREGROUND_STATE, zlib.compress(state.to_thrift(), level=9)
  667. )
  668. if state.keep_alive_timeout:
  669. self._client._keepalive = state.keep_alive_timeout
  670. async def send_command(
  671. self,
  672. thread_id: str,
  673. action: ThreadAction,
  674. client_context: str | None = None,
  675. **kwargs: Any,
  676. ) -> CommandResponse | None:
  677. self.log.debug(f"Preparing to send {action} to {thread_id} with {client_context}")
  678. client_context = client_context or self.state.gen_client_context()
  679. req = {
  680. "thread_id": thread_id,
  681. "client_context": client_context,
  682. "offline_threading_id": client_context,
  683. "action": action.value,
  684. # "device_id": self.state.cookies["ig_did"],
  685. **kwargs,
  686. }
  687. lock_start = time.monotonic()
  688. async with self._message_response_waiter_lock:
  689. lock_wait_dur = time.monotonic() - lock_start
  690. if lock_wait_dur > 1:
  691. self.log.warning(f"Waited {lock_wait_dur:.3f} seconds to send {client_context}")
  692. fut = self._message_response_waiter = asyncio.Future()
  693. self._message_response_waiter_id = client_context
  694. self.log.debug(f"Publishing {action} to {thread_id} with {client_context}")
  695. await self.publish(RealtimeTopic.SEND_MESSAGE, req)
  696. self.log.trace(
  697. f"Request published to {RealtimeTopic.SEND_MESSAGE}, "
  698. f"waiting for response {RealtimeTopic.SEND_MESSAGE_RESPONSE}"
  699. )
  700. try:
  701. resp = await asyncio.wait_for(fut, timeout=30000)
  702. except asyncio.TimeoutError:
  703. self.log.error(f"Request with ID {client_context} timed out!")
  704. raise
  705. return CommandResponse.parse_json(resp.payload.decode("utf-8"))
  706. def send_item(
  707. self,
  708. thread_id: str,
  709. item_type: ThreadItemType,
  710. shh_mode: bool = False,
  711. client_context: str | None = None,
  712. **kwargs: Any,
  713. ) -> Awaitable[CommandResponse]:
  714. return self.send_command(
  715. thread_id,
  716. item_type=item_type.value,
  717. is_shh_mode=str(int(shh_mode)),
  718. action=ThreadAction.SEND_ITEM,
  719. client_context=client_context,
  720. **kwargs,
  721. )
  722. def send_hashtag(
  723. self,
  724. thread_id: str,
  725. hashtag: str,
  726. text: str = "",
  727. shh_mode: bool = False,
  728. client_context: str | None = None,
  729. ) -> Awaitable[CommandResponse]:
  730. return self.send_item(
  731. thread_id,
  732. text=text,
  733. item_id=hashtag,
  734. shh_mode=shh_mode,
  735. item_type=ThreadItemType.HASHTAG,
  736. client_context=client_context,
  737. )
  738. def send_like(
  739. self, thread_id: str, shh_mode: bool = False, client_context: str | None = None
  740. ) -> Awaitable[CommandResponse]:
  741. return self.send_item(
  742. thread_id,
  743. shh_mode=shh_mode,
  744. item_type=ThreadItemType.LIKE,
  745. client_context=client_context,
  746. )
  747. def send_location(
  748. self,
  749. thread_id: str,
  750. venue_id: str,
  751. text: str = "",
  752. shh_mode: bool = False,
  753. client_context: str | None = None,
  754. ) -> Awaitable[CommandResponse]:
  755. return self.send_item(
  756. thread_id,
  757. text=text,
  758. item_id=venue_id,
  759. shh_mode=shh_mode,
  760. item_type=ThreadItemType.LOCATION,
  761. client_context=client_context,
  762. )
  763. def send_media(
  764. self,
  765. thread_id: str,
  766. media_id: str,
  767. text: str = "",
  768. shh_mode: bool = False,
  769. client_context: str | None = None,
  770. ) -> Awaitable[CommandResponse]:
  771. return self.send_item(
  772. thread_id,
  773. text=text,
  774. media_id=media_id,
  775. shh_mode=shh_mode,
  776. item_type=ThreadItemType.MEDIA_SHARE,
  777. client_context=client_context,
  778. )
  779. def send_profile(
  780. self,
  781. thread_id: str,
  782. user_id: str,
  783. text: str = "",
  784. shh_mode: bool = False,
  785. client_context: str | None = None,
  786. ) -> Awaitable[CommandResponse]:
  787. return self.send_item(
  788. thread_id,
  789. text=text,
  790. item_id=user_id,
  791. shh_mode=shh_mode,
  792. item_type=ThreadItemType.PROFILE,
  793. client_context=client_context,
  794. )
  795. def send_reaction(
  796. self,
  797. thread_id: str,
  798. emoji: str,
  799. item_id: str,
  800. reaction_status: ReactionStatus = ReactionStatus.CREATED,
  801. target_item_type: ThreadItemType = ThreadItemType.TEXT,
  802. shh_mode: bool = False,
  803. client_context: str | None = None,
  804. ) -> Awaitable[CommandResponse]:
  805. return self.send_item(
  806. thread_id,
  807. reaction_status=reaction_status.value,
  808. node_type="item",
  809. reaction_type="like",
  810. target_item_type=target_item_type.value,
  811. emoji=emoji,
  812. item_id=item_id,
  813. reaction_action_source="double_tap",
  814. shh_mode=shh_mode,
  815. item_type=ThreadItemType.REACTION,
  816. client_context=client_context,
  817. )
  818. def send_user_story(
  819. self,
  820. thread_id: str,
  821. media_id: str,
  822. text: str = "",
  823. shh_mode: bool = False,
  824. client_context: str | None = None,
  825. ) -> Awaitable[CommandResponse]:
  826. return self.send_item(
  827. thread_id,
  828. text=text,
  829. item_id=media_id,
  830. shh_mode=shh_mode,
  831. item_type=ThreadItemType.REEL_SHARE,
  832. client_context=client_context,
  833. )
  834. def send_text(
  835. self,
  836. thread_id: str,
  837. text: str = "",
  838. urls: list[str] | None = None,
  839. shh_mode: bool = False,
  840. client_context: str | None = None,
  841. replied_to_item_id: str | None = None,
  842. replied_to_client_context: str | None = None,
  843. ) -> Awaitable[CommandResponse]:
  844. args = {
  845. "text": text,
  846. }
  847. item_type = ThreadItemType.TEXT
  848. if urls is not None:
  849. args = {
  850. "link_text": text,
  851. "link_urls": json.dumps(urls or []),
  852. }
  853. item_type = ThreadItemType.LINK
  854. return self.send_item(
  855. thread_id,
  856. **args,
  857. shh_mode=shh_mode,
  858. item_type=item_type,
  859. client_context=client_context,
  860. replied_to_item_id=replied_to_item_id,
  861. replied_to_client_context=replied_to_client_context,
  862. )
  863. def mark_seen(
  864. self, thread_id: str, item_id: str, client_context: str | None = None
  865. ) -> Awaitable[None]:
  866. return self.send_command(
  867. thread_id,
  868. item_id=item_id,
  869. action=ThreadAction.MARK_SEEN,
  870. client_context=client_context,
  871. )
  872. def mark_visual_item_seen(
  873. self, thread_id: str, item_id: str, client_context: str | None = None
  874. ) -> Awaitable[CommandResponse]:
  875. return self.send_command(
  876. thread_id,
  877. item_id=item_id,
  878. action=ThreadAction.MARK_VISUAL_ITEM_SEEN,
  879. client_context=client_context,
  880. )
  881. def indicate_activity(
  882. self,
  883. thread_id: str,
  884. activity_status: TypingStatus = TypingStatus.TEXT,
  885. client_context: str | None = None,
  886. ) -> Awaitable[CommandResponse]:
  887. return self.send_command(
  888. thread_id,
  889. activity_status=activity_status.value,
  890. action=ThreadAction.INDICATE_ACTIVITY,
  891. client_context=client_context,
  892. )
  893. # endregion