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

export Puncher, Responder, PunchHoleError, cleanup, respond

type
  TSRAttempt = ref object of Attempt
    seqNums: seq[uint32]
    future: Future[AsyncSocket]

  TcpSyniResponder* = ref object of Responder

  Request = object
    dstIp: IpAddress
    dstPorts: seq[Port]
    srcIp: IpAddress
    srcPorts: seq[Port]
    seqNums: seq[uint32]

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

proc initTcpSyniResponder*(): TcpSyniResponder =
  randomize()
  TcpSyniResponder()

proc injectSynPackets(attempt: TSRAttempt) {.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 accept(puncher: TcpSyniResponder, 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 = TSRAttempt(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()

method respond*(puncher: TcpSyniResponder, args: string):
               Future[AsyncSocket] {.async.} =
  let req = parseMessage[Request](args)
  let localIp = getPrimaryIPAddr(req.dstIp)
  let existingAttempts = puncher.findAttemptsByLocalAddr(localIp, req.srcPorts[0])
  if existingAttempts.len() == 0:
    echo &"accepting connections from {req.dstIp}:{req.dstPorts[0].int}"
    asyncCheck puncher.accept(localIp, req.srcPorts[0])
  else:
    for a in existingAttempts:
      if a.dstIp == req.dstIp and
         a.dstPorts.any(proc (p: Port): bool = p in req.dstPorts):
        raise newException(PunchHoleError, "hole punching for given parameters already active")
  try:
    let attempt = TSRAttempt(srcIp: localIp, srcPort: req.srcPorts[0],
                             dstIp: req.dstIp,
                             dstPorts: predictPortRange(req.dstPorts),
                             seqNums: req.seqNums,
                             future: newFuture[AsyncSocket]("respond"))
    puncher.attempts.add(attempt)
    await attempt.addFirewallRules() # FIXME: needed?
    await attempt.injectSynPackets()
    await attempt.future or sleepAsync(Timeout)
    await attempt.deleteFirewallRules() # FIXME: needed?
    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)