load_thread_async.py

import collections
from datetime import datetime, timedelta, UTC
import json
from pathlib import Path
from pprint import pprint
import time
import traceback
import httpx
import sys
import asyncio
import inspect
import random
import re
import argparse
import tempfile
import os
import shutil
from trid import gen_trid, init_trid
from aio_pool import AioPool

red_color = "\033[91m"
green_color = "\033[92m"
reset_color = "\033[0m"

SESSION_ID = datetime.now(UTC).isoformat()

USERAGENT = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:143.0) Gecko/20100101 Firefox/143.0'

async def sure(fn, max_sleep=32):
    t = 1
    while True:
        try:
            return await fn()
        except Exception:
            print(f'sure failed', inspect.getsource(fn).strip())
            traceback.print_exc()
            await asyncio.sleep(t)
            t = min(max_sleep, t * 2)


def find_dicts_deep(data, fn_match):
    l = []
    if isinstance(data, dict):
        if fn_match(data):
            l.append(data)
        else:
            for v in data.values():
                l.extend(find_dicts_deep(v, fn_match))
    elif isinstance(data, list):
        for i in data:
            l.extend(find_dicts_deep(i, fn_match))
    return l


def find_all_cursors(data):
    return find_dicts_deep(
        data, lambda d: d.get('__typename') == 'TimelineTimelineCursor')


def find_all_tweets(data):
    return find_dicts_deep(data, lambda d: d.get('__typename') == 'Tweet')


def find_all_users(data):
    return find_dicts_deep(data, lambda d: d.get('__typename') == 'User')


def find_article_referenced_tweets(data):
    q = find_dicts_deep(data, lambda d: 'tweetId' in d)
    return list({i['tweetId'] for i in q if isinstance(i.get('tweetId'), str)})


async def load_tweets_by_ids(session: httpx.AsyncClient, tweet_ids: list[str]):
    headers = {
        'content-type': 'application/json',
        'Accept-Language': 'en-US,en;q=0.5',
        'x-twitter-client-language': 'en',
        'x-twitter-active-user': 'yes',
    }

    params = {
        'variables': f'{{"tweetIds":{json.dumps(tweet_ids)},"includePromotedContent":true,"withBirdwatchNotes":true,"withVoice":true,"withCommunity":true}}',
        'features': '{"creator_subscriptions_tweet_preview_api_enabled":true,"premium_content_api_read_enabled":false,"communities_web_enable_tweet_community_results_fetch":true,"c9s_tweet_anatomy_moderator_badge_enabled":true,"responsive_web_grok_analyze_button_fetch_trends_enabled":false,"responsive_web_grok_analyze_post_followups_enabled":true,"responsive_web_jetfuel_frame":true,"responsive_web_grok_share_attachment_enabled":true,"articles_preview_enabled":true,"responsive_web_edit_tweet_api_enabled":true,"graphql_is_translatable_rweb_tweet_is_translatable_enabled":true,"view_counts_everywhere_api_enabled":true,"longform_notetweets_consumption_enabled":true,"responsive_web_twitter_article_tweet_consumption_enabled":true,"tweet_awards_web_tipping_enabled":false,"responsive_web_grok_show_grok_translated_post":false,"responsive_web_grok_analysis_button_from_backend":true,"creator_subscriptions_quote_tweet_preview_enabled":false,"freedom_of_speech_not_reach_fetch_enabled":true,"standardized_nudges_misinfo":true,"tweet_with_visibility_results_prefer_gql_limited_actions_policy_enabled":true,"longform_notetweets_rich_text_read_enabled":true,"longform_notetweets_inline_media_enabled":true,"payments_enabled":false,"profile_label_improvements_pcf_label_in_post_enabled":true,"rweb_tipjar_consumption_enabled":true,"verified_phone_label_enabled":false,"responsive_web_grok_image_annotation_enabled":true,"responsive_web_grok_imagine_annotation_enabled":true,"responsive_web_grok_community_note_auto_translation_is_enabled":false,"responsive_web_graphql_skip_user_profile_image_extensions_enabled":false,"responsive_web_graphql_timeline_navigation_enabled":true,"responsive_web_enhance_cards_enabled":false}',
    }

    r = await session.get(
        'https://x.com/i/api/graphql/exI9CNWtq0Eoq2aSjneeew/TweetResultsByRestIds',
        params=params,
        headers=headers,
        timeout=20.0,
    )

    if r.status_code != 200:
        r.raise_for_status()

    try:
        j = r.json()
    except Exception:
        print('json decode error', r.status_code, r.content)
        raise

    return j


