Retry all provisioning related HTTP requests

mautrix_instagram/web/provisioning_api.py

@@ -46,6 +46,7 @@ from mauigpapi.errors import (
     IGRateLimitError,
     IGResponseError,
 )
+from mauigpapi.proxy import proxy_with_retry
 from mauigpapi.types import ChallengeStateResponse, LoginResponse
 from mautrix.types import JSON, Serializable, UserID
 from mautrix.util.logging import TraceLogger
@@ -276,7 +277,13 @@ class ProvisioningAPI:
         track(user, "$login_start", {"type": "instagram"})
         api, state = await get_login_state(user, self.device_seed)
         try:
-            resp = await api.login(username, password)
+            resp = await proxy_with_retry(
+                "provisioning_api.login",
+                lambda: api.login(username, password),
+                logger=self.log,
+                proxy_handler=user.proxy_handler,
+                on_proxy_change=user.on_proxy_update,
+            )
         except IGLoginTwoFactorRequiredError as e:
             return self._2fa_required(user, username, state, e)
         except IGChallengeError as e:
@@ -382,7 +389,13 @@ class ProvisioningAPI:
             username = state.login_2fa_username
         track(user, "$login_resend_2fa_sms")
         try:
-            resp = await api.send_two_factor_login_sms(username, identifier=identifier)
+            resp = await proxy_with_retry(
+                "provisioning_api.send_two_factor_login_sms",
+                lambda: api.send_two_factor_login_sms(username, identifier=identifier),
+                logger=self.log,
+                proxy_handler=user.proxy_handler,
+                on_proxy_change=user.on_proxy_update,
+            )
         except IGRateLimitError as e:
             track(user, "$login_resend_2fa_sms_fail", {"error": "ratelimit"})
             try:
@@ -441,8 +454,14 @@ class ProvisioningAPI:
             username = state.login_2fa_username
         track(user, "$login_submit_2fa")
         try:
-            resp = await api.two_factor_login(
-                username, code=code, identifier=identifier, is_totp=is_totp
+            resp = await proxy_with_retry(
+                "provisioning_api.two_factor_login",
+                lambda: api.two_factor_login(
+                    username, code=code, identifier=identifier, is_totp=is_totp
+                ),
+                logger=self.log,
+                proxy_handler=user.proxy_handler,
+                on_proxy_change=user.on_proxy_update,
             )
         except IGBad2FACodeError:
             self.log.debug("%s submitted an incorrect 2-factor auth code", user.mxid)
@@ -482,7 +501,13 @@ class ProvisioningAPI:
         self, user: u.User, state: AndroidState, api: AndroidAPI, err: IGChallengeError, after: str
     ) -> web.Response:
         try:
-            resp = await api.challenge_auto(reset=after == "2fa")
+            resp = await proxy_with_retry(
+                "provisioning_api.challenge_auto",
+                lambda: api.challenge_auto(reset=after == "2fa"),
+                logger=self.log,
+                proxy_handler=user.proxy_handler,
+                on_proxy_change=user.on_proxy_update,
+            )
         except Exception:
             self.log.exception("Challenge reset failed for %s", user.mxid)
             track(user, "$login_failed", {"error": "challenge-reset-fail", "after": after})
@@ -548,7 +573,13 @@ class ProvisioningAPI:
         state: AndroidState = user.command_status["state"]
         track(user, "$login_submit_challenge")
         try:
-            resp = await api.challenge_send_security_code(code=code)
+            resp = await proxy_with_retry(
+                "provisioning_api.challenge_send_security_code",
+                lambda: api.challenge_send_security_code(code=code),
+                logger=self.log,
+                proxy_handler=user.proxy_handler,
+                on_proxy_change=user.on_proxy_update,
+            )
         except IGChallengeWrongCodeError:
             self.log.debug("%s submitted an incorrect challenge code", user.mxid)
             track(user, "$login_failed", {"error": "incorrect-challenge-code"})
@@ -602,7 +633,13 @@ class ProvisioningAPI:
             else "<username not known>"
         )
         try:
-            resp = await api.current_user()
+            resp = await proxy_with_retry(
+                "provisioning_api.finish_login",
+                lambda: api.current_user(),
+                logger=self.log,
+                proxy_handler=user.proxy_handler,
+                on_proxy_change=user.on_proxy_update,
+            )
         except IGChallengeError as e:
             if isinstance(login_resp, ChallengeStateResponse):
                 track(user, "$login_failed", {"error": "repeat-challenge", "after": after})
@@ -694,7 +731,13 @@ class ProvisioningAPI:
         track(user, "$login_start", {"type": "facebook"})
         api, state = await get_login_state(user, self.device_seed)
         try:
-            resp = await api.facebook_signup(fb_access_token)
+            resp = await proxy_with_retry(
+                "provisioning_api.post_fb_login_token",
+                lambda: api.facebook_signup(fb_access_token),
+                logger=self.log,
+                proxy_handler=user.proxy_handler,
+                on_proxy_change=user.on_proxy_update,
+            )
         except IGFBNoContactPointFoundError as e:
             self.log.debug(
                 "%s sent a valid Facebook token, "