Merge pull request #227 from CESNET/pytrap_copyMessage

pytrap: add UnirecTemplate.copyMessage()

Author: cejkato2

pytrap/src/unirecmodule.c

@@ -1687,6 +1687,30 @@ UnirecTemplate_copy(pytrap_unirectemplate *self)
     return (PyObject *) n;
 }
 
+static PyObject *
+UnirecTemplate_copyMessage(pytrap_unirectemplate *self, PyObject *args, PyObject *keywds)
+{
+    pytrap_unirectemplate *src = NULL;
+
+    static char *kwlist[] = {"src", NULL};
+    if (!PyArg_ParseTupleAndKeywords(args, keywds, "O!", kwlist, &pytrap_UnirecTemplate, &src)) {
+        return NULL;
+    }
+
+    if (self->data_obj == NULL || self->data == NULL) {
+        PyErr_SetString(PyExc_TypeError, "Data not allocated by createMessage(), cannot copy content.");
+        return NULL;
+    }
+    if (src->data_obj == NULL) {
+        PyErr_SetString(PyExc_TypeError, "UnirecTemplate src does not have any data to copy.");
+        return NULL;
+    }
+
+    ur_copy_fields(self->urtmplt, self->data, src->urtmplt, src->data);
+
+    Py_RETURN_NONE;
+}
+
 static PyObject *
 UnirecTemplate_strRecord(pytrap_unirectemplate *self)
 {
@@ -2008,6 +2032,15 @@ static PyMethodDef pytrap_unirectemplate_methods[] = {
             "    UnirecTemplate: New copy of object (not just reference).\n"
         },
 
+        {"copyMessage", (PyCFunction) UnirecTemplate_copyMessage, METH_VARARGS | METH_KEYWORDS,
+            "Set the content of the record using src UnirecTemplate argument.\n\n"
+            "The current object must have enough allocated memory using createMessage()!!!\n\n"
+            "Args:\n"
+            "    src (UnirecTemplate): UniRec message with data (previously set by setData)"
+            "Raises:\n"
+            "    TypeError: self or src has no allocated data.\n"
+        },
+
         {"strRecord", (PyCFunction) UnirecTemplate_strRecord, METH_NOARGS,
             "Get values of record in readable format.\n\n"
             "Returns:\n"

pytrap/test/unirectemplate_unittest.py

@@ -858,6 +858,82 @@ def runTest(self):
         self.assertEqual(astr, bstr)
         self.assertEqual(astr, '(ipaddr SRC_IP,time TIME_FIRST,uint32 ABC,uint32 BCD,bytes STREAMBYTES,string TEXT)')
 
+class CopyMessageTest(unittest.TestCase):
+    def runTest(self):
+        import pytrap
+        a = pytrap.UnirecTemplate("ipaddr DST_IP,ipaddr SRC_IP,time TIME_FIRST,time TIME_LAST,uint32 ABC,uint32 BCD,string TEXT,string TEXT2,bytes STREAMBYTES,bytes STREAMBYTES2")
+        a.createMessage(100)
+        a.setFromDict({ "SRC_IP": pytrap.UnirecIPAddr("10.0.0.1"), "DST_IP": pytrap.UnirecIPAddr("10.0.0.1"),
+                  "TIME_FIRST": pytrap.UnirecTime(1669885132, 853),
+                  "TIME_LAST": pytrap.UnirecTime(1669885132, 853),
+                  "ABC": 123, "BCD": 321, "TEXT": "some_text", "TEXT2": "some_text2",
+                  "STREAMBYTES": b"abc\x01\x02\x00\x03", "STREAMBYTES2": b"abc\x02\x03\x00\x04"})
+        astr = str(a)
+        b = pytrap.UnirecTemplate("ipaddr SRC_IP,time TIME_FIRST,uint32 ABC,string TEXT2,bytes STREAMBYTES2,string UNKNOWN")
+        b.createMessage(100)
+        b.copyMessage(a)
+        # print(str(b), b.strRecord())
+        # print(b.getData())
+        self.assertEqual(b.getDict(), {"SRC_IP":
+            pytrap.UnirecIPAddr("10.0.0.1"), "TIME_FIRST": pytrap.UnirecTime(1669885132, 853),
+            "ABC": 123, "TEXT2": "some_text2", "STREAMBYTES2": b"abc\x02\x03\x00\x04", "UNKNOWN": ""})
+        a.STREAMBYTES2 = b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+        a.TEXT2 = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+        b.UNKNOWN = "unknown text"
+        b.copyMessage(a)
+        self.assertEqual(b.STREAMBYTES2, a.STREAMBYTES2)
+        self.assertEqual(b.TEXT2, a.TEXT2)
+        del(a)
+        # print(str(b), b.strRecord())
+        # print(b.getData())
+        self.assertEqual(b.getDict(), {"SRC_IP":
+            pytrap.UnirecIPAddr("10.0.0.1"), "TIME_FIRST": pytrap.UnirecTime(1669885132, 853),
+            "ABC": 123, "TEXT2": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", "STREAMBYTES2": b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", "UNKNOWN": "unknown text"})
+
+class CopyMessageEdgeCasesTest(unittest.TestCase):
+    def runTest(self):
+        import pytrap
+        data = { "SRC_IP": pytrap.UnirecIPAddr("10.0.0.1"), "DST_IP": pytrap.UnirecIPAddr("10.0.0.1"),
+                  "TIME_FIRST": pytrap.UnirecTime(1669885132, 853),
+                  "TIME_LAST": pytrap.UnirecTime(1669885132, 853),
+                  "ABC": 123, "BCD": 321, "TEXT": "some_text", "TEXT2": "some_text2",
+                  "STREAMBYTES": b"abc\x01\x02\x00\x03", "STREAMBYTES2": b"abc\x02\x03\x00\x04"}
+        a = pytrap.UnirecTemplate("ipaddr DST_IP,ipaddr SRC_IP,time TIME_FIRST,time TIME_LAST,uint32 ABC,uint32 BCD,string TEXT,string TEXT2,bytes STREAMBYTES,bytes STREAMBYTES2")
+        b = pytrap.UnirecTemplate("ipaddr SRC_IP,time TIME_FIRST,uint32 ABC,string TEXT2,bytes STREAMBYTES2,string UNKNOWN")
+        with self.assertRaises(TypeError):
+            # both a and b have no allocated data
+            b.copyMessage(a)
+        a.createMessage(100)
+        a.setFromDict(data)
+        with self.assertRaises(TypeError):
+            # b has no allocated data
+            b.copyMessage(a)
+
+        b.createMessage(100)
+        # recreate a
+        a = pytrap.UnirecTemplate("ipaddr DST_IP,ipaddr SRC_IP,time TIME_FIRST,time TIME_LAST,uint32 ABC,uint32 BCD,string TEXT,string TEXT2,bytes STREAMBYTES,bytes STREAMBYTES2")
+        with self.assertRaises(TypeError):
+            # a has no allocated data
+            b.copyMessage(a)
+
+        a.createMessage(100)
+        a.setFromDict(data)
+        b = pytrap.UnirecTemplate("")
+        b.createMessage()
+        b.copyMessage(a)
+        self.assertEqual(b.getDict(), {})
+
+        del(a)
+        del(b)
+
+        a = pytrap.UnirecTemplate("")
+        a.createMessage()
+        b = pytrap.UnirecTemplate("ipaddr DST_IP,time TIME_FIRST,string TEXT")
+        b.createMessage(100)
+        b.copyMessage(a)
+        self.assertEqual(b.getDict(), {"DST_IP": pytrap.UnirecIPAddr("::"), "TIME_FIRST": pytrap.UnirecTime(0, 0), "TEXT": ""})
+
+
 class AllocateBigMessage(unittest.TestCase):
     def runTest(self):
         import pytrap