from nativesockets import Protocol
from net import IpAddress, Port, `$`, `==`
from random import randomize, rand
import asyncdispatch, asyncnet, strformat
import ip_packet
import message
import options
import port_prediction
import puncher
import raw_socket
import utils

export puncher

type
  TcpNutssPuncher* = ref object of Puncher

  InitiateRequest = object
    srcIp: IpAddress
    srcPort: Port
    probedSrcPorts: seq[Port]
    dstIp: IpAddress
    dstPort: Port
    probedDstPorts: seq[Port]

  RespondRequest = object
    dstIp: IpAddress
    dstPort: Port
    probedDstPorts: seq[Port]
    srcIp: IpAddress
    srcPort: Port
    probedSrcPorts: seq[Port]
    extraArgs: string

proc injectSynPackets(attempt: Attempt) {.async.} =
  let injectFd = setupTcpInjectingSocket()
  for dstPort in attempt.dstPorts:
    let synOut = IpPacket(protocol: tcp, ipAddrSrc: attempt.srcIp,
                          ipAddrDst: attempt.dstIp, ipTTL: 2,
                          tcpPortSrc: attempt.srcPort, tcpPortDst: dstPort,
                          tcpSeqNumber: rand(uint32), tcpAckNumber: 0,
                          tcpFlags: {SYN}, tcpWindowSize: 1452 * 10)
    echo &"[{synOut.ipAddrSrc}:{synOut.tcpPortSrc} -> {synOut.ipAddrDst}:{synOut.tcpPortDst}, SEQ {synOut.tcpSeqNumber}] injecting outgoing SYN"
    await injectFd.injectTcpPacket(synOut)

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

proc initTcpNutssPuncher*(): TcpNutssPuncher =
  randomize()
  TcpNutssPuncher()

method parseInitiateRequest*(puncher: TcpNutssPuncher, args: string): Attempt =
  let parsed = parseMessage[InitiateRequest](args)
  let localIp = getPrimaryIPAddr(parsed.dstIp)
  let predictedDstPorts = predictPortRange(parsed.dstPort,
                                           parsed.probedDstPorts)
  let acceptFuture = newFuture[AsyncSocket]("parseInitiateRequest")
  Attempt(protocol: IPPROTO_TCP, srcIp: localIp, srcPort: parsed.srcPort,
          dstIp: parsed.dstIp, dstPorts: predictedDstPorts,
          acceptFuture: some(acceptFuture))

method parseRespondRequest*(puncher: TcpNutssPuncher, args: string): Attempt =
  let parsed = parseMessage[RespondRequest](args)
  let localIp = getPrimaryIPAddr(parsed.dstIp)
  let predictedDstPorts = predictPortRange(parsed.dstPort,
                                           parsed.probedDstPorts)
  Attempt(protocol: IPPROTO_TCP, srcIp: localIp, srcPort: parsed.srcPort,
          dstIp: parsed.dstIp, dstPorts: predictedDstPorts,
          acceptFuture: none(Future[AsyncSocket]))

method initiate*(puncher: TcpNutssPuncher, attempt: Attempt,
                 progress: PunchProgressCb): Future[AsyncSocket] {.async.} =
  assert(attempt.acceptFuture.isSome(), "expected attempt with acceptFuture")
  try:
    let acceptFuture = attempt.acceptFuture.get()
    await injectSynPackets(attempt)
    await progress("") # FIXME: await means we wait until the rendezvous server replied
    await acceptFuture or sleepAsync(Timeout)
    if acceptFuture.finished():
      result = acceptFuture.read()
    else:
      raise newException(PunchHoleError, "timeout")
  except OSError as e:
    raise newException(PunchHoleError, e.msg)

method respond*(puncher: TcpNutssPuncher, attempt: Attempt):
               Future[AsyncSocket] {.async.} =
  assert(attempt.acceptFuture.isNone(), "expected attempt without acceptFuture")
  try:
    let connectFuture = newFuture[AsyncSocket]("respond")
    for dstPort in attempt.dstPorts:
      asyncCheck connect(attempt.srcIp, attempt.srcPort, attempt.dstIp, dstPort,
                         connectFuture)
    await connectFuture or sleepAsync(Timeout)
    if connectFuture.finished():
      result = connectFuture.read()
    else:
      raise newException(PunchHoleError, "timeout")
  except OSError as e:
    raise newException(PunchHoleError, e.msg)