diff --git a/dissect/cstruct/parser.py b/dissect/cstruct/parser.py index 9f465e2..32e23f0 100644 --- a/dissect/cstruct/parser.py +++ b/dissect/cstruct/parser.py @@ -57,6 +57,7 @@ def _tokencollection() -> TokenCollection: TOK = TokenCollection() TOK.add(r"#\[(?P[^\]]+)\](?=\s*)", "CONFIG_FLAG") TOK.add(r"#define\s+(?P[^\s]+)(?P[^\r\n]*)", "DEFINE") + TOK.add(r"#undef\s+(?P[^\s]+)\s*", "UNDEF") TOK.add(r"#ifdef\s+(?P[^\s]+)\s*", "IFDEF") TOK.add(r"#ifndef\s+(?P[^\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: diff --git a/tests/test_parser.py b/tests/test_parser.py index 71fe8ab..f77da47 100644 --- a/tests/test_parser.py +++ b/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