Change how #define values are handled

Author: Schamper

dissect/cstruct/lexer.py

@@ -274,6 +274,38 @@ def _skip_whitespace(self) -> None:
         if match := _RE_WHITESPACE.match(self.data, self._pos):
             self._take(match.end() - self._pos)
 
+    def _skip_comment(self) -> None:
+        """Skip a comment starting at the current position, if present."""
+        if self._current() == "/":
+            peek = self._peek()
+
+            if peek == "*":
+                self._take(2)  # Consume /*
+                end = self.data.find("*/", self._pos)
+                if end != -1:
+                    self._take(end - self._pos + 2)
+                else:
+                    self._take(len(self.data) - self._pos)
+
+            elif peek == "/":
+                self._take(2)  # Consume //
+                end = self.data.find("\n", self._pos)
+                if end != -1:
+                    self._take(end - self._pos)
+                else:
+                    self._take(len(self.data) - self._pos)
+
+    def _skip_whitespace_and_comments(self) -> None:
+        """Skip whitespace and comments."""
+        while True:
+            start_pos = self._pos
+            self._skip_whitespace()
+            if self.eof:
+                break
+            self._skip_comment()
+            if self.eof or self._pos == start_pos:
+                break
+
     def _read_identifier(self) -> str:
         """Read an identifier starting with a letter or underscore, followed by letters, digits, or underscores."""
         if match := _RE_IDENTIFIER.match(self.data, self._pos):
@@ -348,13 +380,6 @@ def _read_string(self) -> str:
 
         return result
 
-    def _read_angle_string(self) -> str:
-        """Read an angle-bracket string for ``#include <...>``."""
-        self._expect("<")  # Consume `<`
-        value = self._read_until(">", or_eof=False)
-        self._expect(">")  # Consume closing `>`
-        return f"<{value}>"
-
     def _read_preprocessor(self) -> None:
         """Read a preprocessor directive starting with ``#``."""
         line = self._line
@@ -373,23 +398,41 @@ def _read_preprocessor(self) -> None:
             return
 
         # Read the keyword after #
-        self._skip_whitespace()
+        self._skip_whitespace_and_comments()
         keyword = self._read_identifier()
 
         if (token_type := _PP_KEYWORDS.get(keyword)) is None:
             raise self._error(f"unknown preprocessor directive '#{keyword}'", line=line)
 
         self._emit(token_type, keyword, line, col)
 
-        if token_type == TokenType.PP_INCLUDE:
+        if token_type == TokenType.PP_DEFINE:
+            self._skip_whitespace_and_comments()
+
+            if not (name := self._read_identifier()):
+                raise self._error("expected identifier after '#define'", line=line)
+            self._emit(TokenType.IDENTIFIER, name, line)
+
+            self._skip_whitespace_and_comments()
+            if self.eof or self._line != line:
+                # No value, just a simple macro definition
+                return
+
+            if (value := self._read_until("\n")).strip():
+                self._emit(TokenType.STRING, value, line)
+
+        elif token_type == TokenType.PP_INCLUDE:
             # Read include path — either "..." or <...>
-            self._skip_whitespace()
+            self._skip_whitespace_and_comments()
 
             ch = self._current()
             if ch == '"' or ch == "'":
                 value = self._read_string()
             elif ch == "<":
-                value = self._read_angle_string()
+                self._expect("<")  # Consume `<`
+                value = self._read_until(">", or_eof=False)
+                self._expect(">")  # Consume closing `>`
+                value = f"<{value}>"
             else:
                 raise self._error("expected include path after '#include'", line=line)
 
@@ -398,34 +441,11 @@ def _read_preprocessor(self) -> None:
     def tokenize(self) -> list[Token]:
         """Tokenize the input data and return a list of tokens."""
         while not self.eof:
-            self._skip_whitespace()
+            self._skip_whitespace_and_comments()
             if self.eof:
                 break
 
             ch = self._current()
