parse TCP flags; consider TCP flags while capturing

ip_packet.nim

@@ -48,6 +48,14 @@ type
     tcp
     other
 
+  TcpFlag* {.size: sizeof(uint8).} = enum
+    FIN
+    SYN
+    RST
+    PSH
+    ACK
+    URG
+
   IpPacket* = object
     case protocol*: Protocol
     of tcp:
@@ -56,13 +64,13 @@ type
       tcpPortSrc*: Port
       tcpPortDst*: Port
       tcpSeqNumber*: uint32
+      tcpFlags*: set[TcpFlag]
     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)
@@ -76,12 +84,14 @@ proc parseEthernetPacket*(input: string): IpPacket =
                           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))
+                        tcpSeqNumber: ntohl(tcpHeader.th_seq),
+                        tcpFlags: cast[set[TcpFlag]](tcpHeader.th_flags))
     else:
       result = IpPacket(protocol: other)
   else:
@@ -108,7 +118,7 @@ proc serialize*(packet: IpPacket): string =
                            th_seq: htonl(packet.tcpSeqNumber),
                            th_ack: 0,
                            th_off: 5,
-                           th_flags: TH_SYN,
+                           th_flags: cast[uint8](packet.tcpFlags),
                            th_win: 1452 * 10,
                            th_sum: 0,
                            th_urp: 0)

tcp_syni.nim

@@ -59,6 +59,7 @@ proc delFirewallRule(srcIp: IpAddress, srcPort: Port,
 
 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:
@@ -70,7 +71,9 @@ proc captureSeqNumbers(puncher: TcpSyniPuncher, rawFd: AsyncFD,
     if parsed.protocol == tcp and
        parsed.tcpIpSrc == puncher.srcIp and
        parsed.tcpPortSrc.int == puncher.srcPort.int and
-       parsed.tcpIpDst == puncher.dstIp:
+       parsed.tcpIpDst == puncher.dstIp and
+       parsed.tcpFlags.contains(SYN) and
+       not parsed.tcpFlags.contains(ACK):
       for i, port in puncher.dstPorts.pairs:
         if parsed.tcpPortDst.int == port.int:
           seqNums.add(parsed.tcpSeqNumber)