import
  asyncdispatch,
  asyncutils,
  network_interface,
  packet_info,
  strformat

from nativesockets import SOCK_RAW, bindAddr, htons
from posix import ioctl, setsockopt, SockAddr, SockLen, SocketHandle

export NetworkInterfaceError

type
  Sockaddr_ll {.importc: "struct sockaddr_ll", pure, final,
                header: "<linux/if_packet.h>".} = object
    sll_family: cushort        # Always AF_PACKET
    sll_protocol: cushort      # Physical-layer protocol
    sll_ifindex: cint          # Interface number
    sll_hatype: cushort        # ARP hardware type
    sll_pkttype: cuchar        # Packet type
    sll_halen: cuchar          # Length of address
    sll_addr: array[8, cuchar] # Physical-layer address

  Packet_mreq {.importc: "struct packet_mreq", pure, final,
                header: "<linux/if_packet.h>".} = object
    mr_ifindex: cint
    mr_type: cushort
    mr_alen: cushort
    mr_address: array[8, cuchar]

  PunchHoleError* = object of ValueError

var
  AF_PACKET {.importc: "AF_PACKET", header: "<sys/socket.h>".}: cushort
  SOL_PACKET {.importc: "SOL_PACKET", header: "<sys/socket.h>".}: cushort
  ETH_P_ALL {.importc: "ETH_P_ALL", header: "<linux/if_ether.h>".}: cushort
  PACKET_ADD_MEMBERSHIP {.importc: "PACKET_ADD_MEMBERSHIP", header: "<linux/if_packet.h>".}: cushort
  PACKET_MR_PROMISC {.importc: "PACKET_MR_PROMISC", header: "<linux/if_packet.h>".}: cushort

proc setupRawSocket(iface: NetworkInterface): AsyncFD =
  # Create a raw packet socket. For accessing outgoing packets we need to use
  # ETH_P_ALL which is needed in network byte order, see packet(7) man page.
  result = createAsyncNativeSocket(AF_PACKET.cint,
                                   SOCK_RAW.cint,
                                   htons(ETH_P_ALL).cint)

  # Limit capturing of packets to the desired network interface, see
  # netdevice(7) man page
  echo "interface: ", iface.name, ", index: ", iface.index
  var sa = Sockaddr_ll(sll_family: AF_PACKET,
                       sll_protocol: htons(ETH_P_ALL),
                       sll_ifindex: iface.index)
  if bindAddr(result.SocketHandle,
              cast[ptr SockAddr](addr sa),
              sizeof(Sockaddr_ll).SockLen) != 0:
    raise newException(PunchHoleError, "cannot bind to interface")
  
  # Enable promiscuous mode, see netdevice(7) man page
  var req = Packet_mreq(mr_ifindex: iface.index, mr_type: PACKET_MR_PROMISC)
  if setsockopt(result.SocketHandle,
                SOL_PACKET.cint,
                PACKET_ADD_MEMBERSHIP.cint,
                addr req,
                sizeof(req).SockLen) != 0:
    raise newException(PunchHoleError, "cannot enable promiscuous mode")

proc captureSequenceNumber(rawFd: AsyncFD,
                           clientAddress: string,
                           clientPort: Port,
                           serverAddress: string,
                           serverPort: Port): Future[uint32] {.async.} =
  while true:
    let packet = await rawFd.recv(4000)
    if packet == "":
      break
    echo "packet len: ", packet.len
    let packetInfo = fromString(packet)
    if packetInfo.protocol == tcp and
       packetInfo.tcpIpSrc == clientAddress and
       packetInfo.tcpPortSrc.int == clientPort.int and
       packetInfo.tcpIpDst == serverAddress and
       packetInfo.tcpPortDst.int == serverPort.int:
      return packetInfo.tcpSeqNumber

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

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

proc punchHoleAsClient*(clientAddress: string,
                        clientPort: Port,
                        serverAddress: string,
                        serverPort: Port) {.async.} =
  let iface = fromIpAddress(clientAddress)
  let rawFd = setupRawSocket(iface)
  let seqNumber = await captureSequenceNumber(rawFd,
                                              clientAddress,
                                              clientPort,
                                              serverAddress,
                                              serverPort)
  echo "captured sequence number: ", seqNumber
  await addFirewallRule(clientAddress, clientPort, serverAddress, serverPort)
  await deleteFirewallRule(clientAddress, clientPort, serverAddress, serverPort)
  closeSocket(rawFd)