test_sessions.nim 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  1. # SPDX-License-Identifier: AGPL-3.0-only
  2. # Test each cookie session in a JSONL file by hitting the UserMedia endpoint.
  3. #
  4. # Usage:
  5. # nim r --path:src tools/test_sessions.nim [sessions_file] [--delay N]
  6. #
  7. # Examples:
  8. # nim r --path:src tools/test_sessions.nim
  9. # nim r --path:src tools/test_sessions.nim sessions_cookie.jsonl --delay 1
  10. import asyncdispatch, httpclient, strutils, uri, os, zippy
  11. import apiutils, auth, consts, types
  12. import experimental/parser/session
  13. const
  14. # jack (id=12) — public account with media, same target as ratelimit_probe.py
  15. testVars = userMediaVars % ["12", "", "20"]
  16. # Build the ApiReq once — same for every session
  17. let
  18. testUrl = ApiUrl(endpoint: graphUserMedia,
  19. params: @[("variables", testVars), ("features", gqlFeatures)])
  20. testReq = ApiReq(cookie: testUrl, oauth: testUrl)
  21. url = testReq.toUrl(SessionKind.cookie)
  22. proc parseCookieSessions(path: string): seq[Session] =
  23. var skipped = 0
  24. for line in path.lines:
  25. let s = line.strip()
  26. if s.len == 0: continue
  27. try:
  28. let sess = parseSession(s)
  29. if sess.kind == SessionKind.cookie:
  30. result.add sess
  31. else:
  32. inc skipped
  33. except Exception as e:
  34. echo " [!] Parse error: ", e.msg
  35. if skipped > 0:
  36. echo " (skipped ", skipped, " non-cookie sessions)"
  37. proc testSession(session: Session): Future[tuple[ok: bool, code: int,
  38. remaining, limit: int, detail: string]] {.async.} =
  39. let headers = await genHeaders(session, url, skipTid = false)
  40. let client = newAsyncHttpClient(headers = headers)
  41. defer: client.close()
  42. try:
  43. let resp = await client.get($url)
  44. var body = await resp.body
  45. if resp.headers.getOrDefault("content-encoding") == "gzip":
  46. body = uncompress(body, dfGzip)
  47. var remaining, limit: int
  48. if resp.headers.hasKey("x-rate-limit-remaining"):
  49. remaining = parseInt(resp.headers["x-rate-limit-remaining"])
  50. if resp.headers.hasKey("x-rate-limit-limit"):
  51. limit = parseInt(resp.headers["x-rate-limit-limit"])
  52. if resp.code == Http200:
  53. return (true, resp.code.int, remaining, limit, "")
  54. else:
  55. let detail = if body.len in 1..120: body else: ""
  56. return (false, resp.code.int, remaining, limit, detail)
  57. except Exception as e:
  58. return (false, 0, 0, 0, e.msg[0 ..< min(e.msg.len, 120)])
  59. proc main() {.async.} =
  60. var
  61. sessionsPath = "sessions.jsonl"
  62. delay = 500 # ms
  63. i = 1
  64. while i <= paramCount():
  65. let arg = paramStr(i)
  66. case arg
  67. of "--delay":
  68. inc i
  69. if i > paramCount():
  70. echo "Error: --delay requires a value (seconds)"; quit(1)
  71. delay = int(parseFloat(paramStr(i)) * 1000)
  72. of "--help", "-h":
  73. echo "Usage: nim r --path:src tools/test_sessions.nim [sessions_file] [--delay N]"
  74. return
  75. else:
  76. sessionsPath = arg
  77. inc i
  78. if not fileExists(sessionsPath):
  79. echo "File not found: ", sessionsPath; quit(1)
  80. setApiProxy("")
  81. setDisableTid(false)
  82. let sessions = parseCookieSessions(sessionsPath)
  83. if sessions.len == 0:
  84. echo "No cookie sessions found in ", sessionsPath; quit(0)
  85. echo "Testing ", sessions.len, " sessions from ", sessionsPath
  86. echo ""
  87. var
  88. nValid, nFail: int
  89. failures: seq[string]
  90. for idx, session in sessions:
  91. let
  92. num = idx + 1
  93. prefix = "[" & $num & "/" & $sessions.len & "] " & session.pretty
  94. res = await testSession(session)
  95. if res.ok:
  96. inc nValid
  97. var extra = ""
  98. if res.limit > 0:
  99. extra = " remaining=" & $res.remaining & "/" & $res.limit
  100. echo prefix, " ✓ ", res.code, extra
  101. else:
  102. inc nFail
  103. let detail = if res.code > 0: $res.code else: "error: " & res.detail
  104. failures.add session.pretty & " " & detail
  105. echo prefix, " ✗ ", detail
  106. if num < sessions.len and delay > 0:
  107. await sleepAsync(delay)
  108. echo ""
  109. echo "Done: ", nValid, " valid, ", nFail, " invalid [", sessions.len, " total]"
  110. if failures.len > 0:
  111. echo ""
  112. echo "Failures:"
  113. for f in failures:
  114. echo " ", f
  115. waitFor main()