Merge pull request #4 from CAIDA/rsdos-attacks

Add code to support reading RS-DOS attack avro data

setup.py

@@ -29,7 +29,7 @@ def no_cythonize(extensions, **_ignore):
     Extension('pyavro_stardust.baseavro', ['src/pyavro_stardust/baseavro.pyx'], language="c++"),
     Extension('pyavro_stardust.flowtuple3', ['src/pyavro_stardust/flowtuple3.pyx'], language="c++"),
     Extension('pyavro_stardust.flowtuple4', ['src/pyavro_stardust/flowtuple4.pyx'], language="c++"),
-    #Extension('pyavro_stardust.rsdos', ['src/pyavro_stardust/rsdos.pyx'])
+    Extension('pyavro_stardust.rsdos', ['src/pyavro_stardust/rsdos.pyx'], language="c++")
 ]
 
 CYTHONIZE = bool(int(os.getenv("CYTHONIZE", 0))) and cythonize is not None

src/pyavro_stardust/baseavro.pxd

@@ -1,4 +1,7 @@
 from libcpp.vector cimport vector
+from cpython cimport array
+import array
+
 
 cdef struct parsedString:
     unsigned int toskip
@@ -14,7 +17,8 @@ cdef struct parsedNumericArrayBlock:
 cdef (unsigned int, long) read_long(const unsigned char[:] buf,
         const unsigned int maxlen)
 cdef parsedString read_string(const unsigned char[:] buf,
-        const unsigned int maxlen)
+        const unsigned int maxlen, int addNullTerm=*)
+
 cdef parsedNumericArrayBlock read_numeric_array(const unsigned char[:] buf,
         const unsigned int maxlen)
 

src/pyavro_stardust/baseavro.pyx

@@ -31,7 +31,8 @@ cdef (unsigned int, long) read_long(const unsigned char[:] buf,
     return (longlen + 1, (n >> 1) ^ -(n & 1))
 
 cdef parsedString read_string(const unsigned char[:] buf,
-        const unsigned int maxlen):
+        const unsigned int maxlen, int addNullTerm=True):
+
     cdef unsigned int skip
     cdef long strlen
     cdef parsedString s
