# tickwatch
# Author: Euxane TRAN-GIRARD
# Licence: EUPL-1.2

import std/sequtils
import std/strutils
import std/random
import std/os
import std/net
import std/nativesockets
import std/posix
import std/times
import std/monotimes


type
  ICMPPacket = object
    ## https://datatracker.ietf.org/doc/html/rfc792
    ## https://datatracker.ietf.org/doc/html/rfc4443#section-4.1
    typ: uint8
    code: uint8
    checksum: uint16
    ident: uint16
    seqNum: uint16
    payload: string

  IPv6PseudoHeader = object
    ## https://datatracker.ietf.org/doc/html/rfc2460#section-8.1
    source: IpAddress
    dest: IpAddress
    length: uint32
    next: uint8


const
  ICMP_ECHO_REQUEST = 8u8
  ICMP_ECHO_REPLY = 0u8
  ICMPv6_ECHO_REQUEST = 128u8
  ICMPv6_ECHO_REPLY = 129u8
  ICMP_IPv4_HEADER_LENGTH = 20
  RANDOM_PAYLOAD_LENGTH = 32
  ICMP_PACKET_BASE_LENGTH = 8
  DEFAULT_TIMEOUT = initDuration(seconds = 1)


func pack16(n: uint16): array[2, uint8] =
  cast[array[2, uint8]](nativesockets.htons(n))

func pack32(n: uint32): array[4, uint8] =
  cast[array[4, uint8]](nativesockets.htonl(n))

func unpack16(a: array[2, uint8]): uint16 =
  nativesockets.ntohs(cast[uint16](a))

func unpack32(a: array[4, uint8]): uint32 =
  nativesockets.ntohl(cast[uint32](a))

func pack(p: ICMPPacket): seq[uint8] =
  result = newSeqOfCap[uint8](ICMP_PACKET_BASE_LENGTH + p.payload.len)
  result.add p.typ
  result.add p.code
  result.add p.checksum.pack16
  result.add p.ident.pack16
  result.add p.seqNum.pack16
  result.add cast[seq[uint8]](p.payload)

func unpackICMPPacket(a: openArray[uint8]): ICMPPacket =
  result.typ = a[0]
  result.code = a[1]
  result.checksum = [a[2], a[3]].unpack16
  result.ident = [a[4], a[5]].unpack16
  result.seqNum = [a[6], a[7]].unpack16
  result.payload = cast[string](a[8..^1])

func pack(h: IPv6PseudoHeader): seq[uint8] =
  result = newSeqOfCap[uint8](40)
  result.add h.source.address_v6
  result.add h.dest.address_v6
  result.add h.length.pack32
  result.add [0u8, 0, 0]
  result.add h.next

func checksum(a: openArray[uint8]): uint16 =
  ## https://datatracker.ietf.org/doc/html/rfc1071
  var sum = 0u32
  for i in countup(0, a.len - 1, 2):
    sum += uint32(a[i]) shl 8
    if i + 1 < a.len:
      sum += uint32(a[i + 1])

  while (sum shr 16) != 0:
    sum = (sum and 0xFFFF) + (sum shr 16)

  not uint16(sum)

func icmpV6Header(source, dest: IpAddress, length: int): IPv6PseudoHeader =
  IPv6PseudoHeader(
    source: source,
    dest: dest,
    length: length.uint32,
    next: posix.IPPROTO_ICMPv6.uint8,
  )

func buildEchoRequest(
  source, dest: IpAddress,
  ident, seqNum: uint16,
  payload: string = "",
): ICMPPacket =
  result = ICMPPacket(ident: ident, seqNum: seqNum, payload: payload)
  case dest.family:
  of IPv6:
    result.typ = ICMPv6_ECHO_REQUEST
    let packedICMP = result.pack
    let pseudoHeader = icmpV6Header(source, dest, packedICMP.len).pack
    result.checksum = checksum(pseudoHeader & packedICMP)
  of IPv4:
    result.typ = ICMP_ECHO_REQUEST
    result.checksum = checksum(result.pack)

