from nativesockets import ntohs, ntohl, htons, htonl
from net import IpAddress, IpAddressFamily, Port, `$`
from posix import InAddr, inet_ntoa

type
  Ether_header {.importc: "struct ether_header", pure, final,
                 header: "<netinet/if_ether.h>".} = object
    ether_dhost: array[6, uint8]
    ether_shost: array[6, uint8]
    ether_type: cushort

  Ip {.importc: "struct ip", pure, final,
       header: "<netinet/ip.h>".} = object
    when cpuEndian == littleEndian:
      ip_hl {.bitsize:4.}: uint   # header length
      ip_v {.bitsize:4.}: uint    # version
    else:
      ip_v {.bitsize:4.}: uint    # version
      ip_hl {.bitsize:4.}: uint   # header length
    ip_tos: uint8                 # type of service
    ip_len: cushort               # total length
    ip_id: cushort                # identification
    ip_off: cushort               # fragment offset field
    ip_ttl: uint8                 # time to live
    ip_p: cuchar                  # protocol
    ip_sum: cushort               # checksum
    ip_src: InAddr                # source address
    ip_dst: InAddr                # destination address

  Tcphdr {.importc: "struct tcphdr", pure, final,
        header: "<netinet/tcp.h>".} = object
    th_sport: uint16               # source port
    th_dport: uint16               # destination port
    th_seq: uint32                 # sequence number
    th_ack: uint32                 # ackkowledgment number
    when cpuEndian == littleEndian:
      th_x2 {.bitsize:4.}: uint8   # (unused)
      th_off {.bitsize:4.}: uint8  # data offset
    else:
      th_off {.bitsize:4.}: uint8  # (unused)
      th_x2 {.bitsize:4.}: uint8   # data offset
    th_flags: uint8                # flags
    th_win: uint16                 # window
    th_sum: uint16                 # checksum
    th_urp: uint16                 # urgent pointer

  Protocol* = enum
    tcp
    other

  IpPacket* = object
    case protocol*: Protocol
    of tcp:
      tcpIpSrc*: IpAddress
      tcpIpDst*: IpAddress
      tcpPortSrc*: Port
      tcpPortDst*: Port
      tcpSeqNumber*: uint32
    else:
      discard

var
  ETHERTYPE_IP {.importc: "ETHERTYPE_IP", header: "<netinet/if_ether.h>".}: cushort
  IPPROTO_TCP {.importc: "IPPROTO_TCP", header: "<netinet/in.h>".}: cint
  TH_SYN {.importc: "TH_SYN", header: "<netinet/tcp.h>".}: uint8

proc parseEthernetPacket*(input: string): IpPacket =
  let etherHeader = cast[ptr Ether_header](input.cstring)
  if ntohs(etherHeader.ether_type) == ETHERTYPE_IP:
    let ipHeader = cast[ptr Ip](cast[int](input.cstring) + sizeof(Ether_header))
    let ipSrcScalar = ntohl(ipHeader.ip_src.s_addr)
    let ipDstScalar = ntohl(ipHeader.ip_dst.s_addr)
    let ipSrc = IpAddress(family: Ipv4,
                          address_v4: cast[array[4, uint8]](ipSrcScalar))
    let ipDst = IpAddress(family: Ipv4,
                          address_v4: cast[array[4, uint8]](ipDstScalar))
    if ipHeader.ip_p.int == IPPROTO_TCP:
      let tcpHeader = cast[ptr Tcphdr](cast[int](ipHeader) + ipHeader.ip_hl.int * 4)
      result = IpPacket(protocol: tcp,
                        tcpIpSrc: ipSrc,
                        tcpIpDst: ipDst,
                        tcpPortSrc: Port(ntohs(tcpHeader.th_sport)),
                        tcpPortDst: Port(ntohs(tcpHeader.th_dport)),
                        tcpSeqNumber: ntohl(tcpHeader.th_seq))
    else:
      result = IpPacket(protocol: other)
  else:
    result = IpPacket(protocol: other)

proc serialize*(packet: IpPacket): string =
  case packet.protocol
  of tcp:
    let srcIp = InAddr(s_addr: htonl(cast[uint32](packet.tcpIpSrc.address_v4)))
    let dstIp = InAddr(s_addr: htonl(cast[uint32](packet.tcpIpDst.address_v4)))
    var ipHeader = Ip(ip_hl: 5,
                      ip_v: 4,
                      ip_tos: 0,
                      ip_len: htons(40),
                      ip_id: htons(54321), # FIXME: random number
                      ip_off: 0,
                      ip_ttl: 64,
                      ip_p: 6.cuchar,
                      ip_sum: htons(0), # done by kernel
                      ip_src: srcIp,
                      ip_dst: dstIp)
    var tcpHeader = Tcphdr(th_sport: htons(packet.tcpPortSrc.uint16),
                           th_dport: htons(packet.tcpPortDst.uint16),
                           th_seq: htonl(packet.tcpSeqNumber),
                           th_ack: 0,
                           th_off: 5,
                           th_flags: TH_SYN,
                           th_win: 1452 * 10,
                           th_sum: 0,
                           th_urp: 0)
    result = newString(sizeof(Ip) + sizeof(Tcphdr))
    copyMem(addr result[0], addr ipHeader, sizeof(Ip))
    copyMem(addr result[sizeof(Ip)], addr tcpHeader, sizeof(Tcphdr))
  else:
    raise newException(ValueError, "protocol not supported")