@@ -45,7 +46,15 @@ cdef parsedString read_string(const unsigned char[:] buf,
 
     s.toskip = skip
     s.strlen = strlen
-    s.start = <unsigned char *>&(buf[skip])
+
+    if addNullTerm:
+        s.start = <unsigned char *>PyMem_Malloc(strlen + 1)
+        memcpy(s.start, &(buf[skip]), strlen)
+        s.start[strlen] = b'\x00'
+    else:
+        s.start = <unsigned char *>PyMem_Malloc(strlen)
+        memcpy(s.start, &(buf[skip]), strlen)
+
     return s
 
 cdef parsedNumericArrayBlock read_numeric_array(const unsigned char[:] buf,
@@ -184,17 +193,14 @@ cdef class AvroRecord:
         if attrind < 0 or <unsigned int>attrind >= self.stringcount:
             return 0
 
-        astr = read_string(buf, maxlen)
+
+        astr = read_string(buf, maxlen, True)
 
         if astr.toskip == 0:
             return 0
 
         self.sizeinbuf += astr.toskip + astr.strlen
-        self.attributes_s[attrind] = <char *>PyMem_Malloc(sizeof(char) * astr.strlen + 1)
-
-        memcpy(self.attributes_s[attrind], astr.start, astr.strlen)
-        self.attributes_s[attrind][astr.strlen] = b'\x00'
-
+        self.attributes_s[attrind] = <char *>astr.start
         return astr.toskip + astr.strlen
 
     @cython.wraparound(False)
@@ -325,10 +331,31 @@ cdef class AvroReader:
     def start(self):
         if self.fh is not None:
             return
-        try:
-            self.fh = wandio.open(self.filepath, 'rb')
-        except:
-            raise
+
+        # Try 'rb' mode if pywandio supports it, else fallback to 'r'
+        # and hope that we're not reading a file off local disk (that
+        # pywandio will try to "decode" into utf-8)
+        #
+        # Future versions of pywandio may allow us to override the
+        # decoding method, in which case we can rework this code to be
+        # less clunky.
+        mode = 'rb'
+        saved = None
+        while mode != 'fail':
+            try:
+                self.fh = wandio.open(self.filepath, mode=mode)
+            except ValueError as e:
+                if mode == 'rb':
+                    mode = 'r'
+                else:
+                    mode = 'fail'
+                    raise
+            except Exception:
+                raise
+
+            if self.fh is not None:
+                break
+
 
         if self.syncmarker is None:
             self._readAvroFileHeader()

src/pyavro_stardust/rsdos.pxd

@@ -0,0 +1,38 @@
+cimport cython
+from pyavro_stardust.baseavro cimport AvroRecord, AvroReader, parsedString
+
+cpdef enum RsdosAttribute:
+    ATTR_RSDOS_TIMESTAMP = 0
+    ATTR_RSDOS_PACKET_LEN = 1
+    ATTR_RSDOS_TARGET_IP = 2
+    ATTR_RSDOS_TARGET_PROTOCOL = 3
+    ATTR_RSDOS_ATTACKER_IP_CNT = 4
+    ATTR_RSDOS_ATTACK_PORT_CNT = 5
+    ATTR_RSDOS_TARGET_PORT_CNT = 6
+    ATTR_RSDOS_PACKET_CNT = 7
+    ATTR_RSDOS_ICMP_MISMATCHES = 8
+    ATTR_RSDOS_BYTE_CNT = 9
+    ATTR_RSDOS_MAX_PPM_INTERVAL = 10
+    ATTR_RSDOS_START_TIME_SEC = 11
+    ATTR_RSDOS_START_TIME_USEC = 12
+    ATTR_RSDOS_LATEST_TIME_SEC = 13
+    ATTR_RSDOS_LATEST_TIME_USEC = 14
+    ATTR_RSDOS_LAST_ATTRIBUTE = 15
+
+cdef class AvroRsdos(AvroRecord):
+
+    cdef unsigned char *packetcontent
+    cdef public unsigned int pktcontentlen
+
+    cpdef dict asDict(self)
+    cpdef void resetRecord(self)
+    cpdef bytes getRsdosPacketString(self)
+    cpdef int setRsdosPacketString(self, const unsigned char[:] buf,
+            const unsigned int maxlen)
+
+cdef class AvroRsdosReader(AvroReader):
+    cdef int _parseNextRecord(self, const unsigned char[:] buf,
+            const unsigned int maxlen)
+
+
+# vim: set sw=4 tabstop=4 softtabstop=4 expandtab :

src/pyavro_stardust/rsdos.pyx

@@ -0,0 +1,111 @@
+
+# cython: language_level=3
+cimport cython
+from cpython.mem cimport PyMem_Free
+from pyavro_stardust.baseavro cimport AvroRecord, read_long, read_string, \
+        AvroReader, parsedString
+
+@cython.final
+cdef class AvroRsdos(AvroRecord):
+    def __init__(self):
+        super().__init__(ATTR_RSDOS_LAST_ATTRIBUTE, 0, 0)
+        self.pktcontentlen = 0
+        self.packetcontent = NULL
+
+    def __str__(self):
+        return "%u %u.%06u %u.%06u %08x %u %u %u %u %u %u %u %u %u" % \
+                (self.attributes_l[<int>ATTR_RSDOS_TIMESTAMP], \
+                 self.attributes_l[<int>ATTR_RSDOS_START_TIME_SEC],
+                 self.attributes_l[<int>ATTR_RSDOS_START_TIME_USEC],
+                 self.attributes_l[<int>ATTR_RSDOS_LATEST_TIME_SEC],
+                 self.attributes_l[<int>ATTR_RSDOS_LATEST_TIME_USEC],
+                 self.attributes_l[<int>ATTR_RSDOS_TARGET_IP],
+                 self.attributes_l[<int>ATTR_RSDOS_TARGET_PROTOCOL],
+                 self.attributes_l[<int>ATTR_RSDOS_PACKET_LEN],
+                 self.attributes_l[<int>ATTR_RSDOS_ATTACKER_IP_CNT],
+                 self.attributes_l[<int>ATTR_RSDOS_ATTACK_PORT_CNT],
+                 self.attributes_l[<int>ATTR_RSDOS_TARGET_PORT_CNT],
+                 self.attributes_l[<int>ATTR_RSDOS_PACKET_CNT],
+                 self.attributes_l[<int>ATTR_RSDOS_BYTE_CNT],
+                 self.attributes_l[<int>ATTR_RSDOS_MAX_PPM_INTERVAL],
+                 self.pktcontentlen)
+
+    cpdef dict asDict(self):
+        if self.pktcontentlen == 0:
+            initpkt = None
+        else:
+            initpkt = self.getRsdosPacketString()
+
+        return {
+            "timestamp": self.attributes_l[<int>ATTR_RSDOS_TIMESTAMP],
+            "start_time_sec": self.attributes_l[<int>ATTR_RSDOS_START_TIME_SEC],
+            "start_time_usec": self.attributes_l[<int>ATTR_RSDOS_START_TIME_USEC],
+            "latest_time_sec": self.attributes_l[<int>ATTR_RSDOS_LATEST_TIME_SEC],
+            "latest_time_usec": self.attributes_l[<int>ATTR_RSDOS_LATEST_TIME_USEC],
+            "target_ip": self.attributes_l[<int>ATTR_RSDOS_TARGET_IP],
+            "target_protocol": self.attributes_l[<int>ATTR_RSDOS_TARGET_PROTOCOL],
+            "packet_len": self.attributes_l[<int>ATTR_RSDOS_PACKET_LEN],
+            "attacker_count": self.attributes_l[<int>ATTR_RSDOS_ATTACKER_IP_CNT],
+            "attack_port_count": self.attributes_l[<int>ATTR_RSDOS_ATTACK_PORT_CNT],
+            "target_port_count": self.attributes_l[<int>ATTR_RSDOS_TARGET_PORT_CNT],
+            "packet_count": self.attributes_l[<int>ATTR_RSDOS_PACKET_CNT],
+            "byte_count": self.attributes_l[<int>ATTR_RSDOS_BYTE_CNT],
+            "max_ppm_interval": self.attributes_l[<int>ATTR_RSDOS_MAX_PPM_INTERVAL],
+            "icmp_mismatches": self.attributes_l[<int>ATTR_RSDOS_ICMP_MISMATCHES],
+            "initial_packet": initpkt,
+        }
+
+
+    cpdef void resetRecord(self):
+        self.pktcontentlen = 0
+        if self.packetcontent != NULL:
+            PyMem_Free(self.packetcontent)
+        super(AvroRsdos, self).resetRecord()
+
+    cpdef bytes getRsdosPacketString(self):
+        return <bytes>self.packetcontent[:self.pktcontentlen]
+
+    cpdef int setRsdosPacketString(self, const unsigned char[:] buf,
+            const unsigned int maxlen):
+
+        cdef parsedString astr
+
+        astr = read_string(buf, maxlen, addNullTerm=False)
+        if astr.toskip == 0:
+            return 0
+        self.packetcontent = astr.start
+        self.pktcontentlen = astr.strlen
+        self.sizeinbuf += astr.toskip + astr.strlen
+        return 1
+
+@cython.final
+cdef class AvroRsdosReader(AvroReader):
+
+    def __init__(self, filepath):
+        super().__init__(filepath)
+        self.currentrec = AvroRsdos()
+
+    cdef int _parseNextRecord(self, const unsigned char[:] buf,
+            const unsigned int maxlen):
+
+        cdef unsigned int offset, offinc
+        cdef RsdosAttribute i
+
+        if maxlen == 0:
+            return 0
+        offset = 0
+
+        self.currentrec.resetRecord()
+
+        for i in range(0, ATTR_RSDOS_LATEST_TIME_USEC + 1):
+            offinc = self.currentrec.parseNumeric(buf[offset:],
+                    maxlen - offset, i)
+            if offinc == 0:
+                return 0
+            offset += offinc
+
+        return self.currentrec.setRsdosPacketString(buf[offset:],
+                maxlen - offset);
+
+
+# vim: set sw=4 tabstop=4 softtabstop=4 expandtab :