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