async def load_tweet(session: httpx.AsyncClient, tweet_id, cursor):
    headers = {
        'content-type': 'application/json',
        'Accept-Language': 'en-US,en;q=0.5',
        'x-twitter-client-language': 'en',
        'x-twitter-active-user': 'yes',
    }

    cursor_str = ''
    if cursor:
        cursor_str = f'"cursor":"{cursor}","referrer":"tweet",'

    params = {
        'variables': f'{{"focalTweetId":"{tweet_id}",{cursor_str}"with_rux_injections":false,"rankingMode":"Relevance","includePromotedContent":true,"withCommunity":true,"withQuickPromoteEligibilityTweetFields":true,"withBirdwatchNotes":true,"withVoice":true}}',
        'features': '{"rweb_video_screen_enabled":false,"payments_enabled":false,"profile_label_improvements_pcf_label_in_post_enabled":true,"responsive_web_profile_redirect_enabled":false,"rweb_tipjar_consumption_enabled":true,"verified_phone_label_enabled":false,"creator_subscriptions_tweet_preview_api_enabled":true,"responsive_web_graphql_timeline_navigation_enabled":true,"responsive_web_graphql_skip_user_profile_image_extensions_enabled":false,"premium_content_api_read_enabled":false,"communities_web_enable_tweet_community_results_fetch":true,"c9s_tweet_anatomy_moderator_badge_enabled":true,"responsive_web_grok_analyze_button_fetch_trends_enabled":false,"responsive_web_grok_analyze_post_followups_enabled":true,"responsive_web_jetfuel_frame":true,"responsive_web_grok_share_attachment_enabled":true,"articles_preview_enabled":true,"responsive_web_edit_tweet_api_enabled":true,"graphql_is_translatable_rweb_tweet_is_translatable_enabled":true,"view_counts_everywhere_api_enabled":true,"longform_notetweets_consumption_enabled":true,"responsive_web_twitter_article_tweet_consumption_enabled":true,"tweet_awards_web_tipping_enabled":false,"responsive_web_grok_show_grok_translated_post":false,"responsive_web_grok_analysis_button_from_backend":true,"creator_subscriptions_quote_tweet_preview_enabled":false,"freedom_of_speech_not_reach_fetch_enabled":true,"standardized_nudges_misinfo":true,"tweet_with_visibility_results_prefer_gql_limited_actions_policy_enabled":true,"longform_notetweets_rich_text_read_enabled":true,"longform_notetweets_inline_media_enabled":true,"responsive_web_grok_image_annotation_enabled":true,"responsive_web_grok_imagine_annotation_enabled":true,"responsive_web_grok_community_note_auto_translation_is_enabled":false,"responsive_web_enhance_cards_enabled":false}',
        'fieldToggles': '{"withArticleRichContentState":true,"withArticlePlainText":false,"withGrokAnalyze":false,"withDisallowedReplyControls":false}',
    }

    r = await session.get(
        'https://x.com/i/api/graphql/bj0Uyae1s0en3ti3POB5fQ/TweetDetail',
        params=params,
        headers=headers,
        timeout=20.0,
    )
    # p = Path(f'./debug_session/{SESSION_ID}/{datetime.now(UTC).isoformat()}.json')
    # p.parent.mkdir(exist_ok=True)
    # p.write_bytes(r.content)
    if r.status_code == 404:
        print('Got 404, request signature failed!')
        raise RateLimitError(datetime.now(UTC) + timedelta(seconds=3))

    if r.status_code == 429:  # ratelimit
        print(r.headers)
        try:
            print(r.json())
        except Exception:
            print('json decode error', r.content)

        if r.headers.get('x-rate-limit-remaining') == '0':
            raise RateLimitError(
                datetime.fromtimestamp(int(r.headers['x-rate-limit-reset']), tz=UTC)
            )
        else:
            # {'errors': [{'code': 88, 'message': 'Rate limit exceeded.'}]}
            raise RateLimitError(datetime.now(UTC) + timedelta(hours=24))

    try:
        j = r.json()
    except Exception:
        print('json decode error', r.status_code, r.content)
        raise

    if j == {'errors': [{'message': 'Invalid or expired token', 'code': 89}]}:
        raise InvalidTokenError

    return j


def get_tweet_ids_belonging_to_thread(tweets, thread_id):
    # should ignore quoted tweets in `tweets` list
    thread = set()
    thread.add(thread_id)

    while True:
        added = 0
        for t in tweets:
            should_add = (
                t['rest_id'] not in thread
                and t['legacy'].get('in_reply_to_status_id_str', -1) in thread
            )
            if not should_add:
                continue
            thread.add(t['rest_id'])
            added += 1

        if added == 0:
            break

    return thread