-
-            # Skip comments
-            if ch == "/":
-                peek = self._peek()
-
-                if peek == "*":
-                    self._take(2)  # Consume /*
-                    end = self.data.find("*/", self._pos)
-                    if end != -1:
-                        self._take(end - self._pos + 2)
-                    else:
-                        self._take(len(self.data) - self._pos)
-                    continue
-
-                if peek == "/":
-                    self._take(2)  # Consume //
-                    end = self.data.find("\n", self._pos)
-                    if end != -1:
-                        self._take(end - self._pos)
-                    else:
-                        self._take(len(self.data) - self._pos)
-                    continue
-
             line = self._line
             col = self._column
 
@@ -439,7 +459,7 @@ def tokenize(self) -> list[Token]:
             elif ch in ("b", "B") and self._peek() in ("'", '"'):
                 # Binary string literal like `b"..."` or `b'...'`
                 self._take()  # Consume `b`
-                self._emit(TokenType.BYTES, f"b'{self._read_string()}'", line, col)
+                self._emit(TokenType.BYTES, f"b{self._read_string()!r}", line, col)
 
             elif ch.isdigit():
                 self._emit(TokenType.NUMBER, self._read_number(), line, col)

dissect/cstruct/parser.py

@@ -159,25 +159,38 @@ def _parse_define(self) -> None:
 
         name_token = self._expect(TokenType.IDENTIFIER)
 
-        # Collect all tokens on the same line as the #define
-        parts = []
-        while (token := self._current()).type != TokenType.EOF and token.line == name_token.line:
-            parts.append(self._take().value)
-
-        value = "".join(parts).strip()
-        try:
-            # Lazy mode, try to evaluate as a Python literal first (for simple constants)
-            value = ast.literal_eval(value)
-        except (ValueError, SyntaxError):
-            pass
-
-        # If it's still a string, try to evaluate it as an expression in the context of current constants
-        if isinstance(value, str):
-            try:
-                value = Expression(value).evaluate(self.cs)
-            except (LexerError, ExpressionParserError):
-                # If evaluation fails, just keep it as a string (e.g. for macro-like constants)
-                pass
+        # If there's a value, it's emitted as a single raw STRING token
+        value = ""
+        if (token := self._current()).type == TokenType.STRING and token.line == name_token.line:
+            value = self._take().value
+
+        if value:
+            if value[0] in ('"', "'"):
+                quote = value[0]
+                if value[-1] != quote:
+                    raise self._error("unterminated string literal", token=token)
+
+                # Remove the surrounding and any duplicate quotes
+                value = "".join(ch for ch in value if ch != quote)
+            elif value[:2].lower() in ("b'", 'b"'):
+                quote = value[1]
+                if value[-1] != quote:
+                    raise self._error("unterminated bytes literal", token=token)
+
+                # Remove the leading b and surrounding quotes
+                value = ast.literal_eval(f"b{value[2:-1]!r}")
+            else:
+                try:
+                    # Lazy mode, try to evaluate as a Python literal first (for simple constants)
+                    value = ast.literal_eval(value)
+                except (ValueError, SyntaxError):
+                    # Try to evaluate it as an expression in the context of current constants
+                    if isinstance(value, str):
+                        try:
+                            value = Expression(value).evaluate(self.cs)
+                        except (LexerError, ExpressionParserError):
+                            # If evaluation fails, just keep it as a string (e.g. for macro-like constants)
+                            pass
 
         self.cs.consts[name_token.value] = value
 

tests/test_lexer.py

@@ -76,7 +76,26 @@
         ),
         ("1 - 1", [TokenType.NUMBER, TokenType.MINUS, TokenType.NUMBER], ["1", "-", "1"]),
         # Preprocessor directives
