diff --git a/pyproject.toml b/pyproject.toml index e9f7967..b4ebe04 100644 --- a/pyproject.toml +++ b/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] diff --git a/twscrape/account.py b/twscrape/account.py index 172900a..4d9d7cb 100644 --- a/twscrape/account.py +++ b/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 diff --git a/twscrape/accounts_pool.py b/twscrape/accounts_pool.py index c41f5be..95f2a8d 100644 --- a/twscrape/accounts_pool.py +++ b/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: diff --git a/twscrape/db.py b/twscrape/db.py index a26adae..862ee10 100644 --- a/twscrape/db.py +++ b/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)})") diff --git a/twscrape/login.py b/twscrape/login.py index 5c8c99b..8285699 100644 --- a/twscrape/login.py +++ b/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":