Support printing NULL-terminated strings.

An example is NULL-terminated XFRM crypto and auth algorithm
names.

Bug: 34812052
Test: new unit test passes
Change-Id: I91c5ddfbab06ca8dd21b289c6981cf9855158713

2 files changed, 37 insertions, 5 deletions

diff --git a/net/test/cstruct.py b/net/test/cstruct.py
index 8fe916a..43c47a2 100644
--- a/net/test/cstruct.py
+++ b/net/test/cstruct.py
@@ -19,6 +19,8 @@ Example usage:
 >>> # Declare a struct type by specifying name, field formats and field names.
 ... # Field formats are the same as those used in the struct module, except:
 ... # - S: Nested Struct.
+... # - A: NULL-padded ASCII string. Like s, but printing ignores contiguous
+... #      trailing NULL blocks at the end.
 ... import cstruct
 >>> NLMsgHdr = cstruct.Struct("NLMsgHdr", "=LHHLL", "length type flags seq pid")
 >>>
@@ -61,11 +63,16 @@ import string
 import struct
 
 
+def CalcSize(fmt):
+  if "A" in fmt:
+    fmt = fmt.replace("A", "s")
+  return struct.calcsize(fmt)
+
 def CalcNumElements(fmt):
   prevlen = len(fmt)
   fmt = fmt.replace("S", "")
   numstructs = prevlen - len(fmt)
-  size = struct.calcsize(fmt)
+  size = CalcSize(fmt)
   elements = struct.unpack(fmt, "\x00" * size)
   return len(elements) + numstructs
 
@@ -93,6 +100,8 @@ def Struct(name, fmt, fieldnames, substructs={}):
     _fieldnames = fieldnames
     # Dict mapping field indices to nested struct classes.
     _nested = {}
+    # List of string fields that are ASCII strings.
+    _asciiz = set()
 
     if isinstance(_fieldnames, str):
       _fieldnames = _fieldnames.split(" ")
@@ -108,11 +117,16 @@ def Struct(name, fmt, fieldnames, substructs={}):
         _nested[index] = substructs[laststructindex]
         laststructindex += 1
         _format += "%ds" % len(_nested[index])
+      elif fmt[i] == "A":
+        # Null-terminated ASCII string.
+        index = CalcNumElements(fmt[:i])
+        _asciiz.add(index)
+        _format += "s"
       else:
          # Standard struct format character.
         _format += fmt[i]
 
-    _length = struct.calcsize(_format)
+    _length = CalcSize(_format)
 
     def _SetValues(self, values):
       super(CStruct, self).__setattr__("_values", list(values))
@@ -178,9 +192,11 @@ def Struct(name, fmt, fieldnames, substructs={}):
 
     def __str__(self):
       def FieldDesc(index, name, value):
-        if isinstance(value, str) and any(
-            c not in string.printable for c in value):
-          value = value.encode("hex")
+        if isinstance(value, str):
+          if index in self._asciiz:
+            value = value.rstrip("\x00")
+          elif any(c not in string.printable for c in value):
+            value = value.encode("hex")
         return "%s=%s" % (name, value)
 
       descriptions = [
diff --git a/net/test/cstruct_test.py b/net/test/cstruct_test.py
index 70ba90c..fdcbd55 100755
--- a/net/test/cstruct_test.py
+++ b/net/test/cstruct_test.py
@@ -94,6 +94,22 @@ class CstructTest(unittest.TestCase):
     unpacked = DoubleNested(expected)
     self.CheckEquals(unpacked, d)
 
+  def testNullTerminatedStrings(self):
+    TestStruct = cstruct.Struct("TestStruct", "B16si16AH",
+                                "byte1 string2 int3 ascii4 word5")
+    nullstr = "hello" + (16 - len("hello")) * "\x00"
+
+    t = TestStruct((2, nullstr, 12345, nullstr, 33210))
+    expected = ("TestStruct(byte1=2, string2=68656c6c6f0000000000000000000000,"
+                " int3=12345, ascii4=hello, word5=33210)")
+    self.assertEquals(expected, str(t))
+
+    embeddednull = "hello\x00visible123"
+    t = TestStruct((2, embeddednull, 12345, embeddednull, 33210))
+    expected = ("TestStruct(byte1=2, string2=68656c6c6f0076697369626c65313233,"
+                " int3=12345, ascii4=hello\x00visible123, word5=33210)")
+    self.assertEquals(expected, str(t))
+
 
 if __name__ == "__main__":
   unittest.main()