-        ("#define", [TokenType.PP_DEFINE], ["define"]),
+        ("#define FOO", [TokenType.PP_DEFINE, TokenType.IDENTIFIER], ["define", "FOO"]),
+        ("#define FOO 42", [TokenType.PP_DEFINE, TokenType.IDENTIFIER, TokenType.STRING], ["define", "FOO", "42"]),
+        (
+            '#define FOO "hello"',
+            [TokenType.PP_DEFINE, TokenType.IDENTIFIER, TokenType.STRING],
+            ["define", "FOO", '"hello"'],
+        ),
+        (
+            "#define FOO ADCRYPT\00",
+            [TokenType.PP_DEFINE, TokenType.IDENTIFIER, TokenType.STRING],
+            ["define", "FOO", "ADCRYPT\00"],
+        ),
+        (
+            """
+            #define FOO (1 \
+            + 2)
+            """,
+            [TokenType.PP_DEFINE, TokenType.IDENTIFIER, TokenType.STRING],
+            ["define", "FOO", "(1             + 2)"],
+        ),
         ("#undef", [TokenType.PP_UNDEF], ["undef"]),
         ("#ifdef", [TokenType.PP_IFDEF], ["ifdef"]),
         ("#ifndef", [TokenType.PP_IFNDEF], ["ifndef"]),

tests/test_parser.py

@@ -216,22 +216,36 @@ def test_typedef_enum(cs: cstruct) -> None:
 
 def test_define(cs: cstruct) -> None:
     cdef = """
-    #define MY_CONST 42
-    #define MY_EXPR (1 + 2 * 3)
-    #define MY_STR "hello"
-    #define MY_BYTES b"world"
-    #define MY_FUNC(x) ( x == 0 )
-    #define MY_TERNARY(x) ( x ? 1 : 0 )
+    #define CONST 42
+    #define EXPR (1 + 2 * 3)
+    #define RAW somevalue
+    #define STR "hello"
+    #define BYTES b"world"
+    #define NULLRAW ADCRYPT\00
+    #define NULLSTR "ADCRYPT\00"
+    #define NULLBYTES b"ADCRYPT\00"
+    #define ARBITRARYBYTES b"\x00\x01\x02"
+    #define MULTILINE (1 + \
+                        2 + \
+                        3)
+    #define FUNC(x) ( x == 0 )
+    #define TERNARY(x) ( x ? 1 : 0 )
     """
     cs.load(cdef)
 
-    assert cs.consts["MY_CONST"] == 42
-    assert cs.consts["MY_EXPR"] == 7
-    assert cs.consts["MY_STR"] == "hello"
-    assert cs.consts["MY_BYTES"] == b"world"
+    assert cs.consts["CONST"] == 42
+    assert cs.consts["EXPR"] == 7
+    assert cs.consts["RAW"] == "somevalue"
+    assert cs.consts["STR"] == "hello"
+    assert cs.consts["BYTES"] == b"world"
+    assert cs.consts["NULLRAW"] == "ADCRYPT\00"
+    assert cs.consts["NULLSTR"] == "ADCRYPT\00"
+    assert cs.consts["NULLBYTES"] == b"ADCRYPT\00"
+    assert cs.consts["ARBITRARYBYTES"] == b"\x00\x01\x02"
+    assert cs.consts["MULTILINE"] == 6
     # We don't evaluate function-like macros yet, so they should be stored as their raw string representation
-    assert cs.consts["MY_FUNC"] == "(x)(x==0)"
-    assert cs.consts["MY_TERNARY"] == "(x)(x?1:0)"
+    assert cs.consts["FUNC"] == "(x) ( x == 0 )"
+    assert cs.consts["TERNARY"] == "(x) ( x ? 1 : 0 )"
 
 
 def test_define_flag_value(cs: cstruct) -> None:
@@ -517,7 +531,7 @@ def test_preprocessor_define_from_enum_in_struct() -> None:
 
     assert cs.consts["PROTO"] == 6
     assert cs.consts["FLAG_SYNACK"] == 3
-    assert cs.consts["FLAG_BIG"] == "flags.PSH|YOMOMMA"
+    assert cs.consts["FLAG_BIG"] == "flags.PSH | YOMOMMA"
 
     assert "type" in cs.packet.fields
     assert "options" in cs.packet.fields