identify.nim

# Copyright (c) 2020-2026 dcrawl contributors
# Licensed under either of
#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
# at your option.
# This file may not be copied, modified, or distributed except according to
# those terms.

{.push raises: [].}

import
  std/[strutils, sets, options, tables, algorithm],
  chronos, chronicles,
  eth/keys as ethKeys,
  eth/p2p/discoveryv5/[node, enr],
  libp2p/[switch, builders, multiaddress, peerid, peerstore],
  libp2p/crypto/[crypto, secp],
  ./devp2p_probe

export switch.Switch
export devp2p_probe

logScope:
  topics = "identify"

type
  IdentifyJob = object
    n: Node
    discoveryCycle: int
    guess: string  # ENR-derived guess captured at enqueue time, "" when --guess off
    sentinel: bool

  IdentifyResult = enum
    irOk, irSkip, irFail

  IdentifyProbe* = ref object
    switch*: Switch
    devp2p*: DevP2pProbe
    queue: AsyncQueue[IdentifyJob]
    next: seq[IdentifyJob]
    processedPeers: HashSet[Node]
    currentInPipeline: int
    cycleNum: int
    discoverDone: bool
    shuttingDown: bool
    workers: seq[Future[void]]
    controller: Future[void]
    statsTask: Future[void]
    agentsFile: File
    maxCycles, timeoutS, minCycleIntervalS, statsIntervalS: int
    cycleOk, cycleFail, cycleSkip: int
    totalOk, totalFail, totalSkip: int
    # Guess-vs-identify accounting (only set when --guess was on at enqueue):
    guessMatch: int               # peers where guess == identified family
    guessMismatch: int            # peers where guess != identified family
    guessMissing: int             # peers we identified but had no guess for
    mismatchTable: Table[string, int]  # "guess->identified" -> count
    startTime: Moment

proc csvQuote(s: string): string =
  if s.len == 0: ""
  else: "'" & s.replace("'", "''") & "'"

proc clientFamily(agentVersion: string): string =
  ## Extract the family name from an agent_version string. Same convention
  ## as plot.py: split on first '/' or ':'. "Lighthouse/v8.x.x/..." -> "Lighthouse",
  ## "SSV-Node:v2.x.x" -> "SSV-Node", "Geth/v1.x.x/..." -> "Geth".
  for i, c in agentVersion:
    if c in {'/', ':'}: return agentVersion[0 ..< i]
  agentVersion

proc recordGuessOutcome(probe: IdentifyProbe, guess, identifiedFamily: string) =
  ## Compare guess vs identified client family and update accounting.
  ## Multi-label guesses use `|` to separate candidates; any candidate
  ## matching the identified family counts as a match.
  if guess.len == 0:
    probe.guessMissing.inc
    return
  if identifiedFamily.len == 0:
    return  # nothing to compare against
  var matched = false
  for cand in guess.split('|'):
    if cand == identifiedFamily:
      matched = true
      break
  if matched:
    probe.guessMatch.inc
  else:
    probe.guessMismatch.inc
    let key = guess & "->" & identifiedFamily
    probe.mismatchTable[key] = probe.mismatchTable.getOrDefault(key) + 1

proc writeAgentRow(probe: IdentifyProbe, job: IdentifyJob, ipStr: string,
                   tcpPort: int, agent, pver, protos, stack: string,
                   disconnectReason: string = "") =
  try:
    probe.agentsFile.writeLine(
      $job.discoveryCycle & "," & job.n.id.toHex & "," &
      ipStr & ":" & $tcpPort & "," &
      csvQuote(agent) & "," & csvQuote(pver) & "," &
      csvQuote(protos) & "," & $probe.cycleNum & "," & stack & "," &
      disconnectReason & "," & job.guess)
    probe.agentsFile.flushFile()
  except IOError as e:
    warn "agents.csv write failed", error=e.msg
  recordGuessOutcome(probe, job.guess, clientFamily(agent))

