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

type
  Attempt* = ref object of RootObj
    ## A hole punching attempt.
    ##
    ## It is created by parsing the arguments of a request in either
    ## ``parseInitiateRequest`` or ``parseRespondRequest`` and must not be
    ## modified afterwards. By inclduing ``some(acceptFuture)``, the puncher
    ## tells the caller to accept connections at the local IP address and
    ## port ``srcPort`` before calling ``initiate``. The desired transport
    ## protocol can be obtained by calling ``getProcotol``. The puncher expects
    ## the caller to complete the future when a connections from
    ## ``dstIp``:``dstPort`` has been accepted.
    srcIp*: IpAddress
    srcPort*: Port
    dstIp*: IpAddress
    dstPorts*: seq[Port]
    acceptFuture*: Option[Future[AsyncSocket]]

  Puncher* = ref object of RootObj
    ## A hole puncher.

  PunchHoleError* = object of ValueError
    attempt: Attempt
    ## An exception indicating that the contained ``attempt`` has failed.

  PunchProgressCb* = proc(extraArgs: string) {.async.}
    ## A callback to allow a puncher report progress.
    ##
    ## When called a status message of type ``progress``, including the given
    ## ``extraArgs`` has to be sent to the application.

const Timeout* = 3000

proc `==`*(a, b: Attempt): bool =
  ## ``==`` for hole punching attempts.
  ##
  ## Two hole punching attempts are considered equal if their ``srcIp``,
  ## ``srcPort`` and ``dstIp`` are equal and their ``dstPorts`` overlap.
  a.srcIp == b.srcIp and a.srcPort == b.srcPort and a.dstIp == b.dstIp and
  a.dstPorts.any(proc (p: Port): bool = p in b.dstPorts)

method cleanup*(attempt: Attempt): Future[void] {.base, async.} =
  ## Cleans up when an attempt finished (either successful or not).
  ##
  ## Does nothing. Override for custom attempt types.
  discard

method getProtocol*(puncher: Puncher): Protocol {.base.} =
  ## Returns the transport protocol the puncher employs. 
  raise newException(CatchableError, "Method without implementation override")

method parseInitiateRequest*(puncher: Puncher, args: string): Attempt {.base.} =
  ## Creates a new hole punching attempt by parsing arguments of an ``initiate``
  ## request.
  ##
  ## ``args`` has to be the arguments after the ``TECHNIQUE`` field, i.e.
  ## ``IP_FROM|PORTS_FROM|IP_TO|PORTS_TO``. My throw a ``ValueError`` if
  ## ``args`` is invalid.
  raise newException(CatchableError, "Method without implementation override")

method parseRespondRequest*(puncher: Puncher, args: string): Attempt {.base.} =
  ## Creates a new hole punching attempt by parsing arguments of a ``respond``
  ## request.
  ##
  ## ``args`` has to be the arguments after the ``TECHNIQUE`` field, i.e.
  ## ``IP_FROM|PORTS_FROM|IP_TO|PORTS_TO|EXTRA_ARGS``. May throw a
  ## ``ValueError`` if ``args`` is invalid.
  raise newException(CatchableError, "Method without implementation override")

method initiate*(puncher: Puncher, attempt: Attempt, progress: PunchProgressCb):
                Future[AsyncSocket] {.base, async.} =
  ## Initiate a hole punching attempt.
  ##
  ## ``attempt`` has to be obtained by calling ``parseInitiateRequest``.
  ## ``progress`` will be called when the attempt has been initiated and a
  ## status message of type ``progress`` has to be sent to the application. The
  ## returned future contains a connected socket and will fail with a
  ## ``PunchHoleError`` if the hole punching fails.
  block: # workaround for https://github.com/nim-lang/Nim/issues/12530
    raise newException(CatchableError, "Method without implementation override")

method respond*(puncher: Puncher, attempt: Attempt):
               Future[AsyncSocket] {.base, async.} =
  ## Respond to a hole punching attempt initiated by another peer.
  ##
  ## ``attempt has to be obtained by calling ``parseRespondRequest``. The
  ## returned future contains a connected socket and will fail with as
  ## ``PunchHoleError`` if the hole punching fails.
  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 =
  result = &"""-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 = &"iptables -I {chain} {rule}"
  discard await asyncExecCmd(firewall_cmd)

proc iptablesDelete(chain: string, rule: string) {.async.} =
  let firewall_cmd = &"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