func validateChecksum(packet: ICMPPacket, source, dest: IpAddress): bool =
  var sourcePacket = packet
  sourcePacket.checksum = 0
  case packet.typ:
  of ICMPv6_ECHO_REPLY:
    let packedICMP = sourcePacket.pack
    let pseudoHeader = icmpV6Header(source, dest, packedICMP.len).pack
    packet.checksum == checksum(pseudoHeader & packedICMP)
  of ICMP_ECHO_REPLY:
    packet.checksum == checksum(sourcePacket.pack)
  else:
    false

func domainProtocol(family: IpAddressFamily): (Domain, Protocol) =
  case family:
  of IPv6: (AF_INET6, IPPROTO_ICMPv6)
  of IPv4: (AF_INET, IPPROTO_ICMP)

func icmpReplyType(family: IpAddressFamily): uint8 =
  case family:
  of IPv6: ICMPv6_ECHO_REPLY
  of IPv4: ICMP_ECHO_REPLY

func stripIPHeader(family: IpAddressFamily, buffer: seq[uint8]): seq[uint8] =
  case family:
  of IPv6: buffer
  of IPv4: buffer[min(ICMP_IPv4_HEADER_LENGTH, buffer.len)..^1]

func toTimeval(d: Duration): Timeval =
  let uSecs = d.inMicroseconds
  Timeval(
    tv_sec: posix.Time(uSecs div (1000 * 1000)),
    tv_usec: posix.Suseconds(uSecs mod (1000 * 1000)),
  )

proc setTimeout(sock: Socket, opt: cint, timeout: Duration) =
  let timeVal = timeout.toTimeval
  let timeValLen = sizeof(timeVal).SockLen
  if sock.getFd().setSockOpt(SOL_SOCKET, opt, timeVal.addr, timeValLen) < 0:
    raiseOsError osLastError()

proc receiveReply(sock: Socket, timeout: Duration): seq[uint8] =
  var buffer: string
  var ipAddr: IpAddress
  var port: Port
  sock.setTimeout(SO_RCVTIMEO, timeout)
  discard sock.recvFrom(buffer, 1024, ipAddr, port)
  stripIPHeader(ipAddr.family, cast[seq[uint8]](buffer))


func toIpAddress(info: ptr AddrInfo): IpAddress =
  case info.ai_family:
  of posix.AF_INET6:
    let sockAddr = cast[ptr posix.SockAddr_in6](info.ai_addr)[]
    IpAddress(
      family: IPv6,
      address_v6: cast[array[0..15, uint8]](sockAddr.sin6_addr.s6_addr),
    )
  of posix.AF_INET:
    let sockAddr = cast[ptr posix.SockAddr_in](info.ai_addr)[]
    IpAddress(
      family: IPv4,
      address_v4: cast[array[0..3, uint8]](sockAddr.sin_addr.s_addr),
    )
  else:
    raise newException(ValueError, "Invalid address info")

func splitDomainHostname(target: string): (Domain, string) =
  let parts = target.split('/', 1)
  case parts[0]:
  of "6": (AF_INET6, parts[1])
  of "4": (AF_INET, parts[1])
  else: (AF_UNSPEC, target)

proc resolve*(target: string): IpAddress =
  try:
    parseIpAddress(target)
  except ValueError:
    let (domain, hostname) = splitDomainHostname(target)
    let addrInfo = getAddrInfo(hostname, Port 0, domain)
    defer: freeAddrInfo(addrInfo)
    addrInfo.toIpAddress


type PingMonitor = ref object
  socket: Socket
  target: IpAddress
  source: IpAddress
  ident: uint16
  payload: string
  seqNum: uint16

proc procIdent*(): uint16 =
  uint16(getCurrentProcessId() and 0xFFFF)

