1use ratatui::text::{Line, Span};
2
3use crate::{
4 app::{
5 instruction::Instruction, log::NotificationLevel, settings::color_settings::ColorSettings,
6 App,
7 },
8 asm::assembler::assemble,
9 get_app_context,
10 headers::{section::Section, Header},
11};
12
13use super::{
14 assembly_line::AssemblyLine, instruction_tag::InstructionTag, section_tag::SectionTag,
15};
16
17impl App {
18 pub(in crate::app) fn find_symbols(&self, filter: &str) -> Vec<(u64, String)> {
19 if filter.is_empty() {
20 return Vec::new();
21 }
22 let symbol_table = self.header.get_symbols();
23 if let Some(symbol_table) = symbol_table {
24 let mut symbols: Vec<(u64, String)> = symbol_table
25 .iter()
26 .filter(|(_, symbol)| symbol.contains(filter))
27 .map(|(address, symbol)| (*address, symbol.clone()))
28 .collect();
29 symbols.sort_by_key(|(_, symbol)| symbol.len());
30 symbols
31 } else {
32 Vec::new()
33 }
34 }
35
36 pub(super) fn instruction_to_line(
37 color_settings: &ColorSettings,
38 instruction: &InstructionTag,
39 selected: bool,
40 header: &Header,
41 address_min_width: usize,
42 ) -> Line<'static> {
43 let symbol_table = header.get_symbols();
44 let mut line = Line::default();
45 line.spans.push(Span::styled(
46 format!("{:>address_min_width$X}", instruction.file_address),
47 if selected {
48 color_settings.assembly_selected
49 } else {
50 color_settings.assembly_address
51 },
52 ));
53 line.spans.push(Span::raw(" "));
54
55 let mnemonic = instruction.instruction.mnemonic();
56 let args = instruction.instruction.operands();
57 let mnemonic_style = match instruction.instruction.mnemonic() {
58 "nop" => color_settings.assembly_nop,
59 ".byte" => color_settings.assembly_bad,
60 _ => color_settings.assembly_default,
61 };
62
63 line.spans
64 .push(Span::styled(mnemonic.to_string(), mnemonic_style));
65 line.spans.push(Span::raw(" "));
66 line.spans.push(Span::raw(args.to_string()));
67 if let Some(symbol_table) = symbol_table {
68 if let Some(symbol) = symbol_table.get(&instruction.instruction.ip()) {
69 line.spans.push(Span::raw(" "));
70 line.spans.push(Span::styled(
71 format!("<{}>", symbol),
72 color_settings.assembly_symbol,
73 ));
74 }
75 }
76 if instruction.instruction.ip() == header.entry_point() {
77 line.spans.push(Span::raw(" "));
78 line.spans.push(Span::styled(
79 "EntryPoint",
80 color_settings.assembly_entry_point,
81 ));
82 }
83 line.spans.push(Span::styled(
84 format!(" @{:X}", instruction.instruction.ip()),
85 color_settings.assembly_virtual_address,
86 ));
87
88 line
89 }
90
91 pub(in crate::app) fn sections_from_bytes(
92 bytes: &[u8],
93 header: &Header,
94 ) -> (Vec<usize>, Vec<AssemblyLine>) {
95 let mut line_offsets = vec![0; bytes.len()];
96 let mut lines = Vec::new();
97 let mut sections = header.get_sections();
98 if sections.is_empty() {
99 sections.push(Section {
100 name: ".text".to_string(),
101 virtual_address: 0,
102 file_offset: 0,
103 size: bytes.len() as u64,
104 });
105 }
106
107 let mut current_byte = 0;
108 for section in sections {
109 if section.file_offset > current_byte as u64 {
110 lines.push(AssemblyLine::SectionTag(SectionTag {
111 name: "Unknown".to_string(),
112 file_address: current_byte as u64,
113 virtual_address: 0,
114 size: section.file_offset as usize - current_byte,
115 }));
116 for _ in 0..section.file_offset as usize - current_byte {
117 line_offsets[current_byte] = lines.len() - 1;
118 current_byte += 1;
119 }
120 }
121 current_byte = section.file_offset as usize;
123 match section.name.as_str() {
124 ".text" | "__text" => {
125 lines.push(AssemblyLine::SectionTag(SectionTag {
126 name: section.name.clone(),
127 file_address: section.file_offset,
128 virtual_address: section.virtual_address,
129 size: section.size as usize,
130 }));
131 let (offsets, instructions) = Self::assembly_from_section(
132 bytes,
133 header,
134 section.virtual_address as usize,
135 current_byte,
136 section.size as usize,
137 lines.len(),
138 );
139 line_offsets.splice(
140 section.file_offset as usize
141 ..section.file_offset as usize + section.size as usize,
142 offsets,
143 );
144 lines.extend(instructions);
145 current_byte += section.size as usize;
146 }
147 name => {
148 lines.push(AssemblyLine::SectionTag(SectionTag {
149 name: name.to_string(),
150 file_address: section.file_offset,
151 virtual_address: section.virtual_address,
152 size: section.size as usize,
153 }));
154 for _ in 0..section.size as usize {
155 line_offsets[current_byte] = lines.len() - 1;
156 current_byte += 1;
157 }
158 }
159 }
160 }
161 if current_byte < bytes.len() {
162 lines.push(AssemblyLine::SectionTag(SectionTag {
163 name: "Unknown".to_string(),
164 file_address: current_byte as u64,
165 virtual_address: 0,
166 size: bytes.len() - current_byte,
167 }));
168 let initial_current_byte = current_byte;
169 for _ in initial_current_byte..bytes.len() {
170 line_offsets[current_byte] = lines.len() - 1;
171 current_byte += 1;
172 }
173 }
174
175 (line_offsets, lines)
176 }
177
178 pub(in crate::app) fn assembly_from_section(
179 bytes: &[u8],
180 header: &Header,
181 starting_ip: usize,
182 starting_file_address: usize,
183 section_size: usize,
184 starting_sections: usize,
185 ) -> (Vec<usize>, Vec<AssemblyLine>) {
186 let mut line_offsets = vec![0; section_size];
187 let mut instructions = Vec::new();
188 let mut current_byte = 0;
189 let decoder = header.get_decoder().expect("Failed to create decoder");
190 let decoded = decoder
191 .disasm_all(
192 &bytes[starting_file_address..starting_file_address + section_size],
193 starting_ip as u64,
194 )
195 .expect("Failed to disassemble");
196 for instruction in decoded.iter() {
197 let instruction_tag = InstructionTag {
198 instruction: Instruction::new(instruction, header.get_symbols()),
199 file_address: current_byte as u64 + starting_file_address as u64,
200 };
201 instructions.push(AssemblyLine::Instruction(instruction_tag));
202 for _ in 0..instruction.len() {
203 line_offsets[current_byte] = starting_sections + instructions.len() - 1;
204 current_byte += 1;
205 }
206 }
207 (line_offsets, instructions)
208 }
209
210 pub(in crate::app) fn bytes_from_assembly(
211 &self,
212 assembly: &str,
213 starting_virtual_address: u64,
214 ) -> Result<Vec<u8>, String> {
215 let bytes = assemble(assembly, starting_virtual_address, &self.header);
216 match bytes {
217 Ok(bytes) => Ok(bytes),
218 Err(e) => Err(e.to_string()),
219 }
220 }
221
222 pub(in crate::app) fn patch_bytes(
223 &mut self,
224 bytes: &[u8],
225 start_from_beginning_of_instruction: bool,
226 ) {
227 let current_instruction = self.get_current_instruction();
228 if let Some(current_instruction) = current_instruction {
229 let current_instruction = current_instruction.clone();
230 let current_ip = match ¤t_instruction {
231 AssemblyLine::Instruction(instruction) => instruction.file_address,
232 AssemblyLine::SectionTag(_) => self.get_cursor_position().global_byte_index as u64,
233 };
234 let instruction_offset = if start_from_beginning_of_instruction {
235 0
236 } else {
237 self.get_cursor_position().global_byte_index - current_ip as usize
238 };
239 let offset = current_ip as usize + instruction_offset;
240 let mut bytes = bytes.to_vec();
241 let mut app_context = get_app_context!(self);
242 app_context.offset = offset;
243 self.plugin_manager.on_edit(&mut bytes, &mut app_context);
244
245 let modified_bytes = self.data.push_change(offset, bytes);
246
247 self.edit_assembly(modified_bytes + instruction_offset);
248 }
249 }
250
251 pub(in crate::app) fn patch(&mut self, assembly: &str) {
252 if let Some(current_instruction) = self.get_current_instruction() {
253 let current_virtual_address =
254 if let AssemblyLine::Instruction(instruction) = current_instruction {
255 instruction.instruction.ip()
256 } else {
257 self.get_cursor_position().global_byte_index as u64
258 };
259 let bytes = self.bytes_from_assembly(assembly, current_virtual_address);
260 match bytes {
261 Ok(bytes) => self.patch_bytes(&bytes, true),
262 Err(e) => {
263 self.log(NotificationLevel::Error, &e);
264 }
265 }
266 }
267 }
268
269 pub(in crate::app) fn get_assembly_view_scroll(&self) -> usize {
270 let cursor_position = self.get_cursor_position();
271 let current_ip = cursor_position
272 .global_byte_index
273 .min(self.assembly_offsets.len() - 1);
274 let current_scroll = self.assembly_offsets[current_ip];
275
276 let visible_lines = self.screen_size.1 - self.vertical_margin;
277 let center_of_view = visible_lines / 2;
278 let view_scroll = (current_scroll as isize - center_of_view as isize).clamp(
279 0,
280 (self.assembly_instructions.len() as isize - visible_lines as isize).max(0),
281 );
282
283 view_scroll as usize
284 }
285
286 pub(in crate::app) fn get_current_instruction(&self) -> Option<&AssemblyLine> {
287 let global_byte_index = self.get_cursor_position().global_byte_index;
288 if global_byte_index >= self.assembly_offsets.len() {
289 return None;
290 }
291 let current_instruction_index = self.assembly_offsets[global_byte_index];
292 Some(&self.assembly_instructions[current_instruction_index])
293 }
294
295 pub(in crate::app) fn get_instruction_at(&self, index: usize) -> &AssemblyLine {
296 let current_instruction_index = self.assembly_offsets[index];
297 &self.assembly_instructions[current_instruction_index]
298 }
299
300 pub(in crate::app) fn edit_assembly(&mut self, modifyied_bytes: usize) {
301 let current_instruction = self.get_current_instruction();
302 if let Some(current_instruction) = current_instruction {
303 let from_byte = current_instruction.file_address() as usize;
304 let virtual_address = current_instruction.virtual_address();
305 let text_section = self.header.get_text_section();
306 let (is_inside_text_section, maximum_code_byte) =
307 if let Some(text_section) = text_section {
308 (
309 from_byte >= text_section.file_offset as usize
310 && from_byte
311 < text_section.file_offset as usize + text_section.size as usize,
312 text_section.file_offset as usize + text_section.size as usize,
313 )
314 } else {
315 (true, self.data.len())
316 };
317 if !is_inside_text_section {
318 return;
319 }
320 let decoder = self.header.get_decoder().expect("Failed to create decoder");
321 let mut offsets = Vec::new();
322 let mut instructions = Vec::new();
323 let mut to_byte = self.data.len();
324
325 let from_instruction = self.assembly_offsets[from_byte];
326 let mut current_byte = from_byte;
327 let mut ip_offset = 0;
328
329 loop {
330 if current_byte >= maximum_code_byte {
331 to_byte = maximum_code_byte;
332 break;
333 }
334 let bytes = &self.data.bytes()[current_byte..maximum_code_byte];
335 let decoded = decoder
336 .disasm_count(bytes, virtual_address + ip_offset, 1)
337 .expect("Failed to disassemble");
338 if decoded.len() == 0 {
339 break;
340 }
341 let instruction = decoded.iter().next().unwrap();
342 ip_offset += instruction.len() as u64;
343 let old_instruction = self.get_instruction_at(current_byte);
344 let instruction_tag = InstructionTag {
345 instruction: Instruction::new(instruction, self.header.get_symbols()),
346 file_address: current_byte as u64,
347 };
348 let new_assembly_line = AssemblyLine::Instruction(instruction_tag.clone());
349 if old_instruction.is_same_instruction(&new_assembly_line)
350 && current_byte - from_byte >= modifyied_bytes
351 {
352 to_byte = old_instruction.file_address() as usize;
353 break;
354 }
355 instructions.push(new_assembly_line);
356 for _ in 0..instruction.len() {
357 offsets.push(from_instruction + instructions.len() - 1);
358 current_byte += 1;
359 }
360 }
361 if from_byte == to_byte {
362 return;
363 }
364
365 let to_instruction = self
366 .assembly_offsets
367 .get(to_byte)
368 .cloned()
369 .unwrap_or(self.assembly_instructions.len());
370
371 let mut original_instruction_count = 1;
372 let mut original_instruction_ip = self.assembly_offsets[from_byte];
373 for i in from_byte..to_byte {
374 if self.assembly_offsets[i] != original_instruction_ip {
375 original_instruction_count += 1;
376 original_instruction_ip = self.assembly_offsets[i];
377 }
378 }
379
380 let new_instruction_count = instructions.len();
381
382 let delta = new_instruction_count as isize - original_instruction_count as isize;
383
384 self.assembly_offsets.splice(from_byte..to_byte, offsets);
385 if delta != 0 {
386 for offset in self.assembly_offsets.iter_mut().skip(to_byte) {
387 *offset = (*offset as isize + delta) as usize;
388 }
389 }
390
391 for i in from_instruction..to_instruction {
392 if let AssemblyLine::Instruction(instruction) = &self.assembly_instructions[i] {
393 self.log(
394 NotificationLevel::Debug,
395 &format!(
396 "Removing instruction \"{}\" at {:X}",
397 instruction.instruction,
398 self.assembly_instructions[i].file_address()
399 ),
400 );
401 } else {
402 break;
403 }
404 }
405 for instruction in instructions.iter() {
406 if let AssemblyLine::Instruction(instruction_tag) = instruction {
407 self.log(
408 NotificationLevel::Debug,
409 &format!(
410 "Adding instruction \"{}\" at {:X}",
411 instruction_tag.instruction,
412 instruction.file_address()
413 ),
414 );
415 }
416 }
417
418 self.assembly_instructions
419 .splice(from_instruction..to_instruction, instructions);
420 }
421 }
422
423 pub(in crate::app) fn parse_header(&mut self) -> Header {
424 let mut app_context = get_app_context!(self);
425 match self.plugin_manager.try_parse_header(&mut app_context) {
426 Some(header) => Header::CustomHeader(header),
427 None => {
428 Header::parse_header(self.data.bytes(), self.filesystem.pwd(), &self.filesystem)
429 }
430 }
431 }
432}
433
434#[cfg(test)]
435mod test {
436 use std::vec;
437
438 use super::*;
439 #[test]
440 fn test_assembly_line() {
441 let file_address = 0xdeadbeef;
442 let virtual_address = 0xcafebabe;
443
444 let al = AssemblyLine::Instruction(InstructionTag {
445 instruction: Instruction {
446 mnemonic: "mov".to_string(),
447 operands: "rax, rbx".to_string(),
448 virtual_address,
449 bytes: vec![0x48, 0x89, 0xd8],
450 },
451 file_address,
452 });
453 let line = al.to_line(
454 &ColorSettings::get_default_dark_theme(),
455 0,
456 &Header::None,
457 0,
458 );
459
460 let contains_mnemonic = line.spans.iter().any(|span| span.content.contains("mov"));
461 assert!(contains_mnemonic);
462 let contains_operands = line
463 .spans
464 .iter()
465 .any(|span| span.content.contains("rax, rbx"));
466 assert!(contains_operands);
467 let comma_count = line
468 .spans
469 .iter()
470 .map(|span| span.content.chars().filter(|c| *c == ',').count())
471 .sum::<usize>();
472 assert_eq!(comma_count, 1);
473 let contains_virtual_address = line
474 .spans
475 .iter()
476 .any(|span| span.content.contains(&format!("{:X}", virtual_address)));
477 assert!(contains_virtual_address);
478 let contains_file_address = line
479 .spans
480 .iter()
481 .any(|span| span.content.contains(&format!("{:X}", file_address)));
482 assert!(contains_file_address);
483
484 let section_size = 0x1000;
485
486 let al = AssemblyLine::SectionTag(SectionTag {
487 name: ".text".to_string(),
488 file_address,
489 virtual_address,
490 size: section_size,
491 });
492
493 let line = al.to_line(
494 &ColorSettings::get_default_dark_theme(),
495 0,
496 &Header::None,
497 0,
498 );
499
500 let contains_section_name = line.spans.iter().any(|span| span.content.contains(".text"));
501 assert!(contains_section_name);
502 let contains_virtual_address = line
503 .spans
504 .iter()
505 .any(|span| span.content.contains(&format!("{:X}", virtual_address)));
506 assert!(contains_virtual_address);
507 let contains_file_address = line
508 .spans
509 .iter()
510 .any(|span| span.content.contains(&format!("{:X}", file_address)));
511 assert!(contains_file_address);
512 let contains_size = line
513 .spans
514 .iter()
515 .any(|span| span.content.contains(&format!("{}B", section_size)));
516 assert!(contains_size);
517 }
518
519 #[test]
520 fn test_disassemble_and_patch() {
521 let data = vec![0x48, 0x89, 0xd8, 0x48, 0x89, 0xc1, 0x48, 0x89, 0xc0];
522 let mut app = App::mockup(data);
523 app.resize_to_size(80, 24);
524 let mut expected_instructions = vec!["mov rax, rbx", "mov rcx, rax", "mov rax, rax"];
525 expected_instructions.reverse();
526 let mut text_found = false;
527 for line in app.assembly_instructions.iter() {
528 match line {
529 AssemblyLine::Instruction(instruction) => {
530 assert!(text_found, "Instructions must be after .text section");
531 let instruction_text = expected_instructions
532 .pop()
533 .expect("There are too many instructions in assembly_instructions");
534 assert!(instruction
535 .instruction
536 .to_string()
537 .contains(instruction_text));
538 }
539 AssemblyLine::SectionTag(section) => {
540 if text_found {
541 panic!("There are too many .text sections in assembly_instructions");
542 }
543 assert_eq!(section.name, ".text");
544 text_found = true;
545 }
546 }
547 }
548 assert!(text_found);
549
550 app.patch("nop; nop; nop;");
551 let expected_data = vec![0x90, 0x90, 0x90, 0x48, 0x89, 0xc1, 0x48, 0x89, 0xc0];
552 let mut expected_instructions = vec!["nop", "nop", "nop", "mov rcx, rax", "mov rax, rax"];
553 expected_instructions.reverse();
554 assert_eq!(app.data.bytes(), expected_data);
555 text_found = false;
556 for line in app.assembly_instructions.iter() {
557 match line {
558 AssemblyLine::Instruction(instruction) => {
559 assert!(text_found, "Instructions must be after .text section");
560 let instruction_text = expected_instructions
561 .pop()
562 .expect("There are too many instructions in assembly_instructions");
563 assert!(instruction
564 .instruction
565 .to_string()
566 .contains(instruction_text));
567 }
568 AssemblyLine::SectionTag(section) => {
569 if text_found {
570 panic!("There are too many .text sections in assembly_instructions");
571 }
572 assert_eq!(section.name, ".text");
573 text_found = true;
574 }
575 }
576 }
577 assert!(text_found);
578
579 app.move_cursor(2, 0, false);
581
582 app.patch("jmp rax");
583 let expected_data = vec![0x90, 0xff, 0xe0, 0x48, 0x89, 0xc1, 0x48, 0x89, 0xc0];
584 let mut expected_instructions = vec!["nop", "jmp rax", "mov rcx, rax", "mov rax, rax"];
585 expected_instructions.reverse();
586 assert_eq!(app.data.bytes(), expected_data);
587 text_found = false;
588 for line in app.assembly_instructions.iter() {
589 match line {
590 AssemblyLine::Instruction(instruction) => {
591 assert!(text_found, "Instructions must be after .text section");
592 let instruction_text = expected_instructions
593 .pop()
594 .expect("There are too many instructions in assembly_instructions");
595 assert!(instruction
596 .instruction
597 .to_string()
598 .contains(instruction_text));
599 }
600 AssemblyLine::SectionTag(section) => {
601 if text_found {
602 panic!("There are too many .text sections in assembly_instructions");
603 }
604 assert_eq!(section.name, ".text");
605 text_found = true;
606 }
607 }
608 }
609 assert!(text_found);
610 }
611
612 #[test]
613 fn test_bad_instruction() {
614 let data = vec![0x06, 0x0e, 0x07];
615 let app = App::mockup(data);
616 for line in app.assembly_instructions.iter() {
617 if let AssemblyLine::Instruction(instruction) = line {
618 let contains_bad_instruction =
619 instruction.instruction.to_string().contains(".byte");
620 assert!(
621 contains_bad_instruction,
622 "Found {} instead of .byte ...",
623 instruction.instruction
624 );
625 }
626 }
627 }
628}