#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Python port of Nitter's tools/test_sessions.nim
#
# Test each cookie session in a JSONL file by hitting the UserMedia endpoint.
#
# Usage:
#   python3 test_sessions.py [sessions_file] [--delay N]
#
# Examples:
#   python3 test_sessions.py
#   python3 test_sessions.py sessions_cookie.jsonl --delay 1
#
# Requires: pip install requests

import base64
import hashlib
import json
import os
import random
import sys
import time
from urllib.parse import quote, urlencode, urlparse

import requests

# ---------------------------------------------------------------- consts.nim

BEARER_TOKEN = "Bearer AAAAAAAAAAAAAAAAAAAAANRILgAAAAAAnNwIzUejRCOuH5E6I8xnZz4puTs%3D1Zv7ttfk8LF81IUq16cHjhLTvJu4FA33AGWWjCpTnA"
BEARER_TOKEN2 = "Bearer AAAAAAAAAAAAAAAAAAAAAFXzAwAAAAAAMHCxpeSDG1gLNLghVe8d74hl6k4%3DRUMF4xAQLsbeBhTSRrCiQpJtxoGWeyHrDb5te2jpGskWDFW82F"

GRAPH_USER_MEDIA = "VyudDWQnr9vJNw7GasFz2g/UserMedia"


def _squash(s: str) -> str:
    return s.replace(" ", "").replace("\n", "")


GQL_FEATURES = _squash("""{
  "rweb_video_screen_enabled": false,
  "rweb_cashtags_enabled": true,
  "profile_label_improvements_pcf_label_in_post_enabled": true,
  "responsive_web_profile_redirect_enabled": false,
  "rweb_tipjar_consumption_enabled": false,
  "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,
  "withQuickPromoteEligibilityTweetFields": true
}""")

USER_MEDIA_VARS = _squash("""{
  "userId": "%s", %s
  "count": %s,
  "includePromotedContent": false,
  "withClientEventToken": false,
  "withBirdwatchNotes": false,
  "withVoice": true
}""")

# jack (id=12) - public account with media, same target as ratelimit_probe.py
TEST_VARS = USER_MEDIA_VARS % ("12", "", "20")

# --------------------------------------------------------------------- tid.nim

DEFAULT_KEYWORD = "obfiowerehiring"
PAIRS_URL = ("https://raw.githubusercontent.com/fa0311/"
             "x-client-transaction-id-pair-dict/refs/heads/main/pair.json")
PAIRS_TTL_SEC = 60 * 60  # refresh every hour

_cached_pairs: list = []
_last_cached = 0


def _get_pair() -> dict:
    global _cached_pairs, _last_cached
    if not _cached_pairs or int(time.time()) - _last_cached > PAIRS_TTL_SEC:
        resp = _http.get(PAIRS_URL, timeout=30)
        if resp.status_code == 200:
            pairs = resp.json()
            if not pairs:
                raise ValueError("Parsing pairs failed: " + resp.text)
            _cached_pairs = pairs
            _last_cached = int(time.time())

    if not _cached_pairs:
        raise ValueError("Failed to fetch x-client-transaction-id pairs")

    return random.choice(_cached_pairs)


def _b64decode_nopad(data: str) -> bytes:
    return base64.b64decode(data + "=" * (-len(data) % 4))


def gen_tid(path: str) -> str:
    pair = _get_pair()

    time_now = int(time.time() - 1682924400)
    time_now_bytes = bytes([
        time_now & 0xFF,
        (time_now >> 8) & 0xFF,
        (time_now >> 16) & 0xFF,
        (time_now >> 24) & 0xFF,
    ])

    data = "GET!" + path + "!" + str(time_now) + DEFAULT_KEYWORD + pair["animationKey"]
    hash_bytes = hashlib.sha256(data.encode()).digest()
    key_bytes = _b64decode_nopad(pair["verification"])
    bytes_arr = key_bytes + time_now_bytes + hash_bytes[:16] + bytes([3])
    random_num = random.randrange(256)
    tid = bytes([random_num]) + bytes(b ^ random_num for b in bytes_arr)

    return base64.b64encode(tid).decode().replace("=", "")


# ------------------------------------------------------- types / auth / session

class Session:
    def __init__(self, kind, id, username, auth_token="", ct0="",
                 oauth_token="", oauth_secret=""):
        self.kind = kind  # "oauth" | "cookie"
        self.id = id
        self.username = username
        self.auth_token = auth_token
        self.ct0 = ct0
        self.oauth_token = oauth_token
        self.oauth_secret = oauth_secret

    def pretty(self) -> str:
        if self.id > 0 and self.username:
            result = f"{self.id} ({self.username})"
        elif self.username:
            result = self.username
        elif self.id > 0:
            result = str(self.id)
        else:
            result = "<unknown>"
        return f"{self.kind} {result}"


def parse_session(raw: str) -> Session:
    s = json.loads(raw)
    kind = s.get("kind") or "oauth"
    username = s.get("username", "") or ""

    if kind == "oauth":
        token = s.get("oauth_token", "") or ""
        # mirrors Nim's token[0 ..< token.find('-')]
        sid = token[:token.find("-")]
        return Session("oauth", int(sid), username,
                       oauth_token=token,
                       oauth_secret=s.get("oauth_token_secret", "") or "")
    if kind == "cookie":
        raw_id = str(s.get("id", "") or "")
        return Session("cookie", int(raw_id) if raw_id else 0, username,
                       auth_token=s.get("auth_token", "") or "",
                       ct0=s.get("ct0", "") or "")
    raise ValueError("Unknown session kind: " + kind)


# ------------------------------------------------------------------ apiutils.nim