proc initPingMonitor*(target: IpAddress, ident: uint16): PingMonitor =
  let (domain, proto) = target.family.domainProtocol
  PingMonitor(
    socket: newSocket(domain, SOCK_RAW, proto),
    target: target,
    ident: ident,
  )

proc buildEchoRequest(mon: PingMonitor): ICMPPacket =
  mon.source = getPrimaryIPAddr(dest=mon.target)
  buildEchoRequest(mon.source, mon.target, mon.ident, mon.seqNum, mon.payload)

proc sendEchoRequest(mon: PingMonitor, timeout: Duration) =
  try:
    let echoRequest = mon.buildEchoRequest
    mon.socket.setTimeout(SO_SNDTIMEO, timeout)
    mon.socket.sendTo(mon.target, Port 0, cast[string](echoRequest.pack))
  except CatchableError:
    discard

func validateICMPReply(mon: PingMonitor, packet: ICMPPacket): bool =
  packet.typ == icmpReplyType(mon.target.family) and
  packet.ident == mon.ident and
  packet.seqNum == mon.seqNum and
  packet.payload == mon.payload and
  packet.validateChecksum(mon.target, mon.source)

proc tryReceiveReply(mon: PingMonitor, timeout: Duration): bool =
  try:
    let data = mon.socket.receiveReply(timeout)
    if data.len < ICMP_PACKET_BASE_LENGTH: return false
    let unpacked = unpackICMPPacket(data)
    mon.validateICMPReply(unpacked)
  except CatchableError:
    false

proc genRandomPayload(length = RANDOM_PAYLOAD_LENGTH): string =
  newSeqWith(length, char.rand).join

proc ping*(mon: PingMonitor, timeout = DEFAULT_TIMEOUT): Duration =
  let startTime = getMonoTime()
  let replyDeadline = startTime + timeout

  inc mon.seqNum
  mon.payload = genRandomPayload()
  mon.sendEchoRequest(timeout)

  while true:
    let receiveTimeout = replyDeadline - getMonoTime()
    if receiveTimeout <= DurationZero:
      break
    if mon.tryReceiveReply(receiveTimeout):
      return getMonoTime() - startTime

  raise newException(TimeoutError, "Ping timeout")


when defined(test):
  import std/unittest
  suite "ping":
    test "pack/unpack endianness":
      check pack16(256) == [1u8, 0]
      check unpack16([1u8, 0]) == 256
      check pack32(256) == [0u8, 0, 1, 0]
      check unpack32([0u8, 0, 1, 0]) == 256

    test "pack/unpack ICMP packet":
      let packet = ICMPPacket(
        typ: 1,
        code: 2,
        checksum: 3,
        ident: 4,
        seqNum: 5,
        payload: cast[string](@[6u8, 7]),
      )
      let packed = pack(packet)
      check packed == [1u8, 2, 0, 3, 0, 4, 0, 5, 6, 7]
      check unpackICMPPacket(packed) == packet

    test "pack IPv6 pseudo-header":
      let pseudoHeader = IPv6PseudoHeader(
        source: parseIpAddress("::1"),
        dest: parseIpAddress("::2"),
        length: 1234,
        next: 123,
      )
      check pseudoHeader.pack == [
        0u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2,
        0, 0, 4, 210, 0, 0, 0, 123,
      ]

    test "packet checksum":
      # numeric example from https://datatracker.ietf.org/doc/html/rfc1071
      let exampleBytes = [0x00u8, 0x01, 0xf2, 0x03, 0xf4, 0xf5, 0xf6, 0xf7]
      check checksum(exampleBytes) == not 0xddf2u16

    test "ping localhost":
      for target in ["127.0.0.1", "::1"]:
        let targetIp = parseIpAddress(target)
        let mon = try:
          initPingMonitor(targetIp, procIdent())
        except OSError:
          skip()  # ping requires raw socket permission
          break
        discard mon.ping()