Add support for #undef (#137)

Author: Schamper

dissect/cstruct/parser.py

@@ -57,6 +57,7 @@ def _tokencollection() -> TokenCollection:
         TOK = TokenCollection()
         TOK.add(r"#\[(?P<values>[^\]]+)\](?=\s*)", "CONFIG_FLAG")
         TOK.add(r"#define\s+(?P<name>[^\s]+)(?P<value>[^\r\n]*)", "DEFINE")
+        TOK.add(r"#undef\s+(?P<name>[^\s]+)\s*", "UNDEF")
         TOK.add(r"#ifdef\s+(?P<name>[^\s]+)\s*", "IFDEF")
         TOK.add(r"#ifndef\s+(?P<name>[^\s]+)\s*", "IFNDEF")
         TOK.add(r"#else\s*", "ELSE")
@@ -154,6 +155,16 @@ def _constant(self, tokens: TokenConsumer) -> None:
 
         self.cstruct.consts[match["name"]] = value
 
+    def _undef(self, tokens: TokenConsumer) -> None:
+        const = tokens.consume()
+        pattern = self.TOK.patterns[self.TOK.UNDEF]
+        match = pattern.match(const.value).groupdict()
+
+        if match["name"] in self.cstruct.consts:
+            del self.cstruct.consts[match["name"]]
+        else:
+            raise ParserError(f"line {self._lineno(const)}: constant {match['name']!r} not defined")
+
     def _enum(self, tokens: TokenConsumer) -> None:
         # We cheat with enums because the entire enum is in the token
         etok = tokens.consume()
@@ -443,6 +454,8 @@ def parse(self, data: str) -> None:
                 self._config_flag(tokens)
             elif token == self.TOK.DEFINE:
                 self._constant(tokens)
+            elif token == self.TOK.UNDEF:
+                self._undef(tokens)
             elif token == self.TOK.TYPEDEF:
                 self._typedef(tokens)
             elif token == self.TOK.STRUCT:

tests/test_parser.py

@@ -163,6 +163,19 @@ def test_typedef_pointer(cs: cstruct) -> None:
     assert cs.PIMAGE_DATA_DIRECTORY.type == cs._IMAGE_DATA_DIRECTORY
 
 
+def test_undef(cs: cstruct) -> None:
+    cdef = """
+    #define MY_CONST 42
+    #undef MY_CONST
+    """
+    cs.load(cdef)
+
+    assert "MY_CONST" not in cs.consts
+
+    with pytest.raises(ParserError, match="line 1: constant 'MY_CONST' not defined"):
+        cs.load("#undef MY_CONST")  # This should raise an error since MY_CONST is not defined
+
+
 def test_conditional_ifdef(cs: cstruct) -> None:
     cdef = """
     #define MY_CONST 42