async def load_tree(pool: AioPool, thread_id, limit_requests=None):
    total_count = 0

    already_fetched = set()  # (tweet_id, cursor)
    already_fetched.add((thread_id, None))

    loaded_tweets = {}  # id -> tweet
    loaded_users = {}  # id -> user

    expected_replies_count = collections.Counter()  # tweet id -> count
    got_replies_count = collections.Counter()

    def put_task_limited(*args):
        if limit_requests is not None and pool.requests_count >= limit_requests:
            return False
        pool.put_task(*args)
        return True

    put_task_limited(thread_id, None)

    while True:
        async for (args, kwargs, response) in pool.results_iter():
            tweet_id, cursor = args
            total_count += 1

            print(
                f'Fetched {tweet_id!r} {"*" if cursor else ""} '
                f'({pool.responses_count} / {total_count})'
            )

            if 'errors' in response:
                print('Errors')
                print(response.get('errors'))
                # continue

            tweets = find_all_tweets(response)
            tweets.extend([i['tweet'] for i in find_dicts_deep(
                response,
                lambda d: d.get('__typename') == 'TweetWithVisibilityResults'
            )])
            users = find_all_users(response)
            cursors = find_all_cursors(response)
            # print('CURSORS', cursors)

            for t in tweets:
                try:
                    tl = t['legacy']
                except KeyError:
                    print('#legacy', t)
                    traceback.print_exc()
                    continue
                tid = t['rest_id']
                expected_replies_count[tid] = max(
                    expected_replies_count[tid], tl['reply_count'])

                if tid in loaded_tweets:
                    continue

                loaded_tweets[tid] = t
                if tl.get('in_reply_to_status_id_str'):
                    got_replies_count[tl['in_reply_to_status_id_str']] += 1

            for u in users:
                if 'rest_id' not in u:
                    print('#rest_id', u)
                    continue
                loaded_users[u['rest_id']] = u

            cursors_terminated = [i['direction'] for i in find_dicts_deep(response, lambda d: d.get('type') == 'TimelineTerminateTimeline')]
            # print('Terminated cursors:', cursors_terminated)
            for c in cursors:
                if c['cursorType'] == 'Top':
                    continue
                if c['cursorType'] in cursors_terminated:
                    continue

                f = (tweet_id, c['value'])
                if f not in already_fetched:
                    already_fetched.add(f)
                    put_task_limited(*f)

        # when no cursors are left to fetch,
        # check which tweets don't have all replies loaded

        thread = get_tweet_ids_belonging_to_thread(
            loaded_tweets.values(),
            thread_id
        )

        added = 0

        for tweet_id, expected in expected_replies_count.items():
            got = got_replies_count[tweet_id]
            if expected > got and tweet_id in thread:
                f = (tweet_id, None)
                if f not in already_fetched:
                    already_fetched.add(f)
                    if put_task_limited(*f):
                        added += 1

        print()
        print('ADDED', added)
        print()

        if added == 0:
            break

    return {
        'thread_tweet_id': thread_id,
        'tweets': list(loaded_tweets.values()),
        'users': list(loaded_users.values()),
    }


class InvalidTokenError(Exception): pass


class RateLimitError(Exception):
    def __init__(self, till, *args: object) -> None:
        super().__init__(*args)
        self.till = till


class SessionManager:
    def __init__(self, accounts) -> None:
        self.sessions = []
        for acc_data in accounts:
            if not acc_data.get('cookies'):
                continue
            csrf_handler = HttpxTwitterCsrf()
            session = httpx.AsyncClient(
                http2=True,
                event_hooks={'request': [csrf_handler]}
            )
            csrf_handler.session = session
            trid = 'ZTpUeXBlRXJyb3I6IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKC4uLikgaXMgbnVsbA=='
            session.headers = {
                'Accept': '*/*',
                'Accept-Language': 'en-US,en;q=0.5',
                'content-type': 'application/json',
                'x-twitter-auth-type': 'OAuth2Session',
                'x-twitter-client-language': 'en',
                'x-twitter-active-user': 'yes',
                # 'x-client-transaction-id': trid,
                'DNT': '1',
                'Sec-Fetch-Dest': 'empty',
                'Sec-Fetch-Mode': 'cors',
                'Sec-Fetch-Site': 'same-origin',
                'authorization': 'Bearer AAAAAAAAAAAAAAAAAAAAANRILgAAAAAAnNwIzUejRCOuH5E6I8xnZz4puTs%3D1Zv7ttfk8LF81IUq16cHjhLTvJu4FA33AGWWjCpTnA',
                'Sec-GPC': '1',
            }
            session.cookies.update(acc_data['cookies'])
            session.headers['User-Agent'] = USERAGENT
            self.sessions.append({
                'invalid': False,
                'session': session,
                'ratelimit_till': datetime.now(UTC) - timedelta(minutes=1)
            })

    def get_session(self):
        for i in random.sample(self.sessions, k=len(self.sessions)):
            if i['invalid']:
                continue

            if i['ratelimit_till'] < datetime.now(UTC):
                return i['session']

    def ratelimit(self, session, till):
        for i in self.sessions:
            if i['session'] == session:
                i['ratelimit_till'] = till
                break

    def invalidate(self, session):
        raise Exception
        for i in self.sessions:
            if i['session'] == session:
                i['invalid'] = True
                break


