Skip to content

Commit 1ea9072

Browse files
committed
Fix assembler normalizer
1 parent 51a1d79 commit 1ea9072

4 files changed

Lines changed: 175 additions & 7 deletions

File tree

ptrlib/connection/tube.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -972,6 +972,21 @@ def sendlineafter(self,
972972
return self.after(delim, blocksize, regex, timeout) \
973973
.sendline(data)
974974

975+
def sendafter(self,
976+
delim: DelimiterT,
977+
data: str | bytes,
978+
blocksize: int = 4096,
979+
regex: RegexDelimiterT | None = None,
980+
timeout: int | float = -1) -> int:
981+
"""Wait for a delimiter (or regex), then send data.
982+
983+
Note:
984+
This method is deprecated.
985+
Use `after(delim, blocksize, regex, timeout).send(...)` instead.
986+
"""
987+
return self.after(delim, blocksize, regex, timeout) \
988+
.send(data)
989+
975990
def shutdown(self, target: typing.Literal['send', 'recv']):
976991
"""Shut down a specific connection.
977992

ptrlib/cpu/intel/assembler.py

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -153,10 +153,11 @@ def assemble_nasm(assembly: str, address: int, bits: PtrlibBitsT = 64) -> bytes:
153153
"""
154154
nasm_path = nasm()
155155

156-
assembly = '\n'.join(_normalize_assembly(assembly)[1])
156+
# NASM does not use the 'ptr' keyword in memory operands; normalize without inserting it.
157+
assembly = '\n'.join(_normalize_assembly(assembly, insert_ptr=False)[1])
157158
assembly = f'bits {bits}\n' + assembly
158159
if address > 0:
159-
assembly = 'org {address}\n' + assembly
160+
assembly = f'org {address}\n' + assembly
160161

161162
fname_s = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.S'
162163
fname_o = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.o'
@@ -187,7 +188,7 @@ def assemble_nasm(assembly: str, address: int, bits: PtrlibBitsT = 64) -> bytes:
187188

188189
raise OSError("Assemble failed")
189190

190-
def _normalize_assembly(assembly: str) -> tuple[bool, list[str]]:
191+
def _normalize_assembly(assembly: str, insert_ptr: bool = True) -> tuple[bool, list[str]]:
191192
"""Normalize assembly syntax.
192193
193194
Args:
@@ -232,7 +233,9 @@ def _normalize_assembly(assembly: str) -> tuple[bool, list[str]]:
232233
# Normalize syntax
233234
has_label = False
234235
re_label = re.compile(r'^[a-zA-Z0-9_.]+:')
235-
re_spec = re.compile(r'(byte|word|dword|qword)\s*\[', re.IGNORECASE)
236+
# Patterns for size specifier with/without 'ptr'
237+
re_spec_bracket = re.compile(r'\b(byte|word|dword|qword)\s*\[', re.IGNORECASE)
238+
re_spec_with_ptr = re.compile(r'\b(byte|word|dword|qword)\s+ptr\s*\[', re.IGNORECASE)
236239
re_many_ws = re.compile(r'[ \t]+')
237240
re_comma = re.compile(r',\s*')
238241
re_lbracket = re.compile(r'\[\s*')
@@ -246,8 +249,16 @@ def _normalize_assembly(assembly: str) -> tuple[bool, list[str]]:
246249
u = re_many_ws.sub(' ', token).strip()
247250
# Ensure single space after commas
248251
u = re_comma.sub(', ', u)
249-
# Ensure "spec [" -> "spec ptr ["
250-
u = re_spec.sub(r'\1 ptr \[', u)
252+
# Handle size-specifiers around memory operands depending on assembler
253+
if insert_ptr:
254+
# GAS/Keystone style: ensure "spec ptr ["
255+
u = re_spec_with_ptr.sub(r'\1 ptr [', u)
256+
u = re_spec_bracket.sub(r'\1 ptr [', u)
257+
else:
258+
# NASM style: remove 'ptr' if present; keep "spec ["
259+
u = re_spec_with_ptr.sub(r'\1 [', u)
260+
# Do not add 'ptr' for plain "spec ["
261+
251262
# Tighten brackets "[ ... ]" -> "[...]" then ensure ", [" spacing
252263
u = re_lbracket.sub('[', u)
253264
u = re_rbracket.sub(']', u)

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
setup(
1212
name='ptrlib',
13-
version='3.0.0',
13+
version='3.0.1',
1414
description='CTF library',
1515
long_description=long_description,
1616
long_description_content_type='text/markdown',
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
import os
2+
import unittest
3+
from logging import getLogger, FATAL
4+
from ptrlib.cpu import CPU
5+
6+
_is_windows = os.name == 'nt'
7+
8+
class TestIntelAssemblerNormalize(unittest.TestCase):
9+
"""Normalization-focused tests for Intel assembler backends.
10+
11+
These tests ensure that _normalize_assembly() handles:
12+
- Size specifiers around memory operands (inserting 'ptr' for GAS)
13+
- Comma and bracket spacing normalization
14+
- Semicolon-based instruction splitting
15+
- // line comments and /* block comments */ removal
16+
- Case-insensitive size specifiers
17+
"""
18+
def setUp(self):
19+
getLogger("ptrlib").setLevel(FATAL)
20+
21+
def test_gcc_size_ptr_insertion_qword(self):
22+
"""qword [mem] vs qword ptr [mem] must assemble identically with GCC."""
23+
if _is_windows:
24+
return
25+
cpu = CPU('intel', 64)
26+
cpu.assembler = 'gcc'
27+
a = 'mov qword [rax], 0x1234'
28+
b = 'mov qword ptr [rax], 0x1234'
29+
self.assertEqual(cpu.assemble(a), cpu.assemble(b))
30+
31+
def test_gcc_size_ptr_insertion_other_sizes(self):
32+
"""byte/word/dword [mem] vs ... ptr [mem] must match with GCC."""
33+
if _is_windows:
34+
return
35+
cpu = CPU('intel', 64)
36+
cpu.assembler = 'gcc'
37+
pairs = [
38+
('mov dword [rax], 0x11223344', 'mov dword ptr [rax], 0x11223344'),
39+
('mov word [rax], 0x3344', 'mov word ptr [rax], 0x3344'),
40+
('mov byte [rax], 0x22', 'mov byte ptr [rax], 0x22'),
41+
]
42+
for a, b in pairs:
43+
self.assertEqual(cpu.assemble(a), cpu.assemble(b))
44+
45+
def test_no_spurious_backslash_produced(self):
46+
"""Ensure 'mov qword [rax], 1' assembles (no '\\[' introduced)."""
47+
if _is_windows:
48+
return
49+
cpu = CPU('intel', 64)
50+
cpu.assembler = 'gcc'
51+
# Should not raise OSError
52+
cpu.assemble('mov qword [rax], 1')
53+
54+
def test_comma_and_bracket_spacing_mov(self):
55+
"""mov rax,[rbx] variants should assemble identically."""
56+
if _is_windows:
57+
return
58+
for assembler in ('gcc', 'nasm'):
59+
cpu = CPU('intel', 64)
60+
cpu.assembler = assembler
61+
a = 'mov rax,[rbx]'
62+
b = 'mov rax, [ rbx ]'
63+
self.assertEqual(cpu.assemble(a), cpu.assemble(b))
64+
65+
def test_bracket_tightening_with_lea(self):
66+
"""[ rbx + 8 ] -> [rbx + 8] equivalence (no 'ptr' involved)."""
67+
if _is_windows:
68+
return
69+
for assembler in ('gcc', 'nasm'):
70+
cpu = CPU('intel', 64)
71+
cpu.assembler = assembler
72+
a = 'lea rax,[ rbx + 8 ]'
73+
b = 'lea rax, [rbx+8]'
74+
self.assertEqual(cpu.assemble(a), cpu.assemble(b))
75+
76+
def test_line_and_block_comment_stripping(self):
77+
"""'// ...' and '/* ... */' must be ignored by the tokenizer."""
78+
if _is_windows:
79+
return
80+
for assembler in ('gcc', 'nasm'):
81+
cpu = CPU('intel', 64)
82+
cpu.assembler = assembler
83+
code = 'nop // line\n/* block\n comment */ nop ; nop'
84+
baseline = 'nop\nnop\nnop'
85+
self.assertEqual(cpu.assemble(code), cpu.assemble(baseline))
86+
87+
def test_semicolon_instruction_split(self):
88+
"""'nop; nop ;nop' equals three nops after normalization."""
89+
if _is_windows:
90+
return
91+
for assembler in ('gcc', 'nasm'):
92+
cpu = CPU('intel', 64)
93+
cpu.assembler = assembler
94+
code = 'nop; nop ; nop'
95+
baseline = 'nop\nnop\nnop'
96+
self.assertEqual(cpu.assemble(code), cpu.assemble(baseline))
97+
98+
def test_whitespace_collapse(self):
99+
"""Multiple spaces/tabs collapse and consistent comma spacing."""
100+
if _is_windows:
101+
return
102+
for assembler in ('gcc', 'nasm'):
103+
cpu = CPU('intel', 64)
104+
cpu.assembler = assembler
105+
code = ' mov rax , [ rbx + 0x10 ] '
106+
baseline = 'mov rax, [rbx+0x10]'
107+
self.assertEqual(cpu.assemble(code), cpu.assemble(baseline))
108+
109+
def test_size_keyword_case_insensitive(self):
110+
"""Ensure case-insensitive match for size keywords for GCC path."""
111+
if _is_windows:
112+
return
113+
cpu = CPU('intel', 64)
114+
cpu.assembler = 'gcc'
115+
a = 'mov QWORD [rax], 1'
116+
b = 'mov qword ptr [rax], 1'
117+
self.assertEqual(cpu.assemble(a), cpu.assemble(b))
118+
119+
120+
def test_nasm_removes_ptr_keyword(self):
121+
"""In NASM mode, 'ptr' should be removed; both forms must assemble identically."""
122+
if _is_windows:
123+
return
124+
cpu = CPU('intel', 64)
125+
cpu.assembler = 'nasm'
126+
pairs = [
127+
('mov qword [rax], 0x1122334455667788', 'mov qword ptr [rax], 0x1122334455667788'),
128+
('mov dword [rax], 0x11223344', 'mov dword ptr [rax], 0x11223344'),
129+
('mov word [rax], 0x3344', 'mov word ptr [rax], 0x3344'),
130+
('mov byte [rax], 0x22', 'mov byte ptr [rax], 0x22'),
131+
]
132+
for a, b in pairs:
133+
self.assertEqual(cpu.assemble(a), cpu.assemble(b))
134+
135+
def test_alias_nasm_function_handles_ptr(self):
136+
"""The nasm() helper should also accept 'ptr' and normalize it away."""
137+
if _is_windows:
138+
return
139+
from ptrlib.cpu.assembler import nasm as nasm_asm
140+
a = nasm_asm('mov qword [rax], 1', bits=64)
141+
b = nasm_asm('mov qword ptr [rax], 1', bits=64)
142+
self.assertEqual(a, b)

0 commit comments

Comments
 (0)