proc tryLibp2pIdentify(probe: IdentifyProbe, job: IdentifyJob): Future[IdentifyResult] {.
    async: (raises: [CancelledError]).} =
  let n = job.n
  let tcpPortOpt = n.record.tryGet("tcp", int)
  if tcpPortOpt.isNone: return irSkip
  let tcpPort = tcpPortOpt.get()

  let pubkeyOpt = n.record.tryGet("secp256k1", seq[byte])
  if pubkeyOpt.isNone: return irSkip

  let skPubRes = secp.SkPublicKey.init(pubkeyOpt.get())
  if skPubRes.isErr: return irSkip
  let lpPub = crypto.PublicKey(scheme: PKScheme.Secp256k1, skkey: skPubRes.get())
  let peerIdRes = PeerId.init(lpPub)
  if peerIdRes.isErr: return irSkip
  let peerId = peerIdRes.get()

  if n.address.isNone: return irSkip
  let ipStr = $n.address.get().ip
  let maRes = MultiAddress.init("/ip4/" & ipStr & "/tcp/" & $tcpPort)
  if maRes.isErr: return irSkip
  let ma = maRes.get()

  let attempt = probe.cycleNum
  let tStart = now(chronos.Moment)

  try:
    await probe.switch.connect(peerId, @[ma]).wait(probe.timeoutS.seconds)
  except CancelledError as e:
    raise e
  except CatchableError as e:
    let dialMs = (now(chronos.Moment) - tStart).milliseconds
    progressLog.info "identify connect failed",
      id=n.id.toHex, error=e.msg, dialMs, attempt, stack="libp2p"
    return irFail

  let dialMs = (now(chronos.Moment) - tStart).milliseconds
  try:
    let agent = probe.switch.peerStore[AgentBook][peerId]
    let pver = probe.switch.peerStore[ProtoVersionBook][peerId]
    let protos = probe.switch.peerStore[ProtoBook][peerId].join(";")
    writeAgentRow(probe, job, ipStr, tcpPort, agent, pver, protos, "libp2p")
    progressLog.info "identify ok",
      id=n.id.toHex, agent, dialMs, attempt, stack="libp2p"
  except CatchableError as e:
    debug "identify peerstore read failed", id=n.id.toHex, error=e.msg

  try:
    await probe.switch.disconnect(peerId)
  except CancelledError as e:
    raise e
  except CatchableError:
    discard

  return irOk

proc tryDevp2pIdentify(probe: IdentifyProbe, job: IdentifyJob): Future[IdentifyResult] {.
    async: (raises: [CancelledError]).} =
  let n = job.n
  if probe.devp2p.isNil: return irSkip
  let tcpPortOpt = n.record.tryGet("tcp", int)
  if tcpPortOpt.isNone: return irSkip
  if n.address.isNone: return irSkip

  let attempt = probe.cycleNum
  let tStart = now(chronos.Moment)
  let helloRes = await tryDevp2pHello(probe.devp2p, n, probe.timeoutS)
  let dialMs = (now(chronos.Moment) - tStart).milliseconds

  if helloRes.isNone:
    debug "devp2p connect/handshake failed",
      id=n.id.toHex, dialMs, attempt, stack="devp2p"
    return irFail

  let r = helloRes.get()
  let drStr =
    if r.disconnectReason.isSome: $int(r.disconnectReason.get())
    else: ""
  let ipStr = $n.address.get().ip
  writeAgentRow(probe, job, ipStr, tcpPortOpt.get(), r.clientId, "", "",
                "devp2p", drStr)
  progressLog.info "identify ok",
    id=n.id.toHex, agent=r.clientId, dialMs, attempt, stack="devp2p",
    disconnectReason = drStr
  return irOk

