import asyncdispatch, asyncnet, strformat
from net import IpAddress, Port, `$`, `==`, getPrimaryIPAddr, toSockAddr
from nativesockets import SockAddr, Sockaddr_storage, SockLen, setSockOptInt
import asyncutils
import ip_packet
import network_interface
import raw_socket

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

type
  TcpSyniPuncher* = object
    srcIp: IpAddress
    srcPort: Port
    dstIp: IpAddress
    dstPorts: array[10, Port]
    seqNums: seq[uint32]

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

  PunchHoleError* = object of ValueError

proc addFirewallRule(srcIp: IpAddress, srcPort: Port,
                     dstIp: IpAddress, dstPort: Port) {.async.} =
  let firewall_cmd = fmt"""iptables -A INPUT \
  -d {srcIp} \
  -p icmp \
  --icmp-type time-exceeded \
  -m conntrack \
  --ctstate RELATED \
  --ctproto tcp \
  --ctorigsrc {srcIp} \
  --ctorigsrcport {srcPort.int} \
  --ctorigdst {dstIp} \
  --ctorigdstport {dstPort.int} \
  -j DROP"""
  let exitcode = await asyncExecCmd(firewall_cmd)
  if exitcode != 0:
    raise newException(PunchHoleError, "cannot add firewall rule")

proc delFirewallRule(srcIp: IpAddress, srcPort: Port,
                     dstIp: IpAddress, dstPort: Port) {.async.} =
  let firewall_cmd = fmt"""iptables -D INPUT \
  -d {srcIp} \
  -p icmp \
  --icmp-type time-exceeded \
  -m conntrack \
  --ctstate RELATED \
  --ctproto tcp \
  --ctorigsrc {srcIp} \
  --ctorigsrcport {srcPort.int} \
  --ctorigdst {dstIp} \
  --ctorigdstport {dstPort.int} \
  -j DROP"""
  let exitcode = await asyncExecCmd(firewall_cmd)
  if exitcode != 0:
    raise newException(PunchHoleError, "cannot delete firewall rule")

proc captureSeqNumbers(puncher: TcpSyniPuncher, rawFd: AsyncFD,
                       cb: PunchProgressCb) {.async.} =
  # FIXME: every sequence number is captured twice (RST too?)
  # FIXME: timeout?
  var seqNums = newSeq[uint32]()
  while seqNums.len < puncher.dstPorts.len:
    let packet = await rawFd.recv(4000)
    if packet == "":
      break
    echo "packet len: ", packet.len
    let parsed = parseEthernetPacket(packet)
    if parsed.protocol == tcp and
       parsed.tcpIpSrc == puncher.srcIp and
       parsed.tcpPortSrc.int == puncher.srcPort.int and
       parsed.tcpIpDst == puncher.dstIp and
       parsed.tcpFlags == {SYN}:
      for i, port in puncher.dstPorts.pairs:
        if parsed.tcpPortDst.int == port.int:
          seqNums.add(parsed.tcpSeqNumber)
          break
  await cb(seqNums)

proc injectSyns(rawFd: AsyncFD, srcIp: IpAddress, srcPort: Port,
                dstIp: IpAddress, dstPort: Port,
                seqNums: seq[uint32]) {.async.} =
  for seqNum in seqNums:
    let ipPacket = IpPacket(protocol: tcp,
                            tcpIpSrc: srcIp,
                            tcpIpDst: dstIp,
                            tcpPortSrc: srcPort,
                            tcpPortDst: dstPort,
                            tcpSeqNumber: seqNum,
                            tcpFlags: {SYN})
    try:
      let packet = serialize(ipPacket)
      var sockaddr: Sockaddr_storage
      var sockaddrLen: SockLen
      toSockAddr(dstIp, dstPort, sockaddr, sockaddrLen)
      await rawFd.sendTo(packet.cstring, packet.len,
                         cast[ptr SockAddr](addr sockaddr), sockaddrLen)
      echo &"injected {srcIP}:{srcPort.int} -> {dstIp}:{dstPort.int} (seq {seqNum})"
    except OSError as e:
      echo "cannot inject {srcIp}:{srcPort.int} -> {dstIp}:{dstPort.int} (seq {seqNum}): ", e.msg

proc initPuncher*(srcPort: Port, dstIp: IpAddress, dstPorts: array[3, Port],
                  seqNums: seq[uint32] = @[]): Future[TcpSyniPuncher] {.async.} =
  let localIp = getPrimaryIPAddr(dstIp)
  # TODO: do real port prediction
  var predictedDstPorts: array[10, Port]
  let basePort = min(dstPorts[1].uint16, uint16.high - 9)
  for i in 0.uint16 .. 9.uint16:
    predictedDstPorts[i] = Port(basePort + i)
  result = TcpSyniPuncher(srcIp: localIp, srcPort: srcPort, dstIp: dstIp,
                          dstPorts: predictedDstPorts, seqNums: seqNums)
  for dstPort in result.dstPorts:
    await addFirewallRule(result.srcIp, result.srcPort, result.dstIp, dstPort)

proc cleanup*(puncher: TcpSyniPuncher) {.async.} =
  for dstPort in puncher.dstPorts:
    await delFirewallRule(puncher.srcIp, puncher.srcPort, puncher.dstIp, dstPort)

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)
  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
    discard

proc doAccept(srcIp: IpAddress, srcPort: Port, dstIp: IpAddress, dstPort: Port,
              seqNums: seq[uint32], future: Future[AsyncSocket]) {.async.} =
  try:
    let sock = newAsyncSocket()
    sock.setSockOpt(OptReuseAddr, true)
    sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, 2)
    sock.bindAddr(srcPort, $srcIp)
    await sock.connect($dstIp, dstPort)
    echo "connected during accept phase"
    sock.close()
  except OSError:
    discard
  try:
    # FIXME: timeout
    let rawFd = setupTcpInjectingSocket()
    asyncCheck injectSyns(rawFd, dstIp, dstPort, srcIp, srcPort, seqNums)
    let sock = newAsyncSocket()
    sock.setSockOpt(OptReuseAddr, true)
    sock.setSockOpt(OptReusePort, true)
    sock.bindAddr(srcPort, $srcIp)
    sock.listen()
    let connectedSock = await sock.accept()
    future.complete(connectedSock)
  except OSError as e:
    echo &"accepting connection {dstIP}:{dstPort.int} -> {srcIp}:{srcPort.int} failed: ", e.msg
    discard

proc connect*(puncher: TcpSyniPuncher,
              progressCb: PunchProgressCb): Future[AsyncSocket] =
  result = newFuture[AsyncSocket]("tcp_syni.connect")
  let iface = fromIpAddress($puncher.srcIp)
  let rawFd = setupEthernetCapturingSocket(iface)
  asyncCheck puncher.captureSeqNumbers(rawFd, progressCb)
  for dstPort in puncher.dstPorts:
    asyncCheck doConnect(puncher.srcIp, puncher.srcPort, puncher.dstIp,
                         dstPort, result)

proc accept*(puncher: TcpSyniPuncher): Future[AsyncSocket] =
  result = newFuture[AsyncSocket]("tcp_syni.accept")
  for dstPort in puncher.dstPorts:
    asyncCheck doAccept(puncher.srcIp, puncher.srcPort, puncher.dstIp, dstPort,
                        puncher.seqNums, result)