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

export puncher

type
  TcpSyniPuncher* = 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]
    seqNums: seq[uint32]

  TcpSyniInitiateAttempt = ref object of Attempt

  TcpSyniRespondAttempt = ref object of Attempt
    seqNums: seq[uint32]

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

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.join(","))

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 connect(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 &"connect {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 injectSynPackets(attempt: TcpSyniRespondAttempt) {.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 initTcpSyniPuncher*(): TcpSyniPuncher =
  randomize()
  TcpSyniPuncher()

method cleanup*(attempt: TcpSyniInitiateAttempt) {.async.} =
  await deleteFirewallRules(attempt)

method cleanup*(attempt: TcpSyniRespondAttempt) {.async.} =
  discard

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

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

method initiate*(puncher: TcpSyniPuncher, attempt: Attempt,
                 progress: PunchProgressCb): Future[AsyncSocket] {.async.} =
  assert(attempt of TcpSyniInitiateAttempt, "unexpected attempt type")
  assert(attempt.acceptFuture.isNone(), "expected attempt without acceptFuture")
  await addFirewallRules(attempt)
  asyncCheck captureSeqNumbers(attempt, progress)
  asyncCheck captureAndResendAck(attempt)
  try:
    let connectFuture = newFuture[AsyncSocket]("connect")
    for dstPort in attempt.dstPorts:
      asyncCheck connect(attempt.srcIp, attempt.srcPort, attempt.dstIp,
                         dstPort, connectFuture)
    await connectFuture or sleepAsync(Timeout)
    await deleteFirewallRules(attempt)
    if connectFuture.finished():
      result = connectFuture.read()
    else:
      raise newException(PunchHoleError, "timeout")
  except OSError as e:
    raise newException(PunchHoleError, e.msg)

method respond*(puncher: TcpSyniPuncher, attempt: Attempt):
               Future[AsyncSocket] {.async.} =
  assert(attempt of TcpSyniRespondAttempt, "unexpected attempt type")
  assert(attempt.acceptFuture.isSome(), "expected attempt with acceptFuture")
  try:
    let acceptFuture = attempt.acceptFuture.get()
    await injectSynPackets(TcpSyniRespondAttempt(attempt))
    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)