proc tryIdentify(probe: IdentifyProbe, job: IdentifyJob): Future[IdentifyResult] {.
    async: (raises: [CancelledError]).} =
  ## Route per ENR keys: pure-EL nodes go to devp2p, everything else
  ## (CL or unknown) goes to libp2p.
  let enrStr = $job.n.record
  let hasEth = " eth:" in enrStr
  let hasEth2 = " eth2:" in enrStr
  if hasEth and not hasEth2 and not probe.devp2p.isNil:
    return await tryDevp2pIdentify(probe, job)
  return await tryLibp2pIdentify(probe, job)

proc worker(probe: IdentifyProbe) {.async: (raises: [CancelledError]).} =
  while not probe.shuttingDown:
    let job = await probe.queue.popFirst()
    if job.sentinel: break

    let res =
      try:
        await tryIdentify(probe, job)
      except CancelledError as e:
        probe.currentInPipeline.dec
        raise e

    case res
    of irOk:
      probe.cycleOk.inc; probe.totalOk.inc
      probe.currentInPipeline.dec
    of irSkip:
      probe.cycleSkip.inc; probe.totalSkip.inc
      probe.currentInPipeline.dec
    of irFail:
      probe.cycleFail.inc; probe.totalFail.inc
      if probe.cycleNum < probe.maxCycles - 1:
        probe.next.add(job)
      probe.currentInPipeline.dec

proc cycleController(probe: IdentifyProbe) {.async: (raises: [CancelledError]).} =
  while probe.cycleNum < probe.maxCycles and not probe.shuttingDown:
    let cycleStart = now(chronos.Moment)

    while true:
      if probe.currentInPipeline == 0:
        if probe.cycleNum > 0 or probe.discoverDone: break
      await sleepAsync(100.milliseconds)

    info "identify cycle complete",
      cycle = probe.cycleNum,
      retries_pending = probe.next.len,
      cycle_duration_s = (now(chronos.Moment) - cycleStart).milliseconds div 1000,
      cycle_ok = probe.cycleOk,
      cycle_fail = probe.cycleFail,
      cycle_skip = probe.cycleSkip

    probe.cycleOk = 0
    probe.cycleFail = 0
    probe.cycleSkip = 0

    if probe.next.len == 0: break
    if probe.cycleNum >= probe.maxCycles - 1: break

    let elapsed = now(chronos.Moment) - cycleStart
    let interval = chronos.seconds(probe.minCycleIntervalS)
    if elapsed < interval:
      await sleepAsync(interval - elapsed)

    var toPromote: seq[IdentifyJob]
    swap(toPromote, probe.next)
    probe.currentInPipeline = toPromote.len
    probe.cycleNum.inc
    info "identify cycle starting", cycle = probe.cycleNum, jobs = toPromote.len
    for job in toPromote:
      try: probe.queue.addLastNoWait(job)
      except AsyncQueueFullError: discard

  probe.shuttingDown = true
  for _ in probe.workers:
    try: probe.queue.addLastNoWait(IdentifyJob(sentinel: true))
    except AsyncQueueFullError: discard

proc statsReporter(probe: IdentifyProbe) {.async: (raises: [CancelledError]).} =
  while not probe.shuttingDown:
    await sleepAsync(probe.statsIntervalS.seconds)
    if probe.shuttingDown: break
    let wallS = (now(chronos.Moment) - probe.startTime).milliseconds div 1000
    info "identify stats",
      wall_s = wallS,
      cycle = probe.cycleNum,
      ok = probe.cycleOk,
      fail = probe.cycleFail,
      skip = probe.cycleSkip,
      retry_pending = probe.next.len,
      in_pipeline = probe.currentInPipeline,
      processed = probe.processedPeers.len,
      total_ok = probe.totalOk,
      guess_match = probe.guessMatch,
      guess_mismatch = probe.guessMismatch
    # Mirror to progressLog so postprocessing/progress.py can plot the
    # guess match/mismatch series over time. Default stream (above) keeps
    # stdout informed; progressLog lands in dcrawl.log only.
    progressLog.info "identify stats",
      wall_s = wallS,
      cycle = probe.cycleNum,
      ok = probe.cycleOk,
      fail = probe.cycleFail,
      skip = probe.cycleSkip,
      retry_pending = probe.next.len,
      in_pipeline = probe.currentInPipeline,
      processed = probe.processedPeers.len,
      total_ok = probe.totalOk,
      guess_match = probe.guessMatch,
      guess_mismatch = probe.guessMismatch