def build_url() -> str:
    params = [("variables", TEST_VARS), ("features", GQL_FEATURES)]
    prefix = "" if GRAPH_USER_MEDIA.startswith("1.1/") else "graphql/"
    return ("https://x.com/i/api/" + prefix + GRAPH_USER_MEDIA
            + "?" + urlencode(params, quote_via=quote))


def gen_headers(session: Session, url: str, skip_tid: bool, disable_tid: bool) -> dict:
    h = {
        "Accept": "*/*",
        "Accept-Encoding": "gzip",
        "Accept-Language": "en-US,en;q=0.9",
        "Content-Type": "application/json",
        "Origin": "https://x.com",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                      "(KHTML, like Gecko) Chrome/142.0.0.0 Safari/537.36",
        "X-Twitter-Active-User": "yes",
        "X-Twitter-Client-Language": "en",
        "Priority": "u=1, i",
    }

    # only cookie sessions are tested, so the oauth branch is never reached
    path = urlparse(url).path
    h["X-Twitter-Auth-Type"] = "OAuth2Session"
    h["X-Csrf-Token"] = session.ct0
    h["Cookie"] = f"auth_token={session.auth_token}; ct0={session.ct0}"
    h["Referer"] = "https://x.com/"
    h["Sec-Ch-Ua"] = '"Google Chrome";v="142", "Chromium";v="142", "Not A(Brand";v="24"'
    h["Sec-Ch-Ua-Mobile"] = "?0"
    h["Sec-Ch-Ua-Platform"] = "Windows"
    h["Sec-Fetch-Dest"] = "empty"
    h["Sec-Fetch-Mode"] = "cors"
    h["Sec-Fetch-Site"] = "same-origin"
    if disable_tid or skip_tid or "/1.1/" in path:
        h["Authorization"] = BEARER_TOKEN2
    else:
        h["Authorization"] = BEARER_TOKEN
        h["X-Client-Transaction-Id"] = gen_tid(path)
    return h


# ----------------------------------------------------------- test_sessions.nim

# no proxy, and ignore HTTP(S)_PROXY from the environment (Nim's client does too)
_http = requests.Session()
_http.trust_env = False


def parse_cookie_sessions(path: str) -> list:
    result = []
    skipped = 0
    with open(path, encoding="utf-8") as f:
        for line in f:
            s = line.strip()
            if not s:
                continue
            try:
                sess = parse_session(s)
                if sess.kind == "cookie":
                    result.append(sess)
                else:
                    skipped += 1
            except Exception as e:
                print("  [!] Parse error:", e)
    if skipped > 0:
        print("  (skipped", skipped, "non-cookie sessions)")
    return result


def test_session(session: Session, url: str):
    """Returns (ok, code, remaining, limit, reset, detail)."""
    headers = gen_headers(session, url, skip_tid=False, disable_tid=False)

    try:
        resp = _http.get(url, headers=headers, timeout=30)
        body = resp.text  # requests transparently handles gzip

        remaining = limit = reset = 0
        if "x-rate-limit-reset" in resp.headers:
            reset = int(resp.headers["x-rate-limit-reset"])
        if "x-rate-limit-remaining" in resp.headers:
            remaining = int(resp.headers["x-rate-limit-remaining"])
        if "x-rate-limit-limit" in resp.headers:
            limit = int(resp.headers["x-rate-limit-limit"])

        if resp.status_code == 200:
            return (True, resp.status_code, remaining, limit, reset, "")
        detail = body if 1 <= len(body) <= 120 else ""
        return (False, resp.status_code, remaining, limit, reset, detail)
    except Exception as e:
        return (False, 0, 0, 0, 0, str(e)[:120])


def main():
    sessions_path = "sessions.jsonl"
    delay = 500  # ms
    args = sys.argv[1:]
    i = 0

    while i < len(args):
        arg = args[i]
        if arg == "--delay":
            i += 1
            if i >= len(args):
                print("Error: --delay requires a value (seconds)")
                sys.exit(1)
            delay = int(float(args[i]) * 1000)
        elif arg in ("--help", "-h"):
            print("Usage: test_sessions.py [sessions_file] [--delay N]")
            return
        else:
            sessions_path = arg
        i += 1

    if not os.path.isfile(sessions_path):
        print("File not found:", sessions_path)
        sys.exit(1)

    sessions = parse_cookie_sessions(sessions_path)
    if not sessions:
        print("No cookie sessions found in", sessions_path)
        sys.exit(0)

    print("Testing", len(sessions), "sessions from", sessions_path)
    print("")

    url = build_url()
    n_valid = n_fail = 0
    failures = []

    for idx, session in enumerate(sessions):
        num = idx + 1
        prefix = f"[{num}/{len(sessions)}] {session.pretty()}"
        ok, code, remaining, limit, reset, detail = test_session(session, url)

        if ok:
            n_valid += 1
            extra = f"  remaining={remaining}/{limit}" if limit > 0 else ""
            if reset > 0:
                mins = max(0, reset - int(time.time())) // 60
                extra += f"  resets in {mins}m ({time.strftime('%H:%M:%S', time.localtime(reset))})"
            print(prefix, "  \u2713 ", code, extra, sep="")
        else:
            n_fail += 1
            d = str(code) if code > 0 else "error: " + detail
            failures.append(session.pretty() + "  " + d)
            print(prefix, "  \u2717 ", d, sep="")

        if num < len(sessions) and delay > 0:
            time.sleep(delay / 1000)

    print("")
    print(f"Done: {n_valid} valid, {n_fail} invalid  [{len(sessions)} total]")

    if failures:
        print("")
        print("Failures:")
        for f in failures:
            print("  " + f)


if __name__ == "__main__":
    main()