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 t!("app.entry_point"),
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: t!("app.unknown_section").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: t!("app.unknown_section").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(&t!("errors.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(&t!("errors.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
365 .header
366 .get_decoder()
367 .expect(&t!("errors.create_decoder"));
368 let mut offsets = Vec::new();
369 let mut instructions = Vec::new();
370 let mut to_byte = self.data.len();
371
372 let from_instruction = self.assembly_offsets[from_byte];
373 let mut current_byte = from_byte;
374 let mut ip_offset = 0;
375
376 loop {
377 if current_byte >= maximum_code_byte {
378 to_byte = maximum_code_byte;
379 break;
380 }
381 let bytes = &self.data.bytes()[current_byte..maximum_code_byte];
382 let decoded = decoder
383 .disasm_count(bytes, virtual_address + ip_offset, 1)
384 .expect(&t!("errors.disassemble"));
385 if decoded.is_empty() {
386 break;
387 }
388 let instruction = decoded.iter().next().unwrap();
389 ip_offset += instruction.len() as u64;
390 let old_instruction = self.get_instruction_at(current_byte);
391 let instruction_tag = InstructionTag {
392 instruction: Instruction::new(instruction, self.header.get_symbols()),
393 file_address: current_byte as u64,
394 };
395 let new_assembly_line = AssemblyLine::Instruction(instruction_tag.clone());
396 if old_instruction.is_same_instruction(&new_assembly_line)
397 && current_byte - from_byte >= modifyied_bytes
398 {
399 to_byte = old_instruction.file_address() as usize;
400 break;
401 }
402 instructions.push(new_assembly_line);
403 for _ in 0..instruction.len() {
404 offsets.push(from_instruction + instructions.len() - 1);
405 current_byte += 1;
406 }
407 }
408 if from_byte == to_byte {
409 return;
410 }
411
412 let to_instruction = self
413 .assembly_offsets
414 .get(to_byte)
415 .cloned()
416 .unwrap_or(self.assembly_instructions.len());
417
418 let mut original_instruction_count = 1;
419 let mut original_instruction_ip = self.assembly_offsets[from_byte];
420 for i in from_byte..to_byte {
421 if self.assembly_offsets[i] != original_instruction_ip {
422 original_instruction_count += 1;
423 original_instruction_ip = self.assembly_offsets[i];
424 }
425 }
426
427 let new_instruction_count = instructions.len();
428
429 let delta = new_instruction_count as isize - original_instruction_count as isize;
430
431 self.assembly_offsets.splice(from_byte..to_byte, offsets);
432 if delta != 0 {
433 for offset in self.assembly_offsets.iter_mut().skip(to_byte) {
434 *offset = (*offset as isize + delta) as usize;
435 }
436 }
437
438 for i in from_instruction..to_instruction {
439 if let AssemblyLine::Instruction(instruction) = &self.assembly_instructions[i] {
440 self.log(
441 NotificationLevel::Debug,
442 t!(
443 "app.messages.removing_instruction",
444 instruction = instruction.instruction,
445 address = self.assembly_instructions[i].file_address() : {:#X}
446 ),
447 );
448 } else {
449 break;
450 }
451 }
452 for instruction in instructions.iter() {
453 if let AssemblyLine::Instruction(instruction_tag) = instruction {
454 self.log(
455 NotificationLevel::Debug,
456 t!(
457 "app.messages.adding_instruction",
458 instruction = instruction_tag.instruction,
459 address = instruction.file_address() : {:#X}
460 ),
461 );
462 }
463 }
464
465 self.assembly_instructions
466 .splice(from_instruction..to_instruction, instructions);
467 }
468 }
469
470 pub(in crate::app) fn parse_header(&mut self) -> Header {
471 let mut app_context = get_app_context!(self);
472 match self.plugin_manager.try_parse_header(&mut app_context) {
473 Some(header) => Header::CustomHeader(header),
474 None => {
475 Header::parse_header(self.data.bytes(), self.filesystem.pwd(), &self.filesystem)
476 }
477 }
478 }
479}
480
481#[cfg(test)]
482mod test {
483 use crate::app::comments::Comments;
484
485 use super::*;
486 #[test]
487 fn test_assembly_line() {
488 let file_address = 0xdeadbeef;
489 let virtual_address = 0xcafebabe;
490
491 let al = AssemblyLine::Instruction(InstructionTag {
492 instruction: Instruction {
493 mnemonic: "mov".to_string(),
494 operands: "rax, rbx".to_string(),
495 virtual_address,
496 bytes: vec![0x48, 0x89, 0xd8],
497 },
498 file_address,
499 });
500 let mut comments = Comments::default();
501 comments.insert(file_address, "This is a comment".into());
502 let line = al.to_line(
503 &ColorSettings::get_default_dark_theme(),
504 0,
505 &Header::None,
506 0,
507 &comments,
508 );
509
510 let contains_mnemonic = line.spans.iter().any(|span| span.content.contains("mov"));
511 assert!(contains_mnemonic);
512 let contains_operands = line
513 .spans
514 .iter()
515 .any(|span| span.content.contains("rax, rbx"));
516 assert!(contains_operands);
517 let comma_count = line
518 .spans
519 .iter()
520 .map(|span| span.content.chars().filter(|c| *c == ',').count())
521 .sum::<usize>();
522 assert_eq!(comma_count, 1);
523 let contains_virtual_address = line
524 .spans
525 .iter()
526 .any(|span| span.content.contains(&format!("{virtual_address:X}")));
527 assert!(contains_virtual_address);
528 let contains_file_address = line
529 .spans
530 .iter()
531 .any(|span| span.content.contains(&format!("{file_address:X}")));
532 assert!(contains_file_address);
533 let contains_comment = line
534 .spans
535 .iter()
536 .any(|span| span.content.contains(comments.get(&file_address).unwrap()));
537 assert!(
538 contains_comment,
539 "Comment {} not found in line {:?}",
540 comments.get(&0).unwrap(),
541 line
542 );
543
544 let section_size = 0x1000;
545
546 let al = AssemblyLine::SectionTag(SectionTag {
547 name: ".text".to_string(),
548 file_address,
549 virtual_address,
550 size: section_size,
551 });
552
553 let line = al.to_line(
554 &ColorSettings::get_default_dark_theme(),
555 0,
556 &Header::None,
557 0,
558 &comments,
559 );
560
561 let contains_section_name = line.spans.iter().any(|span| span.content.contains(".text"));
562 assert!(contains_section_name);
563 let contains_virtual_address = line
564 .spans
565 .iter()
566 .any(|span| span.content.contains(&format!("{virtual_address:X}")));
567 assert!(contains_virtual_address);
568 let contains_file_address = line
569 .spans
570 .iter()
571 .any(|span| span.content.contains(&format!("{file_address:X}")));
572 assert!(contains_file_address);
573 let contains_size = line
574 .spans
575 .iter()
576 .any(|span| span.content.contains(&format!("{section_size}B")));
577 assert!(contains_size);
578 let contains_comment = line
579 .spans
580 .iter()
581 .any(|span| span.content.contains(comments.get(&file_address).unwrap()));
582 assert!(contains_comment);
583 }
584
585 #[test]
586 fn test_disassemble_and_patch() {
587 let data = vec![0x48, 0x89, 0xd8, 0x48, 0x89, 0xc1, 0x48, 0x89, 0xc0];
588 let mut app = App::mockup(data);
589 app.resize_to_size(80, 24);
590 let mut expected_instructions = vec!["mov rax, rbx", "mov rcx, rax", "mov rax, rax"];
591 expected_instructions.reverse();
592 let mut text_found = false;
593 for line in app.assembly_instructions.iter() {
594 match line {
595 AssemblyLine::Instruction(instruction) => {
596 assert!(text_found, "Instructions must be after .text section");
597 let instruction_text = expected_instructions
598 .pop()
599 .expect("There are too many instructions in assembly_instructions");
600 assert!(instruction
601 .instruction
602 .to_string()
603 .contains(instruction_text));
604 }
605 AssemblyLine::SectionTag(section) => {
606 if text_found {
607 panic!("There are too many .text sections in assembly_instructions");
608 }
609 assert_eq!(section.name, ".text");
610 text_found = true;
611 }
612 }
613 }
614 assert!(text_found);
615
616 app.patch("nop; nop; nop;");
617 let expected_data = vec![0x90, 0x90, 0x90, 0x48, 0x89, 0xc1, 0x48, 0x89, 0xc0];
618 let mut expected_instructions = vec!["nop", "nop", "nop", "mov rcx, rax", "mov rax, rax"];
619 expected_instructions.reverse();
620 assert_eq!(app.data.bytes(), expected_data);
621 text_found = false;
622 for line in app.assembly_instructions.iter() {
623 match line {
624 AssemblyLine::Instruction(instruction) => {
625 assert!(text_found, "Instructions must be after .text section");
626 let instruction_text = expected_instructions
627 .pop()
628 .expect("There are too many instructions in assembly_instructions");
629 assert!(instruction
630 .instruction
631 .to_string()
632 .contains(instruction_text));
633 }
634 AssemblyLine::SectionTag(section) => {
635 if text_found {
636 panic!("There are too many .text sections in assembly_instructions");
637 }
638 assert_eq!(section.name, ".text");
639 text_found = true;
640 }
641 }
642 }
643 assert!(text_found);
644
645 app.move_cursor(2, 0, false);
647
648 app.patch("jmp rax");
649 let expected_data = vec![0x90, 0xff, 0xe0, 0x48, 0x89, 0xc1, 0x48, 0x89, 0xc0];
650 let mut expected_instructions = vec!["nop", "jmp rax", "mov rcx, rax", "mov rax, rax"];
651 expected_instructions.reverse();
652 assert_eq!(app.data.bytes(), expected_data);
653 text_found = false;
654 for line in app.assembly_instructions.iter() {
655 match line {
656 AssemblyLine::Instruction(instruction) => {
657 assert!(text_found, "Instructions must be after .text section");
658 let instruction_text = expected_instructions
659 .pop()
660 .expect("There are too many instructions in assembly_instructions");
661 assert!(instruction
662 .instruction
663 .to_string()
664 .contains(instruction_text));
665 }
666 AssemblyLine::SectionTag(section) => {
667 if text_found {
668 panic!("There are too many .text sections in assembly_instructions");
669 }
670 assert_eq!(section.name, ".text");
671 text_found = true;
672 }
673 }
674 }
675 assert!(text_found);
676 }
677
678 #[test]
679 fn test_bad_instruction() {
680 let data = vec![0x06, 0x0e, 0x07];
681 let app = App::mockup(data);
682 for line in app.assembly_instructions.iter() {
683 if let AssemblyLine::Instruction(instruction) = line {
684 let contains_bad_instruction =
685 instruction.instruction.to_string().contains(".byte");
686 assert!(
687 contains_bad_instruction,
688 "Found {} instead of .byte ...",
689 instruction.instruction
690 );
691 }
692 }
693 }
694}