import asyncdispatch, asyncnet, strformat
from net import IpAddress, Port, `$`, `==`, parseIpAddress
from random import randomize, rand
from sequtils import any
import ip_packet
import puncher
import raw_socket
import utils

export PunchHoleError

type
  Attempt = object
    srcIp: IpAddress
    srcPort: Port
    dstIp: IpAddress
    dstPorts: seq[Port]
    seqNums: seq[uint32]
    future: Future[AsyncSocket]

  TcpSyniAcceptPuncher* = Puncher[Attempt]

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

proc initTcpSyniAcceptPuncher*(): TcpSyniAcceptPuncher =
  randomize()
  TcpSyniAcceptPuncher()

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)
    for seqNum in attempt.seqNums:
      let synIn = IpPacket(protocol: tcp, ipAddrSrc: attempt.dstIp,
                           ipAddrDst: attempt.srcIp, ipTTL: 64,
                           tcpPortSrc: dstPort,
                           tcpPortDst: attempt.srcPort,
                           tcpSeqNumber: seqNum, tcpAckNumber: 0,
                           tcpFlags: {SYN}, tcpWindowSize: 1452 * 10)
      echo &"[{synIn.ipAddrSrc}:{synIn.tcpPortSrc} -> {synIn.ipAddrDst}:{synIn.tcpPortDst}, SEQ {synIn.tcpSeqNumber}] injecting incoming SYN"
      await injectFd.injectTcpPacket(synIn)
  closeSocket(injectFd)

proc doAccept(puncher: TcpSyniAcceptPuncher, srcIp: IpAddress,
              srcPort: Port) {.async.} =
  let sock = newAsyncSocket()
  sock.setSockOpt(OptReuseAddr, true)
  sock.bindAddr(srcPort, $(srcIp))
  sock.listen()
  while true:
    let acceptFuture = sock.accept()
    await acceptFuture or sleepAsync(Timeout)
    if acceptFuture.finished():
      let peer = acceptFuture.read()
      let (peerAddr, peerPort) = peer.getPeerAddr()
      let peerIp = parseIpAddress(peerAddr)
      let i = puncher.findAttempt(srcIp, srcPort, peerIp, @[peerPort])
      if i == -1:
        echo "Accepted connection, but no attempt found. Discarding."
        peer.close()
        continue
      else:
        let attempt = puncher.attempts[i]
        attempt.future.complete(peer)
    let attempts = puncher.findAttemptsByLocalAddr(srcIp, srcPort)
    # FIXME: should attempts have timestamps, so we can decide here which ones to delete?
    if attempts.len() <= 1:
      break
  sock.close()

proc accept*(puncher: TcpSyniAcceptPuncher, srcPort: Port, dstIp: IpAddress,
             dstPorts: seq[Port],
             seqNums: seq[uint32]): Future[AsyncSocket] {.async.} =
  let localIp = getPrimaryIPAddr(dstIp)
  let existingAttempts = puncher.findAttemptsByLocalAddr(localIp, srcPort)
  if existingAttempts.len() == 0:
    echo &"accepting connections from {dstIp}:{dstPorts[0].int}"
    asyncCheck puncher.doAccept(localIp, srcPort)
  else:
    for a in existingAttempts:
      if a.dstIp == dstIp and
         a.dstPorts.any(proc (p: Port): bool = p in dstPorts):
        raise newException(PunchHoleError, "hole punching for given parameters already active")
  try:
    let attempt = Attempt(srcIp: localIp, srcPort: srcPort, dstIp: dstIp,
                          dstPorts: predictPortRange(dstPorts),
                          seqNums: seqNums,
                          future: newFuture[AsyncSocket]("accept"))
    puncher.attempts.add(attempt)
    await attempt.addFirewallRules()
    await attempt.injectSynPackets()
    await attempt.future or sleepAsync(Timeout)
    await attempt.deleteFirewallRules()
    puncher.attempts.del(puncher.attempts.find(attempt))
    if attempt.future.finished():
      result = attempt.future.read()
    else:
      raise newException(PunchHoleError, "timeout")
  except OSError as e:
    raise newException(PunchHoleError, e.msg)