nitter/src/tokens.nim

# SPDX-License-Identifier: AGPL-3.0-only
import asyncdispatch, httpclient, times, sequtils, json, math, random
import strutils, strformat
import zippy
import types, agents, consts, http_pool

const
  expirationTime = 3.hours
  maxLastUse = 1.hours
  resetPeriod = 15.minutes
  failDelay = initDuration(minutes=30)

var
  clientPool {.threadvar.}: HttpPool
  tokenPool {.threadvar.}: seq[Token]
  lastFailed: Time

proc getPoolInfo*: string =
  if tokenPool.len == 0: return "token pool empty"

  let avg = tokenPool.mapIt(it.remaining).sum() div tokenPool.len
  return &"{tokenPool.len} tokens, average remaining: {avg}"

proc rateLimitError*(): ref RateLimitError =
  newException(RateLimitError, "rate limited with " & getPoolInfo())

proc fetchToken(): Future[Token] {.async.} =
  if getTime() - lastFailed < failDelay:
    raise rateLimitError()

  let headers = newHttpHeaders({
    "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
    "accept-encoding": "gzip",
    "accept-language": "en-US,en;q=0.5",
    "connection": "keep-alive",
    "user-agent": getAgent(),
    "authorization": auth
  })

  var
    resp: string
    tokNode: JsonNode
    tok: string

  try:
    resp = clientPool.use(headers): await c.postContent(activate)
    tokNode = parseJson(uncompress(resp))["guest_token"]
    tok = tokNode.getStr($(tokNode.getInt))

    let time = getTime()
    result = Token(tok: tok, remaining: 187, reset: time + resetPeriod,
                   init: time, lastUse: time)
  except Exception as e:
    lastFailed = getTime()
    echo "fetching token failed: ", e.msg

template expired(token: Token): untyped =
  let time = getTime()
  token.init < time - expirationTime or
    token.lastUse < time - maxLastUse

template isLimited(token: Token): untyped =
  token == nil or (token.remaining <= 5 and token.reset > getTime()) or
    token.expired

proc release*(token: Token; invalid=false) =
  if token != nil and (invalid or token.expired):
    let idx = tokenPool.find(token)
    if idx > -1: tokenPool.delete(idx)

proc getToken*(): Future[Token] {.async.} =
  for i in 0 ..< tokenPool.len:
    if not result.isLimited: break
    release(result)
    result = tokenPool.sample()

  if result.isLimited:
    release(result)
    result = await fetchToken()
    tokenPool.add result

  if result == nil:
    raise rateLimitError()

proc poolTokens*(amount: int) {.async.} =
  var futs: seq[Future[Token]]
  for i in 0 ..< amount:
    futs.add fetchToken()

  for token in futs:
    var newToken: Token

    try: newToken = await token
    except: discard

    if newToken != nil:
      tokenPool.add newToken

proc initTokenPool*(cfg: Config) {.async.} =
  clientPool = HttpPool()

  while true:
    if tokenPool.countIt(not it.isLimited) < cfg.minTokens:
      await poolTokens(min(4, cfg.minTokens - tokenPool.len))
    await sleepAsync(2000)