Merge pull request #153 from ritikkumarsahu/TOTP

feat: support TOTP

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
   "fake-useragent>=1.4.0",
   "httpx>=0.26.0",
   "loguru>=0.7.0",
+  "pyotp>=2.9.0",
 ]
 
 [project.optional-dependencies]

twscrape/account.py

@@ -24,6 +24,7 @@ class Account(JSONTrait):
     stats: dict[str, int] = field(default_factory=dict)  # queue: requests
     headers: dict[str, str] = field(default_factory=dict)
     cookies: dict[str, str] = field(default_factory=dict)
+    mfa_code: str | None = None
     proxy: str | None = None
     error_msg: str | None = None
     last_used: datetime | None = None

twscrape/accounts_pool.py

@@ -80,6 +80,7 @@ class AccountsPool:
         user_agent: str | None = None,
         proxy: str | None = None,
         cookies: str | None = None,
+        mfa_code: str | None = None
     ):
         qs = "SELECT * FROM accounts WHERE username = :username"
         rs = await fetchone(self._db_file, qs, {"username": username})
@@ -99,6 +100,7 @@ class AccountsPool:
             headers={},
             cookies=parse_cookies(cookies) if cookies else {},
             proxy=proxy,
+            mfa_code=mfa_code,
         )
 
         if "ct0" in account.cookies:

twscrape/db.py

@@ -81,10 +81,14 @@ async def migrate(db: aiosqlite.Connection):
     async def v3():
         await db.execute("ALTER TABLE accounts ADD COLUMN _tx TEXT DEFAULT NULL")
 
+    async def v4():
+        await db.execute("ALTER TABLE accounts ADD COLUMN mfa_code TEXT DEFAULT NULL")
+
     migrations = {
         1: v1,
         2: v2,
         3: v3,
+        4: v4,
     }
 
     # logger.debug(f"Current migration v{uv} (latest v{len(migrations)})")

twscrape/login.py

@@ -2,6 +2,7 @@ import imaplib
 from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any
+import pyotp
 
 from httpx import AsyncClient, Response
 
@@ -119,6 +120,23 @@ async def login_enter_password(ctx: TaskCtx) -> Response:
     return rep
 
 
+async def login_two_factor_auth_challenge(ctx: TaskCtx) -> Response:
+    totp = pyotp.TOTP(ctx.acc.mfa_code)
+    payload = {
+        "flow_token": ctx.prev["flow_token"],
+        "subtask_inputs": [
+            {
+                "subtask_id": "LoginTwoFactorAuthChallenge",
+                "enter_text": {"text": totp.now(), "link": "next_link"},
+            }
+        ],
+    }
+
+    rep = await ctx.client.post(LOGIN_URL, json=payload)
+    raise_for_status(rep, "login_two_factor_auth_challenge")
+    return rep
+
+
 async def login_duplication_check(ctx: TaskCtx) -> Response:
     payload = {
         "flow_token": ctx.prev["flow_token"],
@@ -212,6 +230,8 @@ async def next_login_task(ctx: TaskCtx, rep: Response):
                 return await login_duplication_check(ctx)
             if task_id == "LoginEnterPassword":
                 return await login_enter_password(ctx)
+            if task_id == "LoginTwoFactorAuthChallenge":
+                return await login_two_factor_auth_challenge(ctx)
             if task_id == "LoginEnterUserIdentifierSSO":
                 return await login_enter_username(ctx)
             if task_id == "LoginJsInstrumentationSubtask":