import asyncdispatch, asyncnet, strformat
from net import IpAddress, Port, `$`, `==`
from nativesockets import setSockOptInt
import ip_packet
import network_interface
import port_prediction
import puncher
import raw_socket
import utils

export PunchHoleError

type
  Attempt = object
    srcIp: IpAddress
    srcPort: Port
    dstIp: IpAddress
    dstPorts: seq[Port]

  TcpSyniConnectPuncher* = Puncher[Attempt]

  PunchProgressCb* = proc(seqNums: seq[uint32]) {.async.}

var IPPROTO_IP {.importc: "IPPROTO_IP", header: "<netinet/in.h>".}: cint
var IP_TTL {.importc: "IP_TTL", header: "<netinet/in.h>".}: cint

proc cleanup*(puncher: TcpSyniConnectPuncher) {.async.} =
  while puncher.attempts.len() != 0:
    await puncher.attempts.pop().deleteFirewallRules()

proc initTcpSyniConnectPuncher*(): TcpSyniConnectPuncher =
  TcpSyniConnectPuncher()

proc captureSeqNumbers(attempt: Attempt, cb: PunchProgressCb) {.async.} =
  # FIXME: timeout?
  let iface = getNetworkInterface(attempt.srcIp)
  let captureFd = setupEthernetCapturingSocket(iface)
  var seqNums = newSeq[uint32]()
  while seqNums.len < attempt.dstPorts.len:
    let packet = await captureFd.recv(4000)
    if packet == "":
      break
    let parsed = parseEthernetPacket(packet)
    if parsed.protocol == tcp and
       parsed.ipAddrSrc == attempt.srcIp and
       parsed.tcpPortSrc.int == attempt.srcPort.int and
       parsed.ipAddrDst == attempt.dstIp and
       parsed.tcpFlags == {SYN}:
      for port in attempt.dstPorts:
        if parsed.tcpPortDst.int == port.int:
          seqNums.add(parsed.tcpSeqNumber)
          break
  closeSocket(captureFd)
  await cb(seqNums)

proc captureAndResendAck(attempt: Attempt) {.async.} =
  let iface = getNetworkInterface(attempt.srcIp)
  let captureFd = setupEthernetCapturingSocket(iface)
  let injectFd = setupTcpInjectingSocket()
  block loops:
    while true:
      let packet = await captureFd.recv(4000)
      if packet == "":
        break
      var parsed = parseEthernetPacket(packet)
      if parsed.protocol == tcp and
         parsed.ipAddrSrc == attempt.srcIp and
         parsed.tcpPortSrc.int == attempt.srcPort.int and
         parsed.ipAddrDst == attempt.dstIp and
         parsed.tcpFlags == {ACK}:
        for port in attempt.dstPorts:
          if parsed.tcpPortDst.int == port.int:
            parsed.ipTTL = 64
            echo &"[{parsed.ipAddrSrc}:{parsed.tcpPortSrc.int} -> {parsed.ipAddrDst}:{parsed.tcpPortDst}, SEQ {parsed.tcpSeqNumber}] resending ACK with TTL {parsed.ipTTL}"
            await injectFd.injectTcpPacket(parsed)
            break loops
  closeSocket(captureFd)
  closeSocket(injectFd)

proc doConnect(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress, dstPort: Port,
               future: Future[AsyncSocket]) {.async.} =
  let sock = newAsyncSocket()
  sock.setSockOpt(OptReuseAddr, true)
  sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, 2)
  echo &"doConnect {srcIp}:{srcPort} -> {dstIp}:{dstPort}"
  sock.bindAddr(srcPort, $srcIp)
  try:
    await sock.connect($dstIp, dstPort)
    sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, 64)
    future.complete(sock)
  except OSError as e:
    echo &"connection {srcIP}:{srcPort.int} -> {dstIp}:{dstPort.int} failed: ", e.msg
    sock.close()

proc connect*(puncher: TcpSyniConnectPuncher, srcPort: Port, dstIp: IpAddress,
              dstPorts: seq[Port],
              progressCb: PunchProgressCb): Future[AsyncSocket] {.async.} =
  let localIp = getPrimaryIPAddr(dstIp)
  if puncher.findAttempt(localIp, srcPort, dstIp, dstPorts) != -1:
    raise newException(PunchHoleError, "hole punching for given parameters already active")
  let attempt = Attempt(srcIp: localIp, srcPort: srcPort, dstIp: dstIp,
                        dstPorts: predictPortRange(dstPorts))
  puncher.attempts.add(attempt)
  await attempt.addFirewallRules()
  asyncCheck attempt.captureSeqNumbers(progressCb)
  asyncCheck attempt.captureAndResendAck()
  try:
    let connectFuture = newFuture[AsyncSocket]("connect")
    for dstPort in attempt.dstPorts:
      asyncCheck doConnect(attempt.srcIp, attempt.srcPort, attempt.dstIp,
                           dstPort, connectFuture)
    await connectFuture or sleepAsync(Timeout)
    await attempt.deleteFirewallRules()
    puncher.attempts.del(puncher.attempts.find(attempt))
    if connectFuture.finished():
      result = connectFuture.read()
    else:
      raise newException(PunchHoleError, "timeout")
  except OSError as e:
    raise newException(PunchHoleError, e.msg)