async def worker_fn(worker_id: int, manager: SessionManager, tweet_id, cursor):
    while True:
        print(f'[{worker_id}] fetching {tweet_id!r} {"*" if cursor else ""}')

        while True:
            try:
                session = manager.get_session()
                if session:
                    r = await load_tweet(session, tweet_id, cursor)
                    await asyncio.sleep(0.2)
                    print(f'[{worker_id}] done')
                    return r
            except InvalidTokenError as e:
                print(f'[{worker_id}] invalid token')
                manager.invalidate(session)
                await asyncio.sleep(0.01)
                continue
            except RateLimitError as e:
                print(f'[{worker_id}] rate limited {e.till}')
                manager.ratelimit(session, e.till)
            except httpx.RequestError as e:
                print(f'could not load tweet (network error). Error: {e}')
            except asyncio.exceptions.CancelledError:
                raise
            except Exception:
                print('could not load tweet. Traceback:')
                traceback.print_exc()
                exit(1)
            await asyncio.sleep(5)


class HttpxTwitterCsrf:
    def __init__(self):
        self.session = None

    async def handle(self, request):
        request.headers['x-csrf-token'] = self.session.cookies.get('ct0')
        request.headers['x-client-transaction-id'] = gen_trid(
            request.method, str(request.url))

    def __call__(self, *args, **kwargs):
        return self.handle(*args, **kwargs)


async def main():
    parser = argparse.ArgumentParser(description='Load Twitter thread data')
    parser.add_argument('thread_id_or_url', help='Thread ID or Twitter URL')
    parser.add_argument('output_path', help='Output JSON file path')
    parser.add_argument('--limit-requests', type=int, help='Maximum number of requests to make (default: no limit)')
    args = parser.parse_args()

    input_arg = args.thread_id_or_url
    output_path = args.output_path
    limit_requests = args.limit_requests

    thread_id = None
    if input_arg.isdigit():
        thread_id = input_arg
    else:
        match = re.search(r'/status/(\d+)', input_arg)
        if match:
            thread_id = match.group(1)

    if not thread_id:
        print(f"Could not extract thread ID from {input_arg}", file=sys.stderr)
        exit(1)

    num_workers = 2

    try:
        with open('accounts.json') as f:
            accounts_data = json.load(f)
    except FileNotFoundError:
        print('Create accounts.json with cookies from browser: {"accounts": [{"cookies": {...}}, ...]}')
        exit(1)

    await init_trid({'User-Agent': USERAGENT})

    manager = SessionManager(accounts_data['accounts'])
    worker_args = [(worker_id, manager) for worker_id in range(num_workers)]
    pool = AioPool(num_workers, worker_fn, worker_args)

    st = time.monotonic()
    thread = await load_tree(pool, thread_id, limit_requests)

    print()
    print('*' * 80)
    print('DONE', pool.requests_count, pool.responses_count)
    print(f'TWEETS', len(thread['tweets']))
    print(f'USERS', len(thread['users']))
    print('TOOK', time.monotonic() - st)
    print()

    await pool.shutdown()

    existing_tweet_ids = {t['rest_id'] for t in thread['tweets']}
    missing_tweet_ids = list(set(find_article_referenced_tweets(thread)) - existing_tweet_ids)

    if missing_tweet_ids:
        print(f'Fetching {len(missing_tweet_ids)} article tweets...')
        session = manager.get_session()
        if not session:
            pass
        else:
            entity_response = await load_tweets_by_ids(session, missing_tweet_ids)
            entity_tweets = find_all_tweets(entity_response)
            print(f'Added {len(entity_tweets)} new tweets')
            thread['tweets'].extend(entity_tweets)
            
            entity_users = find_all_users(entity_response)
            existing_user_ids = {u['rest_id'] for u in thread['users']}
            new_users = [u for u in entity_users if u['rest_id'] not in existing_user_ids]
            
            print(f'Added {len(new_users)} new users')
            thread['users'].extend(new_users)

    temp_dir = os.path.dirname(output_path)
    with tempfile.NamedTemporaryFile(mode='w', dir=temp_dir, suffix='.tmp', delete=False, encoding='utf-8') as temp_file:
        json.dump(thread, temp_file, ensure_ascii=False)
        temp_path = temp_file.name

    shutil.move(temp_path, output_path)


if __name__ == '__main__':
    asyncio.run(main())