proc createIdentifyProbe*(
    maxConcurrent, maxCycles, timeoutS, minCycleIntervalS, statsIntervalS: int,
    agentsFile: string,
    devp2p: DevP2pProbe = nil): Future[IdentifyProbe] {.
    async: (raises: [CatchableError]).} =
  let rng = crypto.newRng()
  let s = SwitchBuilder.new()
    .withRng(rng)
    .withAddresses(@[])
    .withMaxConnections(maxConcurrent * 4)
    .withTcpTransport()
    .withYamux()
    .withMplex()
    .withNoise()
    .build()
  await s.start()

  let f = open(agentsFile, fmWrite)
  f.writeLine("cycle,node_id,ip:port,agent_version,protocol_version,protocols,attempt,stack,disconnect_reason,guess")

  let probe = IdentifyProbe(
    switch: s,
    devp2p: devp2p,
    queue: newAsyncQueue[IdentifyJob](),
    agentsFile: f,
    maxCycles: maxCycles,
    timeoutS: timeoutS,
    minCycleIntervalS: minCycleIntervalS,
    statsIntervalS: statsIntervalS,
    mismatchTable: initTable[string, int](),
    startTime: now(chronos.Moment),
  )

  for _ in 0 ..< maxConcurrent:
    probe.workers.add(worker(probe))
  probe.controller = cycleController(probe)
  probe.statsTask = statsReporter(probe)

  return probe

proc enqueue*(probe: IdentifyProbe, n: Node, discoveryCycle: int,
              guess: string = "") =
  if probe.processedPeers.containsOrIncl(n):
    return
  let job = IdentifyJob(
    n: n, discoveryCycle: discoveryCycle, guess: guess)
  try:
    probe.queue.addLastNoWait(job)
    probe.currentInPipeline.inc
  except AsyncQueueFullError:
    discard

proc markDiscoverDone*(probe: IdentifyProbe) =
  probe.discoverDone = true

proc drain*(probe: IdentifyProbe) {.async: (raises: [CancelledError]).} =
  try:
    await probe.controller
  except CancelledError as e:
    raise e
  except CatchableError as e:
    error "identify controller failed", error=e.msg
  await allFutures(probe.workers)
  if not probe.statsTask.isNil:
    await probe.statsTask.cancelAndWait()
  let wallS = (now(chronos.Moment) - probe.startTime).milliseconds div 1000
  info "identify final",
    wall_s = wallS,
    cycles_run = probe.cycleNum + 1,
    total_ok = probe.totalOk,
    total_fail = probe.totalFail,
    total_skip = probe.totalSkip,
    processed = probe.processedPeers.len,
    retries_dropped = probe.next.len,
    guess_match = probe.guessMatch,
    guess_mismatch = probe.guessMismatch,
    guess_missing = probe.guessMissing
  # Top mismatch patterns — surfaces stale rules.
  if probe.mismatchTable.len > 0:
    var pairs: seq[(string, int)] = @[]
    for k, v in probe.mismatchTable.pairs: pairs.add((k, v))
    pairs.sort(proc(a, b: (string, int)): int = cmp(b[1], a[1]))
    let topN = min(10, pairs.len)
    info "guess mismatch breakdown (top)", entries = pairs[0 ..< topN]
  if not probe.agentsFile.isNil:
    probe.agentsFile.close()

{.pop.}