import asyncdispatch, asyncnet, strformat
from net import IpAddress, Port, `$`, `==`
from sequtils import any
import asyncutils

type
  Attempt* = ref object of RootObj
    srcIp*: IpAddress
    srcPort*: Port
    dstIp*: IpAddress
    dstPorts*: seq[Port]

  Puncher* = ref object of RootObj
    attempts*: seq[Attempt]

  Initiator* = ref object of Puncher

  Responder* = ref object of Puncher

  PunchHoleError* = object of ValueError

  PunchProgressCb* = proc(extraArgs: string) {.async.}

const Timeout* = 3000

proc findAttempt*(puncher: Puncher, srcIp: IpAddress, srcPort: Port,
                  dstIp: IpAddress, dstPorts: seq[Port]): int =
  for (index, attempt) in puncher.attempts.pairs():
    if attempt.srcIp == srcIp and attempt.srcPort == srcPort and
       attempt.dstIp == dstIp and
       attempt.dstPorts.any(proc (p: Port): bool = p in dstPorts):
      return index
  return -1

proc findAttemptsByLocalAddr*(puncher: Puncher, address: IpAddress,
                              port: Port): seq[Attempt] =
  for attempt in puncher.attempts:
    if attempt.srcIp == address and attempt.srcPort == port:
      result.add(attempt)

method cleanup*(puncher: Puncher): Future[void] {.base, async.} =
  block: # workaround for https://github.com/nim-lang/Nim/issues/12530
    raise newException(CatchableError, "Method without implementation override")

method initiate*(puncher: Initiator, args: string, progress: PunchProgressCb):
                Future[AsyncSocket] {.base, async.} =
  block: # workaround for https://github.com/nim-lang/Nim/issues/12530
    raise newException(CatchableError, "Method without implementation override")

method respond*(puncher: Responder, args: string):
               Future[AsyncSocket] {.base, async.} =
  block: # workaround for https://github.com/nim-lang/Nim/issues/12530
    raise newException(CatchableError, "Method without implementation override")

proc makeFirewallRule(srcIp: IpAddress, srcPort: Port,
                      dstIp: IpAddress, dstPort: Port): string =
  # FIXME: use & instead of fmt?
  result = fmt"""-w \
  -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"""

proc iptablesInsert(chain: string, rule: string) {.async.} =
  let firewall_cmd = fmt"iptables -I {chain} {rule}"
  discard await asyncExecCmd(firewall_cmd)

proc iptablesDelete(chain: string, rule: string) {.async.} =
  let firewall_cmd = fmt"iptables -D {chain} {rule}"
  discard await asyncExecCmd(firewall_cmd)

proc addFirewallRules*(attempt: Attempt) {.async.} =
  for dstPort in attempt.dstPorts:
    let rule = makeFirewallRule(attempt.srcIp, attempt.srcPort,
                                attempt.dstIp, dstPort)
    try:
      await iptablesInsert("INPUT", rule)
    except OSError as e:
      echo "cannot add firewall rule: ", e.msg
      raise newException(PunchHoleError, e.msg)

proc deleteFirewallRules*(attempt: Attempt) {.async.} =
  for dstPort in attempt.dstPorts:
    let rule = makeFirewallRule(attempt.srcIp, attempt.srcPort,
                                attempt.dstIp, dstPort)
    try:
      await iptablesDelete("INPUT", rule)
    except OSError:
      # At least we tried
      discard