diff --git a/twscrape/accounts_pool.py b/twscrape/accounts_pool.py index 9f483f9..4c12b91 100644 --- a/twscrape/accounts_pool.py +++ b/twscrape/accounts_pool.py @@ -94,6 +94,7 @@ class AccountsPool: account.active = True await self.save(account) + logger.info(f"Account {username} added successfully (active={account.active})") async def delete_accounts(self, usernames: str | list[str]): usernames = usernames if isinstance(usernames, list) else [usernames] @@ -131,9 +132,9 @@ class AccountsPool: """ await execute(self._db_file, qs, data) - async def login(self, account: Account): + async def login(self, account: Account, email_first: bool = False): try: - await login(account) + await login(account, email_first=email_first) logger.info(f"Logged in to {account.username} successfully") return True except Exception as e: @@ -142,7 +143,7 @@ class AccountsPool: finally: await self.save(account) - async def login_all(self): + async def login_all(self, email_first=False): qs = "SELECT * FROM accounts WHERE active = false AND error_msg IS NULL" rs = await fetchall(self._db_file, qs) @@ -152,11 +153,11 @@ class AccountsPool: counter = {"total": len(accounts), "success": 0, "failed": 0} for i, x in enumerate(accounts, start=1): logger.info(f"[{i}/{len(accounts)}] Logging in {x.username} - {x.email}") - status = await self.login(x) + status = await self.login(x, email_first=email_first) counter["success" if status else "failed"] += 1 return counter - async def relogin(self, usernames: str | list[str]): + async def relogin(self, usernames: str | list[str], email_first=False): usernames = usernames if isinstance(usernames, list) else [usernames] usernames = list(set(usernames)) if not usernames: @@ -176,12 +177,12 @@ class AccountsPool: """ await execute(self._db_file, qs) - await self.login_all() + await self.login_all(email_first=email_first) - async def relogin_failed(self): + async def relogin_failed(self, email_first=False): qs = "SELECT username FROM accounts WHERE active = false AND error_msg IS NOT NULL" rs = await fetchall(self._db_file, qs) - await self.relogin([x["username"] for x in rs]) + await self.relogin([x["username"] for x in rs], email_first=email_first) async def reset_locks(self): qs = "UPDATE accounts SET locks = json_object()" diff --git a/twscrape/cli.py b/twscrape/cli.py index 6c09ad7..f2c1e7b 100644 --- a/twscrape/cli.py +++ b/twscrape/cli.py @@ -80,15 +80,16 @@ async def main(args): return if args.command == "login_accounts": - print(await pool.login_all()) + stats = await pool.login_all(email_first=args.email_first) + print(stats) return if args.command == "relogin_failed": - await pool.relogin_failed() + await pool.relogin_failed(email_first=args.email_first) return if args.command == "relogin": - await pool.relogin(args.usernames) + await pool.relogin(args.usernames, email_first=args.email_first) return if args.command == "reset_locks": @@ -164,12 +165,16 @@ def run(): del_accounts = subparsers.add_parser("del_accounts", help="Delete accounts") del_accounts.add_argument("usernames", nargs="+", default=[], help="Usernames to delete") - subparsers.add_parser("login_accounts", help="Login accounts") + login_cmd = subparsers.add_parser("login_accounts", help="Login accounts") + login_cmd.add_argument("--email-first", type=bool, default=False, help="Check email first") relogin = subparsers.add_parser("relogin", help="Re-login selected accounts") relogin.add_argument("usernames", nargs="+", default=[], help="Usernames to re-login") + relogin.add_argument("--email-first", type=bool, default=False, help="Check email first") + + re_failed = subparsers.add_parser("relogin_failed", help="Retry login for failed accounts") + re_failed.add_argument("--email-first", type=bool, default=False, help="Check email first") - subparsers.add_parser("relogin_failed", help="Retry login for failed accounts") subparsers.add_parser("reset_locks", help="Reset all locks") subparsers.add_parser("delete_inactive", help="Delete inactive accounts") diff --git a/twscrape/imap.py b/twscrape/imap.py index 8300bde..c9c2226 100644 --- a/twscrape/imap.py +++ b/twscrape/imap.py @@ -33,14 +33,14 @@ def add_imap_mapping(email_domain: str, imap_domain: str): IMAP_MAPPING[email_domain] = imap_domain -def get_imap_domain(email: str) -> str: +def _get_imap_domain(email: str) -> str: email_domain = email.split("@")[1] if email_domain in IMAP_MAPPING: return IMAP_MAPPING[email_domain] return f"imap.{email_domain}" -def search_email_code(imap: imaplib.IMAP4_SSL, count: int, min_t: datetime | None) -> str | None: +def _wait_email_code(imap: imaplib.IMAP4_SSL, count: int, min_t: datetime | None) -> str | None: for i in range(count, 0, -1): _, rep = imap.fetch(str(i), "(RFC822)") for x in rep: @@ -71,7 +71,7 @@ async def imap_get_email_code( _, rep = imap.select("INBOX") now_count = int(rep[0].decode("utf-8")) if len(rep) > 0 and rep[0] is not None else 0 if now_count > was_count: - code = search_email_code(imap, now_count, min_t) + code = _wait_email_code(imap, now_count, min_t) if code is not None: return code @@ -87,8 +87,8 @@ async def imap_get_email_code( raise e -async def imap_try_login(email: str, password: str): - domain = get_imap_domain(email) +async def imap_login(email: str, password: str): + domain = _get_imap_domain(email) imap = imaplib.IMAP4_SSL(domain) try: diff --git a/twscrape/login.py b/twscrape/login.py index 9b84410..f2fc6f3 100644 --- a/twscrape/login.py +++ b/twscrape/login.py @@ -4,7 +4,7 @@ from httpx import AsyncClient, HTTPStatusError, Response from .account import Account from .constants import LOGIN_URL -from .imap import imap_get_email_code, imap_try_login +from .imap import imap_get_email_code, imap_login from .logger import logger from .utils import raise_for_status @@ -117,6 +117,9 @@ async def login_confirm_email(client: AsyncClient, acc: Account, prev: dict, ima async def login_confirm_email_code(client: AsyncClient, acc: Account, prev: dict, imap): + if not imap: + imap = await imap_login(acc.email, acc.password) + now_time = datetime.now(timezone.utc) - timedelta(seconds=30) value = await imap_get_email_code(imap, acc.email, now_time) @@ -181,14 +184,16 @@ async def next_login_task(client: AsyncClient, acc: Account, rep: Response, imap return None -async def login(acc: Account, fresh=False) -> Account: +async def login(acc: Account, email_first=False) -> Account: log_id = f"{acc.username} - {acc.email}" - if acc.active and not fresh: + if acc.active: logger.info(f"account already active {log_id}") return acc - # check if email is valid first - imap = await imap_try_login(acc.email, acc.email_password) + if email_first: + imap = await imap_login(acc.email, acc.email_password) + else: + imap = None client = acc.make_client() guest_token = await get_guest_token(client)