From 145976e52f6f2a408aa1093b4232d52bdc2d4be9 Mon Sep 17 00:00:00 2001
From: Schamper <1254028+Schamper@users.noreply.github.com>
Date: Thu, 31 Jul 2025 15:15:02 +0200
Subject: [PATCH] Add support for `#undef`

---
 dissect/cstruct/parser.py | 15 ++++++++++++++-
 tests/test_parser.py      | 13 +++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/dissect/cstruct/parser.py b/dissect/cstruct/parser.py
index 8f214f6..ea28889 100644
--- a/dissect/cstruct/parser.py
+++ b/dissect/cstruct/parser.py
@@ -54,7 +54,8 @@ def __init__(self, cs: cstruct, compiled: bool = True, align: bool = False):
     def _tokencollection() -> TokenCollection:
         TOK = TokenCollection()
         TOK.add(r"#\[(?P<values>[^\]]+)\](?=\s*)", "CONFIG_FLAG")
-        TOK.add(r"#define\s+(?P<name>[^\s]+)\s+(?P<value>[^\r\n]+)\s*", "DEFINE")
+        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"typedef(?=\s)", "TYPEDEF")
         TOK.add(r"(?:struct|union)(?=\s|{)", "STRUCT")
         TOK.add(
@@ -99,6 +100,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()
@@ -382,6 +393,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:
diff --git a/tests/test_parser.py b/tests/test_parser.py
index d049756..bba0f96 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -164,3 +164,16 @@ def test_typedef_pointer(cs: cstruct) -> None:
     assert cs.IMAGE_DATA_DIRECTORY is cs._IMAGE_DATA_DIRECTORY
     assert issubclass(cs.PIMAGE_DATA_DIRECTORY, Pointer)
     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