From 9fa6343faefad5b6392bf3a2d3e7aea20959b830 Mon Sep 17 00:00:00 2001
From: Vlad Pronsky <v.pronsky@gmail.com>
Date: Thu, 18 Apr 2024 18:39:39 +0300
Subject: [PATCH] fix otp code

---
 twscrape/accounts_pool.py | 2 +-
 twscrape/login.py         | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/twscrape/accounts_pool.py b/twscrape/accounts_pool.py
index 95f2a8d..905e8aa 100644
--- a/twscrape/accounts_pool.py
+++ b/twscrape/accounts_pool.py
@@ -80,7 +80,7 @@ class AccountsPool:
         user_agent: str | None = None,
         proxy: str | None = None,
         cookies: str | None = None,
-        mfa_code: str | None = None
+        mfa_code: str | None = None,
     ):
         qs = "SELECT * FROM accounts WHERE username = :username"
         rs = await fetchone(self._db_file, qs, {"username": username})
diff --git a/twscrape/login.py b/twscrape/login.py
index 8285699..0b37b00 100644
--- a/twscrape/login.py
+++ b/twscrape/login.py
@@ -2,8 +2,8 @@ import imaplib
 from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any
-import pyotp
 
+import pyotp
 from httpx import AsyncClient, Response
 
 from .account import Account
@@ -121,6 +121,9 @@ async def login_enter_password(ctx: TaskCtx) -> Response:
 
 
 async def login_two_factor_auth_challenge(ctx: TaskCtx) -> Response:
+    if ctx.acc.mfa_code is None:
+        raise ValueError("MFA code is required")
+
     totp = pyotp.TOTP(ctx.acc.mfa_code)
     payload = {
         "flow_token": ctx.prev["flow_token"],
@@ -133,7 +136,7 @@ async def login_two_factor_auth_challenge(ctx: TaskCtx) -> Response:
     }
 
     rep = await ctx.client.post(LOGIN_URL, json=payload)
-    raise_for_status(rep, "login_two_factor_auth_challenge")
+    rep.raise_for_status()
     return rep