├── .gitignore
├── LICENSE
├── README.md
├── README.zh_CN.md
├── TimeTravelEmulator.py
└── requirements.txt


/.gitignore:
--------------------------------------------------------------------------------
1 | .vscode
2 | *.log
3 | *.pyc
4 | test.py
5 | TODO
6 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2025 Krietz7
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # Time Travel Emulator
  2 | 
  3 | The `TimeTravelEmulator` is a powerful IDA Pro plugin that brings the concept of time-travel debugging to your reverse engineering workflow. By integrating with the Unicorn emulation framework, it provides a unique capability to record and replay program execution, allowing for detailed analysis of runtime behavior.
  4 | 
  5 | ## Language
  6 | 
  7 | English | [中文](README.zh_CN.md)
  8 | 
  9 | 
 10 | ## Features
 11 | 
 12 | * **Emulation powered by Unicorn**: Utilizes the Unicorn CPU emulator to execute code snippets or entire functions within IDA Pro.
 13 | * **Comprehensive State Capture**: Records detailed snapshots of the CPU registers and memory changes at each instruction executed.
 14 | * **Time Travel Debugging**:
 15 |     * **Forward and Backward Navigation**: Seamlessly step forward (F3) and backward (F2) through the execution history.
 16 |     * **State Switching**: Jump to any captured program state (Q) to analyze the exact conditions at a specific point in time.
 17 |     * **State Switching by ID**: Jump to any captured program state by entering its unique `state_id` (I). The `state_id` is formatted as `$<instruction address>#<execution count>`, allowing precise navigation.
 18 | * **Visualizing Differences**: When switching between states, the plugin highlights the memory and register differences in the plugin interface,
 19 | * **Memory and Register Tracking**: Provides visibility into how memory and register values evolve during execution.
 20 | * **Configurable Emulation**:
 21 |     * Set custom emulation ranges (start and end addresses).
 22 |     * Load initial register values from the current debugger state (if a debugger is attached).
 23 |     * Configure emulation step limits and timeouts.
 24 |     * Option to set custom preprocessing code to set up the Unicorn environment before emulation.
 25 | 
 26 | 
 27 | ## Requirements
 28 | 
 29 | - IDA Pro version >= 7.7
 30 | - Python version >= 3.8
 31 | 
 32 | 
 33 | ## Installation
 34 | 
 35 | 1.  Use `pip install bsdiff4 capstone sortedcontainers unicorn` to install nessesary dependencies for this plugin in your IDAPython.
 36 | 2.  Place the `TimeTravelEmulator.py` file into your IDA Pro `plugins` directory.
 37 | 3.  Restart IDA Pro.
 38 | 
 39 | 
 40 | ## Quick Start
 41 | 
 42 | <img width="483" height="585" alt="image" src="https://github.com/user-attachments/assets/315d56fe-7599-49b0-a9c1-b2285d43b8a8" />
 43 | 
 44 | Press the hotkey `Shift+T` to open the `TimeTravel Emulator: Emulator Settings` dialog.
 45 | 
 46 | After the setting is completed, You can click "Emulate" to start the time-travel emulation.
 47 | 
 48 | <img width="2300" height="1020" alt="image" src="https://github.com/user-attachments/assets/053cee98-c4a0-4146-b3ab-c28bfcef4ad3" />
 49 | 
 50 | Once emulation is completed, the plugin will create a independent view, you can use the following hotkeys to navigate through the recorded states in this view:
 51 | * `F3`: Move to the next recorded state.
 52 | * `F2`: Move to the previous recorded state.
 53 | * `Q`: Switch to a specific state by the instruction position where the cursor is located.
 54 | 
 55 | ## More Details
 56 | ```cpp
 57 | #include <iostream>
 58 | 
 59 | int main() {
 60 |     volatile int a = 5;
 61 |     volatile int b = 3;
 62 |     volatile int result = 0;
 63 | 
 64 |     int i = 0;
 65 |     while (i < 32) {
 66 |         result = result + i;
 67 |         b += a;
 68 |         i+= 1;
 69 |     }
 70 | 
 71 |      printf("Result: %d\n", result);
 72 | 
 73 |     return 0;
 74 | }
 75 | ```
 76 | 
 77 | Source code of a simple program for emulation.
 78 | 
 79 | ### Emulator Settings dialog
 80 | <img width="2312" height="1281" alt="image" src="https://github.com/user-attachments/assets/f30f80ac-c19a-4dc0-9fd0-a61201993bf1" />
 81 | 
 82 | Opening the emulation settings dialog by selecting code and using `Shift+T`.
 83 | 
 84 | In the settings dialog, configure the emulation parameters:
 85 | 
 86 | * **Emulation Execute Range**: Specify a start and end address, or select a function as the emulation range.
 87 | * **Emulate step limit**: Set a limit on the number of instructions to emulate.
 88 | * **Emulate time out**: Set a time limit for the emulation run.
 89 | * **load registers**: Choose whether to load current register values (effective in debug mode).
 90 | * **Set Stack value**: Configure special stack frame register values.
 91 | 
 92 | * **Skip interrupts**: Skip `int` instructions during simulation
 93 | * **Skip unloaded calls**: Skip the call instructions whose target address not loaded
 94 | * **Skip thunk functions**: Skip thunk functions during simulation
 95 | 
 96 | * **Log level & Log file path**: Set logging level and save location for logs.
 97 | 
 98 | <img width="1652" height="971" alt="image" src="https://github.com/user-attachments/assets/a57468de-f946-448c-8bcd-8300ef4d79e9" />
 99 | 
100 | * **Set custom preprocessing code**: Add custom Python code to execute before emulation. This can be used to set up memory, register values, or add hooks.
101 | 
102 | 
103 | ### Time Travel Emulator View
104 | ![Opening Time Travel Emulator View](https://github.com/user-attachments/assets/51ed9758-250c-44a7-a5b5-a98a99dfd250)
105 | 
106 | Click "Emulate" to start the emulation. When the emulation is completed, a new window will open displaying the disassembly, register, and memory views, starting from the first emulated state.
107 | 
108 | #### Disassembly View
109 | <img width="2325" height="1054" alt="image" src="https://github.com/user-attachments/assets/125d0e5f-747a-4ea1-a6e0-2e6943731aa0" />
110 | 
111 | This is the core view, highlighting the currently executed assembly instruction line. The number on the far left indicates how many times the current instruction has been executed throughout the emulation.
112 | 
113 | Use the following hotkeys for navigation:
114 |   * `F3`: Move to the next recorded state.
115 |   * `F2`: Move to the previous recorded state.
116 |   * `Q`:  Jump to the state corresponding to the instruction at the cursor.
117 |   * `I`:  Prompt to enter a `state_id` (e.g., `$0x401000#5`) to jump to a specific state.
118 | 
119 | #### Register View and Memory View
120 | 
121 | ![Register View and Memory View examples](https://github.com/user-attachments/assets/fbcfc301-c639-4f85-a4cf-891909189011)
122 | 
123 | 
124 | When switching between states, these views will highlight the changed register values and memory bytes respectively, making it easy to identify differences.
125 | 
126 | <img width="2559" height="799" alt="image" src="https://github.com/user-attachments/assets/ab080307-4298-4438-8280-c506974520c2" />
127 | 
128 | In the memory view, you can select a range of bytes and press `E` to print the content to the console.
129 | 
130 | #### Auxiliary Views
131 | 
132 | The plugin also provides the following auxiliary views:
133 | 
134 | ##### State Chooser
135 | 
136 | <img width="2547" height="1253" alt="image" src="https://github.com/user-attachments/assets/e60b7573-3259-4f23-8216-f725ffc984b1" />
137 | 
138 | Use the hotkey `C` to open the state chooser view. This view displays all saved states during the emulation. Double-click an entry to jump to the corresponding state.
139 | 
140 | ##### Memory Page Chooser
141 | 
142 | <img width="2242" height="1194" alt="image" src="https://github.com/user-attachments/assets/14a158b3-6249-4cf5-956e-67e69061085b" />
143 | 
144 | The plugin employs a lazy-loading mechanism for memory pages during emulation, meaning pages are mapped and loaded only when accessed by an instruction.
145 | 
146 | Use hotkey `M` to open this view and quickly ascertain the memory pages loaded in the current state.
147 | 
148 | 
149 | ##### Difference Chooser
150 | 
151 | ![Difference chooser use examples](https://github.com/user-attachments/assets/ad40cd4b-c82a-49ca-a07d-ff22705394dd)
152 | 
153 | Use hotkey `D` to open the difference chooser view. This view automatically updates when switching states, providing a clear visual representation of memory and register changes between the two states.
154 | 
155 | 
156 | #### Debugging Mode
157 | 
158 | <img width="2552" height="921" alt="image" src="https://github.com/user-attachments/assets/e0d2697b-3260-4e67-8aeb-8ddcaec43a80" />
159 | 
160 | The plugin supports emulation in IDA's debugging mode and can automatically load the current register values.
161 | 
162 | #### Dynamic Disassembly
163 | 
164 | <img width="2454" height="976" alt="image" src="https://github.com/user-attachments/assets/97a2dc00-077c-4a1b-bffa-9a9221b9ee9a" />
165 | 
166 | 
167 | When the simulation execution flow enters a segment of data that the IDA does not recognize as executable code, the plugin disassembles the data using the Capstone engine.
168 | 
169 | 
170 | ## Supported Architectures
171 | 
172 | The plugin currently supports x86 (32-bit) and x64 (64-bit) architectures.
173 | More architectures will be added in the future.
174 | 
175 | 
176 | ## Contributing
177 | 
178 | Contributions are welcome! Please feel free to open issues or submit pull requests.
179 | 


--------------------------------------------------------------------------------
/README.zh_CN.md:
--------------------------------------------------------------------------------
  1 | # Time Travel Emulator
  2 | 
  3 | `TimeTravelEmulator` 是一款功能强大的 IDA Pro 插件，通过集成 Unicorn 仿真框架，实现时间回溯调试的功能。通过记录和回放程序执行，可以实现对运行时行为的详细分析。
  4 | 
  5 | ## 语言
  6 | 
  7 | [English](README.md) | 中文
  8 | 
  9 | 
 10 | ## 功能特性
 11 | 
 12 | * **基于Unicorn的仿真引擎**：集成 Unicorn 仿真框架，用于在 IDA Pro 中执行代码片段或整个函数。
 13 | * **全面的状态捕获**：记录执行每条指令时CPU寄存器和内存变化的详细快照。
 14 | * **时间回溯调试**：
 15 |     * **前进和后退导航**：无缝地前进（F3）和后退（F2）执行历史记录。
 16 |     * **状态切换**：跳转到任何已捕获的程序状态（Q），分析特定时间点的确切条件。
 17 |     * **基于状态ID的切换**：通过唯一的 `state_id` 跳转（I）到任何已捕获的程序状态(`state_id` 格式为 `$<指令地址>#<执行次数>`)，以实现精确导航。
 18 | * **差异可视化** ：在状态之间切换时，插件会高亮显示插件界面中的内存和寄存器差异，
 19 | * **内存和寄存器跟踪**：提供对程序执行过程中内存和寄存器值变化的可见性。
 20 | * **可配置的仿真**：
 21 |     * 设置自定义仿真范围（起始地址和结束地址）。
 22 |     * 从当前调试器状态加载初始寄存器值（在调试模式下生效）。
 23 |     * 配置仿真步数限制和时间限制。
 24 |     * 设置自定义的预处理代码，以设置 Unicorn 环境。
 25 | 
 26 | 
 27 | ## 运行要求
 28 | 
 29 | - IDA Pro 版本 >= 7.7
 30 | - Python 版本 >= 3.8
 31 | 
 32 | 
 33 | ## 安装
 34 | 
 35 | 1.  在你的IDAPython中使用指令 `pip install bsdiff4 capstone sortedcontainers unicorn` 安装必要的依赖包。
 36 | 2.  将 `TimeTravelEmulator.py` 文件放入 IDA Pro 的 `plugins` 目录。
 37 | 3.  重启 IDA Pro。
 38 | 
 39 | 
 40 | ## 快速入门
 41 | 
 42 | <img width="483" height="585" alt="image" src="https://github.com/user-attachments/assets/315d56fe-7599-49b0-a9c1-b2285d43b8a8" />
 43 | 
 44 | 按下快捷键 `Shift+T` 打开 `EmuTrace: Emulator Settings` (模拟器设置) 对话框。
 45 | 
 46 | 设置完成后，点击 "Emulate" 开始时间回溯仿真。
 47 | 
 48 | 
 49 | <img width="2300" height="1020" alt="image" src="https://github.com/user-attachments/assets/053cee98-c4a0-4146-b3ab-c28bfcef4ad3" />
 50 | 
 51 | 仿真完成后，插件将创建一个独立的视图，您可以使用以下快捷键浏览该视图中记录的状态：
 52 | * `F3`: 移动到下一个录制的状态。
 53 | * `F2`: 移动到上一个录制的状态。
 54 | * `Q`: 通过光标所在的位置跳转到对应指令处的状态。
 55 | 
 56 | ## 更多信息
 57 | ```cpp
 58 | #include <iostream>
 59 | 
 60 | int main() {
 61 |     volatile int a = 5;
 62 |     volatile int b = 3;
 63 |     volatile int result = 0;
 64 | 
 65 |     int i = 0;
 66 |     while (i < 32) {
 67 |         result = result + i;
 68 |         b += a;
 69 |         i+= 1;
 70 |     }
 71 | 
 72 |      printf("Result: %d\n", result);
 73 | 
 74 |     return 0;
 75 | }
 76 | ```
 77 | 
 78 | 用于仿真的示例程序源代码。
 79 | 
 80 | ### 模拟器设置对话框
 81 | <img width="2312" height="1281" alt="image" src="https://github.com/user-attachments/assets/f30f80ac-c19a-4dc0-9fd0-a61201993bf1" />
 82 | 
 83 | 选择要仿真的代码并使用快捷键 `Shift+T` 打开仿真设置对话框。
 84 | 
 85 | 在设置对话框中，配置仿真参数：
 86 | * **Start address & End address**: 模拟运行的起始地址和终止地址
 87 | * **Select funtion range**: 选择某一函数作为模拟执行的运行范围
 88 | 
 89 | * **Emulate step limit**:  模拟运行的运行步数限制
 90 | * **Emulate time out**: 模拟运行的运行时间限制
 91 | 
 92 | * **load registers**: 是否加载当前的寄存器值(在调试模式下生效)
 93 | * **Set Stack value**: 是否设置特殊的栈帧寄存器值
 94 | 
 95 | * **Skip interrupts**：在模拟过程中跳过`int`指令
 96 | * **Skip unloaded calls**：跳过目标地址未被加载的`call`指令
 97 | * **Skip thunk functions**：在模拟过程中跳过thunk函数
 98 | 
 99 | * **Log level & Log file path**: 日志记录等级和保存位置
100 | 
101 | <img width="1652" height="971" alt="image" src="https://github.com/user-attachments/assets/a57468de-f946-448c-8bcd-8300ef4d79e9" />
102 | 
103 | * **Set custom preprocessing code**: 设置自定义的预处理代码，这些代码将会在模拟运行前执行。你可以使用此功能提前设置内存、寄存器值或添加Hook
104 | 
105 | 
106 | ### 时间旅行仿真器视图
107 | 
108 | ![Opening Time Travel Emulator View](https://github.com/user-attachments/assets/51ed9758-250c-44a7-a5b5-a98a99dfd250)
109 | 
110 | 
111 | 点击 "Emulate" 开始仿真。仿真完成后，将打开一个新窗口，显示从第一个模拟状态开始的反汇编、寄存器和内存视图。
112 | 
113 | #### 反汇编视图
114 | <img width="2325" height="1054" alt="image" src="https://github.com/user-attachments/assets/125d0e5f-747a-4ea1-a6e0-2e6943731aa0" />
115 | 
116 | 这是插件的核心视图，标注了当前执行的汇编指令行，最左侧的数字指示当前汇编指令行在整个模拟执行过程中的执行次数。
117 | 
118 | 使用以下快捷键进行导航：
119 | 
120 |   * `F3`: 移动到下一个录制的状态。
121 |   * `F2`: 移动到上一个录制的状态。
122 |   * `Q`:  跳转到光标所在指令处的状态。
123 |   * `I`:  提示输入 state_id（例如，$0x401000#5）以跳转到特定状态。
124 | 
125 | #### 寄存器视图和内存视图
126 | 
127 | ![Register View and Memory View examples](https://github.com/user-attachments/assets/fbcfc301-c639-4f85-a4cf-891909189011)
128 | 
129 | 
130 | 当切换状态时，寄存器视图和内存视图分别高亮改变的寄存器值和内存字节变化。
131 | 
132 | 
133 | <img width="2559" height="799" alt="image" src="https://github.com/user-attachments/assets/ab080307-4298-4438-8280-c506974520c2" />
134 | 
135 | 在内存视图中可以选择字节范围，按 `E` 将内容打印到控制台。
136 | 
137 | 
138 | #### 辅助视图
139 | 
140 | 插件还提供了以下辅助视图：
141 | 
142 | 
143 | ##### 状态选择器
144 | 
145 | 
146 | <img width="2547" height="1253" alt="image" src="https://github.com/user-attachments/assets/e60b7573-3259-4f23-8216-f725ffc984b1" />
147 | 
148 | 使用 `C` 快捷键打开状态选择器视图。此视图显示了模拟执行期间保存的所有状态。双击条目可跳转到相应状态。
149 | 
150 | 
151 | ##### 内存页选择器
152 | 
153 | <img width="2242" height="1194" alt="image" src="https://github.com/user-attachments/assets/14a158b3-6249-4cf5-956e-67e69061085b" />
154 | 
155 | 该插件在仿真期间对内存页面采用延迟加载机制，这意味着只有在被指令访问时才会映射和加载这些页面。
156 | 
157 | 使用快捷键 `M` 打开此视图并快速确定当前状态下加载的内存页。
158 | 
159 | 
160 | ##### 差异选择器
161 | 
162 | ![Difference chooser use examples](https://github.com/user-attachments/assets/ad40cd4b-c82a-49ca-a07d-ff22705394dd)
163 | 
164 | 使用快捷键 `D` 打开差异选择器视图，该视图将会在每次切换状态时同步更新，直观显示这两个状态之间的内存和寄存器变化
165 | 
166 | 
167 | 
168 | #### 调试模式
169 | 
170 | <img width="2552" height="921" alt="image" src="https://github.com/user-attachments/assets/e0d2697b-3260-4e67-8aeb-8ddcaec43a80" />
171 | 
172 | 该插件支持在 IDA 的调试模式下进行仿真，并可以自动加载当前寄存器值。
173 | 
174 | 
175 | #### 动态反汇编
176 | 
177 | <img width="2454" height="976" alt="image" src="https://github.com/user-attachments/assets/97a2dc00-077c-4a1b-bffa-9a9221b9ee9a" />
178 | 
179 | 当模拟执行流进入 IDA 未识别为可执行代码的数据段时，插件将使用 Capstone 引擎对这些数据进行反汇编。
180 | 
181 | 
182 | ## 支持的架构
183 | 
184 | 该插件目前支持 x86 (32-bit) 和 x64 (64-bit) 架构。
185 | 未来将提供更多架构支持
186 | 
187 | 
188 | ## 贡献
189 | 
190 | 欢迎贡献！请随时提交问题或拉取请求。
191 | 


--------------------------------------------------------------------------------
/TimeTravelEmulator.py:
--------------------------------------------------------------------------------
   1 | import idaapi
   2 | import ida_ida
   3 | import ida_kernwin
   4 | import ida_dbg
   5 | import idc
   6 | import ida_name
   7 | import ida_bytes
   8 | import ida_lines
   9 | import ida_segment
  10 | 
  11 | import logging
  12 | import bisect
  13 | from abc import ABC, abstractmethod
  14 | from collections import defaultdict
  15 | from typing import Callable, Dict, Iterator, List, Literal, Optional, Set, Tuple, Union
  16 | from copy import deepcopy
  17 | from dataclasses import dataclass
  18 | from re import split
  19 | 
  20 | import bsdiff4
  21 | from sortedcontainers import SortedDict, SortedList
  22 | from unicorn import *
  23 | from unicorn.x86_const import *
  24 | from capstone import *
  25 | from PyQt5 import QtCore, QtWidgets
  26 | 
  27 | 
  28 | 
  29 | VERSION = '1.1.1'
  30 | 
  31 | PLUGIN_NAME = 'TimeTravelEmulator'
  32 | PLUGIN_HOTKEY = 'Shift+T'
  33 | 
  34 | NEXT_STATE_ACTION_SHORTCUT = "F3"
  35 | PREV_STATE_ACTION_SHORTCUT = "F2"
  36 | CURSOR_STATE_ACTION_SHORTCUT = "Q"
  37 | 
  38 | 
  39 | EXECUTE_INSN_HILIGHT_COLOR = 0xFFD073
  40 | CHANGE_HIGHLIGHT_COLOR   = 0xFFD073
  41 | BYTE_CHANGE_HIGHTLIGHT = ida_kernwin.CK_EXTRA11
  42 | 
  43 | 
  44 | 
  45 | # Define page size and page mask, usually 4 kb
  46 | PAGE_SIZE = 0x1000
  47 | PAGE_MASK = ~(PAGE_SIZE - 1)
  48 | 
  49 | # Define default page permission
  50 | DEFAULT_PAGE_PERMISSION = UC_PROT_WRITE | UC_PROT_READ
  51 | 
  52 | # Define default stack and frame values
  53 | DEFAULT_STACK_POINT_VALUE = 0x70000000
  54 | DEFAULT_BASE_POINT_VALUE = DEFAULT_STACK_POINT_VALUE
  55 | 
  56 | 
  57 | UNICORN_ARCH_MAP =      {
  58 |             "x64": [UC_ARCH_X86, UC_MODE_64],
  59 |             "x86": [UC_ARCH_X86, UC_MODE_32],
  60 | }
  61 | 
  62 | CAPSTONE_ARCH_MAP = {
  63 |     "x64": [CS_ARCH_X86, CS_MODE_64],
  64 |     "x86": [CS_ARCH_X86, CS_MODE_32]
  65 | }
  66 | 
  67 | IDA_PROC_TO_ARCH_MAP = {
  68 |     ("metapc", 64) : "x64",
  69 |     ("metapc", 32) : "x86",
  70 | }
  71 | 
  72 | ARCH_TO_INSN_POINTER_MAP = {
  73 |     "x64" : UC_X86_REG_RIP,
  74 |     "x86" : UC_X86_REG_EIP,
  75 | }
  76 | 
  77 | 
  78 | UNICORN_REGISTERS_MAP = {
  79 |     "x64" : {
  80 |         # General Purpose Registers (GPRs)
  81 |         "RAX": UC_X86_REG_RAX,
  82 |         "RBX": UC_X86_REG_RBX,
  83 |         "RCX": UC_X86_REG_RCX,
  84 |         "RDX": UC_X86_REG_RDX,
  85 |         "RSI": UC_X86_REG_RSI,
  86 |         "RDI": UC_X86_REG_RDI,
  87 |         "RBP": UC_X86_REG_RBP,
  88 |         "RSP": UC_X86_REG_RSP,
  89 |         "RIP": UC_X86_REG_RIP,
  90 |         "R8": UC_X86_REG_R8,
  91 |         "R9": UC_X86_REG_R9,
  92 |         "R10": UC_X86_REG_R10,
  93 |         "R11": UC_X86_REG_R11,
  94 |         "R12": UC_X86_REG_R12,
  95 |         "R13": UC_X86_REG_R13,
  96 |         "R14": UC_X86_REG_R14,
  97 |         "R15": UC_X86_REG_R15,
  98 |         # Instruction Pointer
  99 |         # Flags Register
 100 |         "Rflags": UC_X86_REG_EFLAGS, # In 64-bit, EFLAGS is extended to RFLAGS, but the constant remains EFLAGS
 101 |         # Segment Registers
 102 |         "CS": UC_X86_REG_CS,
 103 |         "SS": UC_X86_REG_SS,
 104 |         "DS": UC_X86_REG_DS,
 105 |         "ES": UC_X86_REG_ES,
 106 |         "FS": UC_X86_REG_FS,
 107 |         "GS": UC_X86_REG_GS,
 108 |     },
 109 |     "x86": {
 110 |         # General Purpose Registers (GPRs)
 111 |         "EAX": UC_X86_REG_EAX,
 112 |         "EBX": UC_X86_REG_EBX,
 113 |         "ECX": UC_X86_REG_ECX,
 114 |         "EDX": UC_X86_REG_EDX,
 115 |         "ESI": UC_X86_REG_ESI,
 116 |         "EDI": UC_X86_REG_EDI,
 117 |         "EBP": UC_X86_REG_EBP,
 118 |         "ESP": UC_X86_REG_ESP,
 119 |         # Instruction Pointer
 120 |         "EIP": UC_X86_REG_EIP,
 121 |         # Flags Register
 122 |         "Eflags": UC_X86_REG_EFLAGS,
 123 |         # Segment Registers
 124 |         "CS": UC_X86_REG_CS,
 125 |         "SS": UC_X86_REG_SS,
 126 |         "DS": UC_X86_REG_DS,
 127 |         "ES": UC_X86_REG_ES,
 128 |         "FS": UC_X86_REG_FS,
 129 |         "GS": UC_X86_REG_GS,
 130 |     }
 131 | }
 132 | 
 133 | IP_REG_NAME_MAP = {
 134 |     "x64": "RIP",
 135 |     "x86": "EIP"
 136 | }
 137 | 
 138 | IDA_PERM_TO_UC_PERM_MAP = {
 139 |     ida_segment.SEGPERM_EXEC : UC_PROT_EXEC,
 140 |     ida_segment.SEGPERM_WRITE : UC_PROT_WRITE,
 141 |     ida_segment.SEGPERM_READ : UC_PROT_READ
 142 | }
 143 | 
 144 | 
 145 | def get_bitness() -> Union[None, Literal[64], Literal[32]]:
 146 |     if ida_ida.inf_is_64bit():
 147 |         return 64
 148 |     elif ida_ida.inf_is_32bit_exactly():
 149 |         return 32
 150 | 
 151 | def get_is_be() -> bool:
 152 |         return ida_ida.inf_is_be()
 153 | 
 154 | def get_arch() -> str:
 155 |     proc_name = idaapi.inf_get_procname()
 156 |     proc_bitness = get_bitness()
 157 | 
 158 |     if proc_bitness == None:
 159 |         return ""
 160 |     if (proc_name, proc_bitness) not in IDA_PROC_TO_ARCH_MAP:
 161 |         return ""
 162 |     return IDA_PROC_TO_ARCH_MAP[(proc_name, proc_bitness)]
 163 | 
 164 | 
 165 | def get_arch_x64_regs_value() -> Dict[int, int]:
 166 |     if not idaapi.is_debugger_on():
 167 |         return {}
 168 |     arch = get_arch()
 169 |     if arch not in UNICORN_REGISTERS_MAP:
 170 |         return {}
 171 | 
 172 |     result: Dict[int, int] = {}
 173 |     regs_map =  UNICORN_REGISTERS_MAP[arch]
 174 |     for reg_name, uc_reg_const in regs_map.items():
 175 |         if reg_name == "Rflags":
 176 |             flag_positions = {
 177 |                 "CF": 0,    # Carry Flag
 178 |                 "PF": 2,    # Parity Flag
 179 |                 "AF": 4,    # Auxiliary Carry Flag
 180 |                 "ZF": 6,    # Zero Flag
 181 |                 "SF": 7,    # Sign Flag
 182 |                 "TF": 8,    # Trap Flag
 183 |                 "IF": 9,    # Interrupt Enable Flag
 184 |                 "DF": 10,   # Direction Flag
 185 |                 "OF": 11,   # Overflow Flag
 186 |                 "IOPL": 12, # I/O Privilege Level (Usually two bit)
 187 |                 "NT": 14,   # Nested Task Flag
 188 |                 "RF": 16,   # Resume Flag
 189 |                 "VM": 17,   # Virtual-8086 Mode Flag
 190 |                 "AC": 18,   # Alignment Check / Access Control Flag
 191 |                 "VIF": 19,  # Virtual Interrupt Flag
 192 |                 "VIP": 20,  # Virtual Interrupt Pending
 193 |                 "ID": 21    # ID Flag
 194 |             }
 195 |             rflags_value = 0
 196 |             for flag_name, bit_pos in flag_positions.items():
 197 |                 try:
 198 |                     flag_val = ida_dbg.get_reg_val(flag_name)
 199 | 
 200 |                     if flag_val is not None:
 201 |                         if flag_name == "IOPL":
 202 |                             rflags_value |= ((flag_val & 0x3) << bit_pos)
 203 |                         else:
 204 |                             if flag_val == 1:
 205 |                                 rflags_value |= (1 << bit_pos)
 206 |                 except Exception as e:
 207 |                     continue
 208 |             result[uc_reg_const] = rflags_value
 209 |         elif reg_name in ["FS", "GS"]:
 210 |             continue
 211 |         else:
 212 |             try:
 213 |                 reg_value = ida_dbg.get_reg_val(reg_name)
 214 |                 result[uc_reg_const] = reg_value
 215 |             except Exception as e:
 216 |                 tte_log_err(f"Error getting register value for {reg_name}: {e}")
 217 |                 pass
 218 |     return result
 219 | 
 220 | 
 221 | def get_arch_x86_regs_value() -> Dict[int, int]:
 222 |     if not idaapi.is_debugger_on():
 223 |         return {}
 224 |     arch = get_arch()
 225 |     if arch not in UNICORN_REGISTERS_MAP:
 226 |         return {}
 227 | 
 228 |     result: Dict[int, int] = {}
 229 |     regs_map =  UNICORN_REGISTERS_MAP[arch]
 230 |     for reg_name, uc_reg_const in regs_map.items():
 231 |         if reg_name == "Eflags":
 232 |             flag_positions = {
 233 |                 "CF": 0,    # Carry Flag
 234 |                 "PF": 2,    # Parity Flag
 235 |                 "AF": 4,    # Auxiliary Carry Flag
 236 |                 "ZF": 6,    # Zero Flag
 237 |                 "SF": 7,    # Sign Flag
 238 |                 "TF": 8,    # Trap Flag
 239 |                 "IF": 9,    # Interrupt Enable Flag
 240 |                 "DF": 10,   # Direction Flag
 241 |                 "OF": 11,   # Overflow Flag
 242 |                 "IOPL": 12, # I/O Privilege Level (Usually two bit)
 243 |                 "NT": 14,   # Nested Task Flag
 244 |                 "RF": 16,   # Resume Flag
 245 |                 "VM": 17,   # Virtual-8086 Mode Flag
 246 |                 "AC": 18,   # Alignment Check / Access Control Flag
 247 |                 "VIF": 19,  # Virtual Interrupt Flag
 248 |                 "VIP": 20,  # Virtual Interrupt Pending
 249 |                 "ID": 21    # ID Flag
 250 |             }
 251 |             rflags_value = 0
 252 |             for flag_name, bit_pos in flag_positions.items():
 253 |                 try:
 254 |                     flag_val = ida_dbg.get_reg_val(flag_name)
 255 | 
 256 |                     if flag_val is not None:
 257 |                         if flag_name == "IOPL":
 258 |                             rflags_value |= ((flag_val & 0x3) << bit_pos)
 259 |                         else:
 260 |                             if flag_val == 1:
 261 |                                 rflags_value |= (1 << bit_pos)
 262 |                 except Exception as e:
 263 |                     continue
 264 |             result[uc_reg_const] = rflags_value
 265 |         elif reg_name in ["CS", "SS", "DS", "ES", "FS", "GS"]:
 266 |             continue
 267 |         else:
 268 |             try:
 269 |                 reg_value = ida_dbg.get_reg_val(reg_name)
 270 |                 result[uc_reg_const] = reg_value
 271 |             except Exception as e:
 272 |                 tte_log_err(f"Error getting register value for {reg_name}: {e}")
 273 |                 pass
 274 |     return result
 275 | 
 276 | 
 277 | def get_page_slice(page_start: int, page_size: int) -> List[Tuple[int, int]]:
 278 |     """
 279 |     Get a slice of segments address of a memory page
 280 | 
 281 |     :param page_start: Start address of the memory page. Must be page aligned.
 282 |     :param page_size: Size of the memory page
 283 |     :return: A list, where the value is a list of tuples (start_ea, end_ea) of the segments slice in the page.
 284 |     """
 285 |     result: List[Tuple[int, int]] = []
 286 |     seg = ida_segment.get_first_seg()
 287 |     page_end = page_start + page_size
 288 |     while seg:
 289 |         if seg.start_ea >= page_end:
 290 |             break
 291 |         # Check for overlap:
 292 |         # (seg.start_ea <= end_address) and (seg.end_ea >= start_address)
 293 |         if seg.start_ea < page_end - 1 and seg.end_ea > page_start:
 294 |             max_addr = max(seg.start_ea, page_start)
 295 |             min_addr = min(seg.end_ea, page_end - 1)
 296 |             result.append((max_addr, min_addr))
 297 |         seg = ida_segment.get_next_seg(seg.end_ea -1 ) # -1 to ensure next segment is properly found
 298 | 
 299 |     return result
 300 | 
 301 | 
 302 | def get_segment_prem(addr: int) -> int: # [ ] TODO: The segments of the program are not always page-aligned, so this situation needs to be considered
 303 |     seg = ida_segment.getseg(addr)
 304 |     if seg is not None:
 305 |         ida_perm = seg.perm
 306 |         uc_perm = 0
 307 |         for ida_bit, uc_bit in IDA_PERM_TO_UC_PERM_MAP.items():
 308 |             if ida_perm & ida_bit:
 309 |                 uc_perm |= uc_bit
 310 |         return uc_perm
 311 |     return DEFAULT_PAGE_PERMISSION
 312 | 
 313 | 
 314 | def catch_dict_patch(
 315 |         base_dict: Dict[str, int],
 316 |         target_dict: Dict[str, int]
 317 | ) -> Dict[str, int]:
 318 |     """
 319 |     Compare the two dictionaries and return different key-value pairs of target_dict relative to base_dict.
 320 |     Prerequisite: Two dictionaries have exactly the same set of keys.
 321 |     Returns the key-value pair contains only the changed values.
 322 |     """
 323 |     return {key: target_dict[key] for key in base_dict if base_dict[key] != target_dict[key]}
 324 | 
 325 | 
 326 | def apply_dict_patch(
 327 |         base_dict: Dict[str, int],
 328 |         patch: Optional[Dict[str, int]]
 329 | ) -> Dict[str, int]:
 330 |     result = base_dict.copy()
 331 |     if patch is not None:
 332 |         result.update(patch)
 333 |     return result
 334 | 
 335 | 
 336 | def catch_bytes_patch(
 337 |         base_bytes_dict: Dict[int, Tuple[int, bytearray]], #
 338 |         target_bytes_dict: Dict[int, Tuple[int, bytearray]]
 339 | ) -> Tuple[Dict[int, Tuple[int, bytes]], Dict[int, Tuple[int, bytearray]]]:
 340 |     """
 341 |     Compare the two byte dictionaries and return different key-value pairs of target_bytes relative to base_bytes.
 342 | 
 343 |     :param base_bytes_dict: Original dictionary of bytes  {addr: (permission, data)}
 344 |     :param target_bytes_dict: Dictionary of bytes to be compared with base_bytes_dict  {addr: (permission, data)}
 345 |     :return patches: Dictionary of binary diffs for modified keys  {addr: (permission, patch)}
 346 |     :return new_entries: Dictionary of new key-value pairs  {addr: (permission, data)}
 347 |     """
 348 |     patches: Dict[int, Tuple[int, bytes]] = {}
 349 |     new_entries: Dict[int, Tuple[int, bytearray]] = {}
 350 | 
 351 | 
 352 |     for addr, (permossion, target_bytes) in target_bytes_dict.items():
 353 |         entry = base_bytes_dict.get(addr)
 354 |         if entry is None:
 355 |             new_entries[addr] = (permossion, target_bytes)
 356 |             continue
 357 |         _, base_bytes = entry
 358 |         if base_bytes != target_bytes:
 359 |             patches[addr] = (permossion, bsdiff4.diff(bytes(base_bytes), bytes(target_bytes)))
 360 |         else:
 361 |             patches[addr] = (permossion, b'')
 362 | 
 363 |     return patches, new_entries
 364 | 
 365 | 
 366 | def apply_bytes_patch(
 367 |         base_bytes_dict: Dict[int, Tuple[int, bytearray]],
 368 |         patches: Dict[int, Tuple[int, bytes]],
 369 |         new_entries: Dict[int, Tuple[int, bytearray]]
 370 | ) -> Dict[int, Tuple[int, bytearray]]:
 371 |     """
 372 |     Apply the generated patches to the base data dictionary and merge new entries.
 373 | 
 374 |     :param base_bytes_dict: Original dictionary of bytes  {addr: (permission, data)}
 375 |     :param patches: Dictionary of binary diffs for modified keys  {addr: (permission, patch)}
 376 |     :param new_entries: Dictionary of new key-value pairs  {addr: (permission, data)}
 377 |     :return updated_dict: Updated dictionary with applied diffs and new entries  {addr: (permission, data)}
 378 |     """
 379 |     updated_dict: Dict[int, Tuple[int, bytearray]] = dict(base_bytes_dict)
 380 | 
 381 |     for addr, (permossion, patch) in patches.items():
 382 |         if patch == b'':
 383 |             continue
 384 |         key_data = updated_dict.get(addr)
 385 |         assert key_data is not None, f"Error: key {addr} not found in base_bytes_dict"
 386 |         _, original_data = key_data
 387 |         try:
 388 |             patched_data = bsdiff4.patch(bytes(original_data), patch)
 389 |             updated_dict[addr] = (permossion, bytearray(patched_data))
 390 |         except Exception as e:
 391 |             raise RuntimeError(f"Error applying patch to key {addr}: {e}")
 392 | 
 393 |     updated_dict.update(new_entries)
 394 |     return updated_dict
 395 | 
 396 | 
 397 | 
 398 | 
 399 | class TTE_Logger:
 400 |     _instance = None
 401 |     _initialized = False
 402 | 
 403 |     def __new__(cls, *args, **kwargs):
 404 |         if cls._instance is None:
 405 |             cls._instance = super(TTE_Logger, cls).__new__(cls)
 406 |         return cls._instance
 407 | 
 408 |     def start(self, log_level, log_file=None):
 409 |         """
 410 |         Start the logger
 411 |         :param log_file: Log file path
 412 |         :param level: Log Level
 413 |         """
 414 |         if self._initialized:
 415 |             tte_log_info("Logger already started.")
 416 |             return
 417 | 
 418 |         # Set logging
 419 |         self.logger = logging.getLogger('TTE_Logger')
 420 |         self.logger.setLevel(log_level)
 421 |         formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
 422 | 
 423 |         # Create handler
 424 |         if log_file:
 425 |             file_handler = logging.FileHandler(log_file)
 426 |             file_handler.setFormatter(formatter)
 427 |             self.logger.addHandler(file_handler)
 428 | 
 429 |         self._initialized = True
 430 |         self.log(logging.INFO, "Logger started.")
 431 | 
 432 |     def stop(self):
 433 |         """
 434 |         Stop the logger
 435 |         """
 436 |         if not self._initialized:
 437 |             return
 438 | 
 439 |         self.log(logging.INFO, "Logger stopping...")
 440 |         handlers = self.logger.handlers[:]
 441 |         for handler in handlers:
 442 |             handler.close()
 443 |             self.logger.removeHandler(handler)
 444 |         self._initialized = False
 445 | 
 446 |     def log(self, level, message):
 447 |         """
 448 |         Logging
 449 |         :param level: Log Level (logging.INFO, logging.ERROR)
 450 |         :param message: Information to be recorded
 451 |         """
 452 |         if not self._initialized:
 453 |             return
 454 | 
 455 |         self.logger.log(level, message)
 456 | 
 457 | def tte_log_info(message):
 458 |     """
 459 |     Logging with INFO level
 460 |     :param message: Information to be recorded
 461 |     """
 462 |     TTE_Logger().log(logging.INFO,message)
 463 | 
 464 | def tte_log_dbg(message):
 465 |     """
 466 |     Logging with DEBUG level
 467 |     :param message: Information to be recorded
 468 |     """
 469 |     TTE_Logger().log(logging.DEBUG,message)
 470 | 
 471 | def tte_log_warn(message):
 472 |     """
 473 |     Logging with WARNING level
 474 |     :param message: Information to be recorded
 475 |     """
 476 |     TTE_Logger().log(logging.WARNING,message)
 477 | 
 478 | def tte_log_err(message):
 479 |     """
 480 |     Logging with ERROR level
 481 |     :param message: Information to be recorded
 482 |     """
 483 |     TTE_Logger().log(logging.ERROR,message)
 484 | 
 485 | 
 486 | @dataclass
 487 | class EmuSettings():
 488 |     start: int = -1
 489 |     end: int = -1
 490 | 
 491 |     is_load_registers: bool = False
 492 |     is_skip_interrupts = False
 493 |     is_skip_unloaded_calls = True
 494 |     is_skip_trunk_funcs: bool = False
 495 |     is_set_stack_value: bool = False
 496 | 
 497 |     time_out: int = 0
 498 |     count: int = 500
 499 |     log_level: int = logging.WARNING
 500 |     log_file_path: Optional[str] = None
 501 | 
 502 |     preprocessing_code: str = ""
 503 | 
 504 | 
 505 | class EmuSettingsForm(idaapi.Form):
 506 | 
 507 |     def __init__(self, start_ea, end_ea) -> None:
 508 |         self.emu_settings = EmuSettings()
 509 | 
 510 |         self.i_start_address: Optional[ida_kernwin.Form.NumericInput] =  None
 511 |         self.i_end_address: Optional[ida_kernwin.Form.NumericInput] = None
 512 | 
 513 |         self.i_emulate_step_limit: Optional[ida_kernwin.Form.NumericInput] =  None
 514 |         self.i_time_out: Optional[ida_kernwin.Form.NumericInput] = None
 515 | 
 516 |         self.c_configs_group: Optional[ida_kernwin.Form.ChkGroupControl] = None
 517 |         self.r_load_register: Optional[ida_kernwin.Form.ChkGroupItemControl] = None
 518 | 
 519 | 
 520 |         self.r_skip_interrupts: Optional[ida_kernwin.Form.ChkGroupItemControl] = None
 521 |         self.r_skip_unloaded_calls: Optional[ida_kernwin.Form.ChkGroupItemControl] = None
 522 |         self.r_skip_thunk_funcs: Optional[ida_kernwin.Form.ChkGroupItemControl] = None
 523 |         self.r_set_stack_value: Optional[ida_kernwin.Form.ChkGroupItemControl] = None
 524 | 
 525 |         self.i_log_level: Optional[ida_kernwin.Form.DropdownListControl] = None
 526 |         self.i_log_file_path: Optional[ida_kernwin.Form.FileInput] = None
 527 | 
 528 | 
 529 |         super().__init__(
 530 |             r'''STARTITEM {id:i_start_address}
 531 | BUTTON YES* Emulate
 532 | TimeTravel Emulator: Emulator Settings
 533 | 
 534 |             {FormChangeCb}
 535 |             Emulation Execute Range:
 536 |             <Start address  :{i_start_address}>
 537 |             <End address    :{i_end_address}>
 538 |             <Select function range    :{b_select_function}>
 539 | 
 540 |             Configs:
 541 |             <Emulate step limit  :{i_emulate_step_limit}>
 542 |             <Emluate time out    :{i_time_out}>
 543 | 
 544 |             <load registers:{r_load_register}>
 545 |             <Set stack value:{r_set_stack_value}>
 546 |             <Skip interrupts:{r_skip_interrupts}>
 547 |             <Skip unloaded calls:{r_skip_unloaded_calls}>
 548 |             <Skip thunk functions:{r_skip_thunk_funcs}>{c_configs_group}>
 549 | 
 550 |             <Log level:{i_log_level}>
 551 |             <Log file path:{i_log_file_path}>
 552 | 
 553 | 
 554 |             Advanced Configs:
 555 |             <Set custom preprocessing code:{b_set_preprocessing_code}>
 556 | 
 557 |             ''',
 558 |             {
 559 |                 'FormChangeCb': self.FormChangeCb(self._on_form_change),
 560 |                 'i_start_address': self.NumericInput(self.FT_ADDR, value=start_ea,  swidth = 30),
 561 |                 'i_end_address': self.NumericInput(self.FT_ADDR, value=end_ea, swidth = 30),
 562 |                 'b_select_function': self.ButtonInput(self._open_select_function_dialog),
 563 | 
 564 |                 'i_emulate_step_limit': self.NumericInput(self.FT_DEC, value=500, swidth = 30),
 565 |                 "i_time_out": self.NumericInput(self.FT_DEC, value=0, swidth = 30),
 566 |                 'c_configs_group': self.ChkGroupControl(("r_load_register", "r_set_stack_value", "r_skip_interrupts", "r_skip_unloaded_calls", "r_skip_thunk_funcs")),
 567 | 
 568 | 
 569 |                 'i_log_level': self.DropdownListControl(
 570 |                     items=["DEBUG", "INFO", "WARNING", "ERROR"],
 571 |                     readonly=True,
 572 |                     selval=2
 573 |                 ),
 574 |                 'i_log_file_path': self.FileInput(save=True, swidth=30),
 575 | 
 576 |                 'b_set_preprocessing_code': self.ButtonInput(self._add_preprocessing_code)
 577 |             }
 578 |          )
 579 |         self.Compile()
 580 |         self.set_default_values()
 581 | 
 582 | 
 583 |     def _on_form_change(self, fid: int):
 584 |         assert self.r_load_register is not None, "r_load_register is not initialized"
 585 |         assert self.r_skip_thunk_funcs is not None, "r_skip_thunk_funcs is not initialized"
 586 | 
 587 |         # Init
 588 |         if fid == -1:
 589 |             if not idaapi.is_debugger_on():
 590 |                 self.EnableField(self.r_load_register, False)
 591 |             else:
 592 |                 self.EnableField(self.r_load_register, True)
 593 | 
 594 | 
 595 |         # Click Yes
 596 |         elif fid == -2:
 597 |             ok = self._set_emu_range()
 598 |             if not ok:
 599 |                 return 0
 600 |             self._set_setting_result()
 601 |             return 1
 602 | 
 603 |         return 1
 604 | 
 605 | 
 606 |     def set_default_values(self):
 607 |         assert self.r_load_register is not None \
 608 |             and self.r_set_stack_value is not None \
 609 |             and self.r_skip_interrupts is not None \
 610 |             and self.r_skip_unloaded_calls is not None \
 611 |             and self.r_skip_thunk_funcs is not None \
 612 | 
 613 |         self.preprocessing_code: str = "mu: unicorn.Uc = emu_executor.get_mu()"
 614 | 
 615 |         if idaapi.is_debugger_on():
 616 |             self.r_load_register.checked = True
 617 |             self.r_set_stack_value.checked = False
 618 |         else:
 619 |             self.r_load_register.checked = False
 620 |             self.r_set_stack_value.checked = True
 621 |             self.r_skip_thunk_funcs.checked = True
 622 | 
 623 |         self.r_skip_interrupts.checked = True
 624 |         self.r_skip_unloaded_calls.checked = True
 625 | 
 626 | 
 627 |     def _open_select_function_dialog(self, code = 0):
 628 |         target_func =  ida_kernwin.choose_func("Select target function range",1)
 629 |         if not target_func:
 630 |             return
 631 | 
 632 |         self._set_emu_range(target_func.start_ea, target_func.end_ea)
 633 | 
 634 | 
 635 |     def _add_preprocessing_code(self, code = 0):
 636 | 
 637 |         class PreprocessingCodeForm(ida_kernwin.Form):
 638 |             def __init__(self, default_multiline_text: str = ""):
 639 |                 self.i_multiline_text: Optional[ida_kernwin.Form.MultiLineTextControl] = None
 640 |                 self.i_load_file: Optional[ida_kernwin.Form.ButtonInput] = None
 641 |                 self.i_save_file: Optional[ida_kernwin.Form.ButtonInput] = None
 642 |                 super().__init__(
 643 |                     r"""STARTITEM {id:i_multiline_text}
 644 | BUTTON YES* OK
 645 | BUTTON NO Cancel
 646 | Preprocessing Code Input
 647 | 
 648 |         {FormChangeCb}
 649 |         <##Enter your preprocessing code here:{i_multiline_text}>
 650 |         <Load from file:{i_load_file}> <Save to file:{i_save_file}>
 651 |         """, {
 652 |                     'FormChangeCb': self.FormChangeCb(self.OnFormChange),
 653 |                     'i_multiline_text': self.MultiLineTextControl(text=default_multiline_text,
 654 |                                                                   flags = ida_kernwin.Form.MultiLineTextControl.TXTF_FIXEDFONT | \
 655 |                                                                   ida_kernwin.Form.MultiLineTextControl.TXTF_ACCEPTTABS,
 656 |                                                                   tabsize = 4,
 657 |                                                                   width = 120,
 658 |                                                                   swidth = 100),
 659 |                     'i_load_file': self.ButtonInput(self.OnLoadFile),
 660 |                     'i_save_file': self.ButtonInput(self.OnSaveFile),
 661 |                 })
 662 |                 self.Compile()
 663 |                 self.user_input_text = ""
 664 | 
 665 |             def OnLoadFile(self, code=0):
 666 |                 file_path = ida_kernwin.ask_file(False, "*.*", "Select a file to load")
 667 |                 if file_path:
 668 |                     try:
 669 |                         with open(file_path, 'r') as f:
 670 |                             content = f.read()
 671 |                             textctrl = ida_kernwin.textctrl_info_t(content)
 672 |                             self.SetControlValue(self.i_multiline_text, textctrl)
 673 |                             idaapi.msg(f"Loaded content from {file_path}\n")
 674 |                     except Exception as e:
 675 |                         ida_kernwin.warning(f"Failed to load file: {e}")
 676 |                 else:
 677 |                     idaapi.msg("No file selected for loading.\n")
 678 | 
 679 |             def OnSaveFile(self, code=0):
 680 |                 file_path = ida_kernwin.ask_file(True, "*.*", "Select a file to save to")
 681 |                 if file_path:
 682 |                     try:
 683 |                         textctrl = self.GetControlValue(self.i_multiline_text)
 684 |                         with open(file_path, 'w') as f:
 685 |                             f.write(textctrl.text)
 686 |                         idaapi.msg(f"Saved content to {file_path}\n")
 687 |                     except Exception as e:
 688 |                         ida_kernwin.warning(f"Failed to save file: {e}")
 689 |                 else:
 690 |                     idaapi.msg("No file selected for saving.\n")
 691 | 
 692 |             def OnFormChange(self, fid):
 693 |                 assert self.i_multiline_text is not None, "i_multiline_text is not initialized"
 694 | 
 695 |                 if fid == -2: # OnFormChange is called with -2 on form close
 696 |                     self.user_input_text = self.GetControlValue(self.i_multiline_text).text
 697 |                 return 1
 698 | 
 699 |             def get_user_text(self):
 700 |                 """
 701 |                 Public method to retrieve the user's input after the form is closed.
 702 |                 """
 703 |                 return self.user_input_text
 704 | 
 705 |         form = PreprocessingCodeForm(self.preprocessing_code)
 706 |         ok = form.Execute()
 707 |         if ok == 1:
 708 |             self.preprocessing_code = form.get_user_text()
 709 |         form.Free()
 710 | 
 711 | 
 712 |     def _set_emu_range(self, start: int = -1, end: int = -1):
 713 |         # set range by parameters
 714 |         if start != -1 and end != -1:
 715 |             start_t = start
 716 |             end_t = end
 717 | 
 718 |         # set range by edit input
 719 |         else:
 720 |             try:
 721 |                 if self.i_start_address is None or self.i_end_address is None:
 722 |                     ida_kernwin.warning("Address input controls are not initialized.")
 723 |                     return 0
 724 |                 start_t: Optional[int] =  self.GetControlValue(self.i_start_address) # type: ignore
 725 |                 end_t: Optional[int] =  self.GetControlValue(self.i_end_address) # type: ignore
 726 |                 if not start_t or not end_t:
 727 |                     raise ValueError()
 728 |             except ValueError:
 729 |                 ida_kernwin.warning("Invalid Input: Please enter valid hexadecimal addresses.")
 730 |                 return 0
 731 | 
 732 |         if start_t < ida_ida.inf_get_min_ea() or end_t > ida_ida.inf_get_max_ea():
 733 |             y = ida_kernwin.ask_yn(2, "Invalid Range: Address out of range, continute?")
 734 |             if y == 0:
 735 |                 return 0
 736 | 
 737 |         # set range
 738 |         self.sim_range_start = start_t
 739 |         self.sim_range_end = end_t
 740 | 
 741 |         if start != -1 and end != -1:
 742 |             self.SetControlValue(self.i_start_address, start_t)
 743 |             self.SetControlValue(self.i_end_address, end_t)
 744 |         return 1
 745 | 
 746 | 
 747 |     def _set_setting_result(self):
 748 |         if self.i_emulate_step_limit is None or \
 749 |            self.c_configs_group is None or \
 750 |            self.i_time_out is None or \
 751 |            self.r_load_register is None or \
 752 |            self.r_skip_interrupts is None or \
 753 |            self.r_skip_unloaded_calls is None or \
 754 |            self.r_skip_thunk_funcs is None or \
 755 |            self.r_set_stack_value is None or \
 756 |            self.i_log_level is None or \
 757 |            self.i_log_file_path is None:
 758 |             ida_kernwin.warning("Form controls are not initialized.")
 759 |             return None
 760 | 
 761 |         self.emu_settings.start = self.sim_range_start
 762 |         self.emu_settings.end = self.sim_range_end
 763 | 
 764 |         self.emu_settings.count = self.GetControlValue(self.i_emulate_step_limit) # type: ignore
 765 |         self.emu_settings.time_out = self.GetControlValue(self.i_time_out) # type: ignore
 766 | 
 767 |         self.emu_settings.is_load_registers = self.GetControlValue(self.r_load_register) # type: ignore
 768 |         self.emu_settings.is_skip_interrupts = self.GetControlValue(self.r_skip_interrupts) # type: ignore
 769 |         self.emu_settings.is_skip_unloaded_calls = self.GetControlValue(self.r_skip_unloaded_calls) # type: ignore
 770 |         self.emu_settings.is_skip_trunk_funcs = self.GetControlValue(self.r_skip_thunk_funcs) # type: ignore
 771 |         self.emu_settings.is_set_stack_value = self.GetControlValue(self.r_set_stack_value) # type: ignore
 772 | 
 773 |         self.emu_settings.preprocessing_code = self.preprocessing_code
 774 | 
 775 |         log_level_map = {
 776 |             0: logging.DEBUG,
 777 |             1: logging.INFO,
 778 |             2: logging.WARNING,
 779 |             3: logging.ERROR,
 780 |         }
 781 | 
 782 |         log_level_value: int = self.GetControlValue(self.i_log_level) # type: ignore
 783 |         self.emu_settings.log_level = log_level_map.get(log_level_value, logging.WARNING)
 784 |         self.emu_settings.log_file_path = self.GetControlValue(self.i_log_file_path) # type: ignore
 785 | 
 786 | 
 787 |     def GetSetting(self):
 788 |         return self.emu_settings
 789 | 
 790 | 
 791 | 
 792 | class EmuExecutor():
 793 |     """
 794 |     A simulation executor class used to manage the operation of the entire program, including initialization, loading, running, saving, etc.
 795 |     """
 796 | 
 797 |     def __init__(self, settings: EmuSettings) -> None:
 798 |         self._is_initialized = False
 799 | 
 800 |         self.arch = get_arch()
 801 |         self.unicorn_arch, self.unicorn_mode = UNICORN_ARCH_MAP[self.arch]
 802 |         self.is_be = get_is_be()
 803 | 
 804 |         self.settings = settings
 805 |         tte_log_info("Create EmuExecutor: arch-{}, mode-{}".format(self.unicorn_arch, self.unicorn_mode))
 806 | 
 807 |         self.loaded_pages: Set[int] = set() # Set(page_start)
 808 | 
 809 |         self.emu_map_mem_callback = []
 810 |         self.emu_run_end_callback = []
 811 |         self.mu_hook_handlers = []
 812 | 
 813 | 
 814 |     def init(self) -> None:
 815 |         """
 816 |         Initialize Unicorn emulation, the method must be called before calling other methods.
 817 |         """
 818 |         if self._is_initialized:
 819 |             tte_log_info("EmuExecutor already initialized.")
 820 |             return
 821 | 
 822 |         # Unicorn instance creation
 823 |         self.mu = unicorn.Uc(self.unicorn_arch, self.unicorn_mode)
 824 | 
 825 |         # Memory mapping and data loading
 826 |         self._map_and_load_binary(self.settings.start, self.settings.end)
 827 | 
 828 |         # Register initialization
 829 |         self._set_regs_init_value()
 830 | 
 831 |         # Hooks setting
 832 |         if self.settings.is_skip_interrupts:
 833 |             self._hook_skip_interrupt()
 834 |         if self.settings.is_skip_trunk_funcs:
 835 |             self._hook_skip_thunk_funcs()
 836 |         if self.settings.is_skip_unloaded_calls:
 837 |             self._hook_skip_unloaded_call()
 838 |         self._hook_mem_unmapped()
 839 | 
 840 |         self._is_initialized = True
 841 | 
 842 | 
 843 |     @staticmethod
 844 |     def execute_preprocessing_code(preprocessing_code, emu_executor: 'EmuExecutor') -> int:
 845 |         """
 846 |         Execute the preprocessing code.
 847 |         """
 848 |         if(len(preprocessing_code) == 0):
 849 |             return 1
 850 |         # idaapi.msg("Executing preprocessing code...\n")
 851 |         try:
 852 |             exec(preprocessing_code)
 853 |         except Exception as e:
 854 |             idaapi.msg(f"Error executing preprocessing code: {e}\n")
 855 |             return idaapi.ask_yn(0, f"Error executing preprocessing code: {e}\nDo you want to continue?")
 856 | 
 857 |         mem_regions_iter = emu_executor.get_mu().mem_regions()
 858 |         for page_start, page_end, perm in mem_regions_iter:
 859 |             current_addr = page_start
 860 |             while current_addr < page_end:
 861 |                 emu_executor.loaded_pages.add(current_addr & PAGE_MASK)
 862 |                 current_addr += PAGE_SIZE
 863 | 
 864 |         return 1
 865 | 
 866 | 
 867 |     def get_mu(self):
 868 |         return self.mu
 869 | 
 870 | 
 871 |     def add_mu_hook(self, htype: int, callback, user_data = None, begin: int = 1, end: int = 0) -> None:
 872 |         handler = self.mu.hook_add(htype, callback, user_data, begin, end)
 873 |         self.mu_hook_handlers.append(handler)
 874 | 
 875 | 
 876 |     def _is_memory_mapped(self, address) -> bool:
 877 |         # return any(map_start <= address < map_end for map_start, map_end, _ in self.mu.mem_regions())
 878 |         return (address & PAGE_MASK) in self.loaded_pages
 879 | 
 880 | 
 881 |     def _map_memory(self, map_start, map_size) -> None:
 882 |         """
 883 |         Map memory pages for the given range.
 884 | 
 885 |         :param map_start: Start address of the mapping. Must be page aligned.
 886 |         :param map_size: Size of the mapping. Must be page aligned.
 887 |         """
 888 |         for page_start in range(map_start, map_start + map_size, PAGE_SIZE):
 889 |             if self._is_memory_mapped(page_start):
 890 |                 continue
 891 |             try:
 892 |                 perm = get_segment_prem(page_start)
 893 |                 self.mu.mem_map(page_start, PAGE_SIZE, perm)
 894 |                 tte_log_info(f"Map memory: 0x{page_start:X}~0x{page_start + PAGE_SIZE - 1 :X} by permission {perm}")
 895 |             except UcError as e:
 896 |                 tte_log_warn(f"Mapping memory failed: {e}")
 897 | 
 898 |             for callback in self.emu_map_mem_callback:
 899 |                 callback(self.mu)
 900 | 
 901 | 
 902 |     def _load_binary(self, load_seg_start, load_seg_end) -> None:
 903 |         """
 904 |         Map memory pages for the given range.
 905 | 
 906 |         :param load_seg_start: Start address of the binary data to load.
 907 |         :param load_seg_end: End address of the binary data to load.
 908 |         """
 909 | 
 910 |         seg_data = ida_bytes.get_bytes(load_seg_start, load_seg_end - load_seg_start + 1)
 911 |         if seg_data:
 912 |             try:
 913 |                 self.mu.mem_write(load_seg_start, seg_data)
 914 |             except UcError as e:
 915 |                 tte_log_info(f"Writing segment failed 0x{load_seg_start:X}~0x{load_seg_end:X}: {e}")
 916 |         else:
 917 |             tte_log_warn(f"Warning: Section 0x{load_seg_start:X}~0x{load_seg_end:X} No data")
 918 | 
 919 | 
 920 |     def _map_and_load_binary(self, start_ea, end_ea) -> None:
 921 |         tte_log_dbg(f"Start map and load binary: 0x{start_ea:X}~0x{end_ea:X}")
 922 | 
 923 |         end_ea = min(end_ea, start_ea + PAGE_SIZE * 8)
 924 |         map_start = start_ea & PAGE_MASK
 925 |         map_size = (end_ea - map_start + PAGE_SIZE - 1) & PAGE_MASK
 926 | 
 927 |         self._map_memory(map_start, map_size)
 928 |         # if not is_address_range_loaded(start_ea, end_ea):
 929 |         #     return
 930 | 
 931 |         for page_start in range(map_start, map_start + map_size, PAGE_SIZE):
 932 |             if page_start not in self.loaded_pages:
 933 |                 load_range_list = get_page_slice(page_start, PAGE_SIZE)
 934 |                 for start_ea, end_ea in load_range_list:
 935 |                     self._load_binary(start_ea, end_ea )
 936 |                     tte_log_info(f"Load binary segments {ida_segment.get_segm_name(idaapi.getseg(start_ea))}: 0x{start_ea:X}~0x{end_ea:X}")
 937 | 
 938 | 
 939 |     def _hook_mem_unmapped(self) -> None:
 940 |         def cb_map_mem_unmapped(uc, access, address, size, value, user_data) -> bool:
 941 |             if access == UC_MEM_READ_UNMAPPED:
 942 |                 tte_log_info(f"Hook callback: Try to read the unmapped memory at 0x{address:X}, size {size}")
 943 |             elif access == UC_MEM_WRITE_UNMAPPED:
 944 |                 tte_log_info(f"Hook callback: Try to write the unmapped memory at 0x{address:X}, size {size}, value 0x{value:X}")
 945 |             elif access == UC_MEM_FETCH_UNMAPPED:
 946 |                 tte_log_info(f"Hook callback: Try to fetch the unmapped memory at 0x{address:X}, size {size}")
 947 | 
 948 |             try:
 949 |                 tte_log_info(f"Hook callback: Map memory during running: 0x{address:X}, 0x{size:X}")
 950 |                 self._map_and_load_binary(address, address + size)
 951 |                 return True
 952 | 
 953 |             except UcError as e:
 954 |                 tte_log_err(f"Hook callback: Mapping memory failed: {e}")
 955 |                 return False
 956 | 
 957 |         self.add_mu_hook(UC_HOOK_MEM_READ_UNMAPPED | UC_HOOK_MEM_WRITE_UNMAPPED | UC_HOOK_MEM_FETCH_UNMAPPED, cb_map_mem_unmapped)
 958 | 
 959 | 
 960 |     def _hook_skip_interrupt(self) -> None:
 961 |         def cb_skip_interrupt(uc, intno, user_data) -> bool:
 962 |             tte_log_info(f"Hook callback: Skip interrupt: {intno}")
 963 |             return True
 964 | 
 965 |         self.add_mu_hook(UC_HOOK_INTR, cb_skip_interrupt)
 966 | 
 967 | 
 968 |     def _hook_skip_thunk_funcs(self) -> None:
 969 | 
 970 |         def execute_x86_ret(uc):
 971 |             if self.unicorn_mode == UC_MODE_64:
 972 |                 stack_pointer_value = uc.reg_read(UC_X86_REG_RSP)
 973 |                 return_address_bytearray = uc.mem_read(stack_pointer_value, 8)
 974 |                 return_address_int = int.from_bytes(return_address_bytearray, byteorder='big' if self.is_be else 'little',  signed=False)
 975 | 
 976 |                 uc.reg_write(UC_X86_REG_RIP, return_address_int)
 977 |                 uc.reg_write(UC_X86_REG_RSP, stack_pointer_value + 8)
 978 | 
 979 |                 tte_log_info(f"Hook callback: Jump to return address: 0x{return_address_int:X}")
 980 |                 return
 981 | 
 982 |             elif self.unicorn_mode == UC_MODE_32:
 983 |                 stack_pointer_value = uc.reg_read(UC_X86_REG_ESP)
 984 |                 return_address_bytearray = uc.mem_read(stack_pointer_value, 4)
 985 |                 return_address_int = int.from_bytes(return_address_bytearray, byteorder='big' if self.is_be else 'little',  signed=False)
 986 | 
 987 |                 uc.reg_write(UC_X86_REG_EIP, return_address_int)
 988 |                 uc.reg_write(UC_X86_REG_ESP, stack_pointer_value + 4)
 989 | 
 990 |                 tte_log_info(f"Hook callback: Jump to return address: 0x{return_address_int:X}")
 991 |                 return
 992 | 
 993 |             else:
 994 |                 return
 995 | 
 996 |         def cb_skip_thunk_func_call(uc, address, size, user_data):
 997 | 
 998 |             flags = idc.get_func_attr(address, idc.FUNCATTR_FLAGS)
 999 |             if flags == idaapi.BADADDR:
1000 |                 return False
1001 |             is_thunk_function = (flags & idaapi.FUNC_THUNK) != 0
1002 | 
1003 |             # If the function is a thunk function, detect the return address on the stack and jump to it.
1004 |             if is_thunk_function:
1005 |                 tte_log_info(f"Hook callback: Skip thunk function at 0x{address:X}")
1006 |                 if self.unicorn_arch == UC_ARCH_X86:
1007 |                     execute_x86_ret(uc)
1008 | 
1009 |         self.add_mu_hook(UC_HOOK_CODE, cb_skip_thunk_func_call)
1010 | 
1011 | 
1012 |     def _hook_skip_unloaded_call(self) -> None:
1013 | 
1014 |         def get_x86_call_target(uc, insn_bytes, address, size):
1015 |             target_address = -1
1016 |             if insn_bytes[0] == 0xE8: # Direct call
1017 |                 relative_offset = int.from_bytes(insn_bytes[1:5], byteorder = 'big' if self.is_be else 'little', signed=True)
1018 |                 target_address = address + size + relative_offset
1019 |                 tte_log_dbg(f"Direct call to 0x{target_address:X}")
1020 | 
1021 |             elif insn_bytes[0] == 0xFF: # Indirect call
1022 |                 insn = None
1023 |                 try:
1024 |                     insn = next(InstrctionParser().parse_instructions(insn_bytes))
1025 |                 except StopIteration:
1026 |                     return
1027 | 
1028 |                 op_str = insn.op_str
1029 |                 if '[' not in op_str: # register indirect calls
1030 |                     register_id = UNICORN_REGISTERS_MAP[self.arch].get(op_str.upper(), -1)
1031 |                     if register_id != -1:
1032 |                         target_address = uc.reg_read(register_id)
1033 | 
1034 |                 else: # memory indirect call
1035 |                     pass # TODO: support memory indirect call calculation
1036 |                     # reg_or_mem = op_str.strip('[]')
1037 |                     # mem_addr = uc.reg_read(self.get_register_id(reg_or_mem))
1038 |                     # target_address = int.from_bytes(uc.mem_read(mem_addr, uc.mode // 8), byteorder='little')
1039 | 
1040 |                 tte_log_dbg(f"Indirect call to 0x{target_address:X}")
1041 |             return target_address
1042 | 
1043 |         def cb_check_call_target_loaded_and_skip(uc, address, size, user_data):
1044 |             insn_bytes = uc.mem_read(address, size)
1045 | 
1046 |             target_address = -1
1047 |             if self.unicorn_arch == UC_ARCH_X86:
1048 |                 if insn_bytes[0] not in [0xE8, 0xFF]: # Not a call instruction
1049 |                     return False
1050 |                 target_address = get_x86_call_target(uc, insn_bytes, address, size)
1051 | 
1052 |             if target_address == -1:
1053 |                 return
1054 | 
1055 |             elif not ida_bytes.is_loaded(target_address):
1056 |                 tte_log_dbg(f"Hook callback: Skip unloaded call target: 0x{target_address:X}")
1057 |                 current_insn_address = uc.reg_read(ARCH_TO_INSN_POINTER_MAP[self.arch])
1058 |                 uc.reg_write(ARCH_TO_INSN_POINTER_MAP[self.arch], current_insn_address + size)
1059 | 
1060 |         self.add_mu_hook( UC_HOOK_CODE, cb_check_call_target_loaded_and_skip)
1061 | 
1062 | 
1063 |     def _set_regs_init_value(self) -> None:
1064 |         if not self.settings.is_load_registers and self.settings.is_set_stack_value and self.unicorn_arch == UC_ARCH_X86: # Set default stack value in x86 mode
1065 |             offset = 0
1066 |             while(self._is_memory_mapped(DEFAULT_STACK_POINT_VALUE + offset)):
1067 |                 offset += 0x1000000
1068 |             if self.unicorn_mode == UC_MODE_64:
1069 |                 stack_point, frame_point  = UNICORN_REGISTERS_MAP[self.arch]['RSP'], UNICORN_REGISTERS_MAP[self.arch]['RBP']
1070 |             elif self.unicorn_mode == UC_MODE_32:
1071 |                 stack_point, frame_point  = UNICORN_REGISTERS_MAP[self.arch]['ESP'], UNICORN_REGISTERS_MAP[self.arch]['EBP']
1072 |             else:
1073 |                 raise AssertionError("Invalid unicorn mode")
1074 | 
1075 |             self.mu.reg_write(stack_point, DEFAULT_STACK_POINT_VALUE + offset)
1076 |             self.mu.reg_write(frame_point, DEFAULT_BASE_POINT_VALUE + offset)
1077 |             tte_log_info(f"Init regs: Sets registers default stack regs value: sp = {DEFAULT_STACK_POINT_VALUE + offset:X}, bp = {DEFAULT_BASE_POINT_VALUE + offset:X}")
1078 | 
1079 |         elif self.settings.is_load_registers and self.unicorn_arch == UC_ARCH_X86 and self.unicorn_mode == UC_MODE_64:
1080 |             for reg_id, reg_value in get_arch_x64_regs_value().items():
1081 |                 if reg_value is not None:
1082 |                     self.mu.reg_write(reg_id, reg_value)
1083 |                     tte_log_info(f"Init regs: Sets register ID {reg_id} value: 0x{reg_value:X}")
1084 |                 else:
1085 |                     tte_log_info(f"Init regs: Skip register ID {reg_id} value: None")
1086 | 
1087 |         elif self.settings.is_load_registers and self.unicorn_arch == UC_ARCH_X86 and self.unicorn_mode == UC_MODE_32:
1088 |             for reg_id, reg_value in get_arch_x86_regs_value().items():
1089 |                 if reg_value is not None:
1090 |                     self.mu.reg_write(reg_id, reg_value)
1091 |                     tte_log_info(f"Init regs: Sets register ID {reg_id} value: 0x{reg_value:X}")
1092 |                 else:
1093 |                     tte_log_info(f"Init regs: Skip register ID {reg_id} value: None")
1094 | 
1095 | 
1096 |     CUSTOM_HOOK_MEM_MAP = 0
1097 |     CUSTOM_HOOK_EXECUTE_END = 1
1098 |     def add_custom_hook(self, htype: int, callback):
1099 |         if htype == EmuExecutor.CUSTOM_HOOK_MEM_MAP:
1100 |             self.emu_map_mem_callback.append(callback)
1101 |         elif htype == EmuExecutor.CUSTOM_HOOK_EXECUTE_END:
1102 |             self.emu_run_end_callback.append(callback)
1103 |         else:
1104 |             raise AssertionError("Invalid hook type")
1105 | 
1106 |     def get_mem_regions(self) -> Iterator[Tuple[int, int, int]]:
1107 |         return self.mu.mem_regions()
1108 | 
1109 |     def run(self) -> None:
1110 |         """
1111 |         Run the emulation from the specified address until the specified address.
1112 |         """
1113 |         tte_log_info(f"Run emulation: 0x{self.settings.start:X}~0x{self.settings.end:X}, count {self.settings.count}")
1114 | 
1115 |         try:
1116 |             self.mu.emu_start(self.settings.start, self.settings.end, 0, self.settings.count)
1117 |         except UcError as e:
1118 |             tte_log_err(f"Emulation failed: {e}")
1119 | 
1120 |         tte_log_info(f"Emulation ended: 0x{self.mu.reg_read(ARCH_TO_INSN_POINTER_MAP[self.arch]):X}")
1121 |         tte_log_dbg("Emulation: Start calling end callback functions.")
1122 |         for callback in self.emu_run_end_callback:
1123 |             callback(self.mu)
1124 | 
1125 |         tte_log_info("Emulation completed successfully.")
1126 | 
1127 |     def destroy(self) -> None:
1128 |         """
1129 |         End the emulate
1130 |         """
1131 |         if not self._is_initialized:
1132 |             tte_log_info("EmuExecutor not initialized.")
1133 |             return
1134 | 
1135 |         self.mu.emu_stop()
1136 |         for handler in self.mu_hook_handlers:
1137 |             self.mu.hook_del(handler)
1138 | 
1139 |         self.emu_run_end_callback.clear()
1140 |         self.emu_map_mem_callback.clear()
1141 |         self.loaded_pages.clear()
1142 | 
1143 |         self._is_initialized = False
1144 | 
1145 | 
1146 | StateList = List[Tuple[str, "EmuState"]]
1147 | RegistersDict = Dict[str, int] # {reg_name: reg_value}
1148 | RegistersPatchDict = Dict[str, int] # {reg_name: patch_reg_value}
1149 | MemoryPatchList= List[Tuple[int, int, bytes]] # [(page_start, page_size, patch_bytes)]
1150 | MemoryPagesDict = Dict[int, Tuple[int, bytearray]] # {page_start: (page_permissions, page_data)}
1151 | MemoryBsdiffPatchDict = Dict[int, Tuple[int, bytes]] # {page_start: (page_permissions, bsdiff_patch_bytes)}
1152 | 
1153 | RegistersDiffsDict = Dict[str, Tuple[int, int]] # {reg_name: (reg_value, patch_value)}
1154 | 
1155 | 
1156 | 
1157 | class EmuState(ABC):
1158 |     """
1159 |     Class that saves Unicorn execution status
1160 |     Saved status includes:
1161 |     1. Register value(or Register patch)
1162 |     2. Memory page(or Memory patch)
1163 |     3. Memory page comparison
1164 |     4. Command address
1165 |     5. Number of executions
1166 |     """
1167 | 
1168 |     # State type constant, Represents a specific state in Unicorn simulation execution.
1169 |     # The full state or difference from the previous state can be stored.
1170 |     STATE_TYPE_FULL  = 1  # Full status
1171 |     STATE_TYPE_LIGHT_PATCH = 2 # Light Patch Status
1172 |     STATE_TYPE_HEAVY_PATCH = 3 # Heavy Patch Status
1173 | 
1174 |     def __init__(self,
1175 |                  state_id: str,
1176 |                  prev_state_id: str = "",
1177 |                  instruction: bytes = b"",
1178 |                  instruction_address: int = -1,
1179 |                  execution_count: int = -1) -> None:
1180 |         self.state_id = state_id
1181 |         self.prev_state_id = prev_state_id
1182 |         self.type = None
1183 | 
1184 |         self.instruction_address = instruction_address
1185 |         self.instruction = instruction
1186 |         self.execution_count = execution_count
1187 | 
1188 |         self.memory_patches: Optional[MemoryPatchList] = []
1189 | 
1190 |     @abstractmethod
1191 |     def generate_full_state(self, states_dict: Dict[str, 'EmuState']) -> Optional['FullEmuState']:
1192 |         pass
1193 | 
1194 | 
1195 | class FullEmuState(EmuState):
1196 |     def __init__(self,
1197 |                  state_id: str,
1198 |                  prev_state_id: str = "",
1199 |                  instruction: bytes = b"",
1200 |                  instruction_address: int = -1,
1201 |                  execution_count: int = -1) -> None:
1202 |         super().__init__(state_id,
1203 |                          prev_state_id,
1204 |                          instruction,
1205 |                          instruction_address,
1206 |                          execution_count)
1207 |         self.type = EmuState.STATE_TYPE_FULL
1208 | 
1209 |         self.registers_map: RegistersDict = {} # {reg_name: reg_value}
1210 |         self.memory_patches: Optional[MemoryPatchList] = None # [(page_start, page_size, patch_bytes)]
1211 |         self.memory_pages: MemoryPagesDict = {} # {page_start: (page_permissions, page_data)}
1212 | 
1213 | 
1214 |     def set_data(self,
1215 |                  registers_map: RegistersDict,
1216 |                  memory_patches: Optional[MemoryPatchList],
1217 |                  memory_pages: MemoryPagesDict) -> None:
1218 |         """
1219 |         Set the full state of the emulator.
1220 |         :param registers: Dictionary of register values.
1221 |         :param memory_pages: Dictionary of memory pages.
1222 |         """
1223 |         self.registers_map = registers_map
1224 |         self.memory_patches = memory_patches # As a prompt in viewer only, it will not be used for memory patching
1225 |         self.memory_pages = memory_pages
1226 | 
1227 |     def generate_full_state(self, states_dict: Dict[str, EmuState]) -> Optional['FullEmuState']:
1228 |         return deepcopy(self)
1229 | 
1230 | 
1231 | class HeavyPatchEmuState(EmuState):
1232 |     def __init__(self,
1233 |                  state_id: str,
1234 |                  prev_state_id: str = "",
1235 |                  instruction: bytes = b"",
1236 |                  instruction_address: int = -1,
1237 |                  execution_count: int = -1) -> None:
1238 |         super().__init__(state_id,
1239 |                          prev_state_id,
1240 |                          instruction,
1241 |                          instruction_address,
1242 |                          execution_count)
1243 |         self.type = EmuState.STATE_TYPE_HEAVY_PATCH
1244 | 
1245 |         self.base_full_state_id: Optional[str] = None
1246 | 
1247 |         self.reg_patches: RegistersPatchDict = {} # {reg_name: patch_value}
1248 |         self.memory_patches: Optional[MemoryPatchList] = None # [(page_start, page_size, patch_bytes)]
1249 | 
1250 |         # Memory patch stores binary differential data generated by bsdiff4
1251 |         self.mem_bsdiff_patches: MemoryBsdiffPatchDict = {} # {page_start: (page_permissions, bsdiff_patch_bytes)}
1252 |         self.new_pages: MemoryPagesDict = {} # {page_start: (page_permissions, page_data)}
1253 | 
1254 |     def set_data(self,
1255 |                  base_full_state_id: str,
1256 |                  reg_patches: RegistersPatchDict,
1257 |                  memory_patches: Optional[MemoryPatchList],
1258 |                  mem_bsdiff_patches: MemoryBsdiffPatchDict,
1259 |                  new_pages: MemoryPagesDict):
1260 |         """
1261 |         Sets the patch status of the emulator.
1262 | 
1263 |         :param reg_patches: Dictionary of register patches.
1264 |         :param mem_bsdiff_patches: Dictionary of memory patches, need to be applied by bsdiff4.
1265 |         """
1266 |         self.base_full_state_id = base_full_state_id
1267 | 
1268 |         self.reg_patches = reg_patches
1269 |         self.memory_patches = memory_patches # As a prompt in viewer only, it will not be used for memory patching
1270 |         self.mem_bsdiff_patches = mem_bsdiff_patches
1271 |         self.new_pages = new_pages
1272 | 
1273 |     def generate_full_state(self, states_dict: Dict[str, EmuState]) -> Optional[FullEmuState]:
1274 |         tte_log_dbg(f"Generate full state for heavy path state {self.state_id}")
1275 |         target_state = FullEmuState(self.state_id, self.prev_state_id,self.instruction, self.instruction_address, self.execution_count)
1276 | 
1277 |         assert self.base_full_state_id is not None, "Generate full State: Cannot apply memory patch: base state id not set."
1278 |         assert self.base_full_state_id in states_dict, f"Generate full State: Cannot apply memory patch: base state {self.base_full_state_id} not found."
1279 |         base_full_state: Optional[EmuState] = states_dict[self.base_full_state_id].generate_full_state(states_dict)
1280 | 
1281 |         assert base_full_state is not None, f"Generate full State: Cannot get full state for {self.base_full_state_id}: base state not found."
1282 |         assert isinstance(base_full_state, FullEmuState), f"Generate full State: Cannot apply memory patch: base state is not a full state."
1283 | 
1284 |         target_registers_map = apply_dict_patch(base_full_state.registers_map, self.reg_patches)
1285 |         target_memory_pages = apply_bytes_patch(base_full_state.memory_pages, self.mem_bsdiff_patches, self.new_pages)
1286 |         target_state.set_data(target_registers_map, None,  target_memory_pages)
1287 | 
1288 |         return deepcopy(target_state)
1289 | 
1290 | 
1291 | class LightPatchEmuState(EmuState):
1292 |     def __init__(self,
1293 |                  state_id: str,
1294 |                  prev_state_id: str = "",
1295 |                  instruction: bytes = b"",
1296 |                  instruction_address: int = -1,
1297 |                  execution_count: int = -1) -> None:
1298 |         super().__init__(state_id,
1299 |                          prev_state_id,
1300 |                          instruction,
1301 |                          instruction_address,
1302 |                          execution_count)
1303 |         self.type = EmuState.STATE_TYPE_LIGHT_PATCH
1304 | 
1305 |         self.base_full_state_id: Optional[str] = None
1306 | 
1307 |         self.reg_patches: RegistersPatchDict = {} # {reg_name: patch_value}
1308 |         self.memory_patches: MemoryPatchList = [] # List[[address, size, value]...]
1309 | 
1310 | 
1311 |     def set_data(self,
1312 |                  base_full_state_id: str,
1313 |                  reg_patches: RegistersPatchDict,
1314 |                  memory_patches: MemoryPatchList):
1315 |         """
1316 |         Set the patch state of the emulator.
1317 |         :param reg_patches: Dictionary of register patches.
1318 |         :param memory_patches: Dictionary of memory patches.
1319 |         """
1320 |         self.base_full_state_id = base_full_state_id
1321 | 
1322 |         self.reg_patches = reg_patches
1323 |         self.memory_patches = memory_patches
1324 |         tte_log_dbg(message=f"state:{self.state_id}, memory_patches {self.memory_patches}")
1325 | 
1326 | 
1327 |     def generate_full_state(self, states_dict: Dict[str, EmuState]) -> Optional[FullEmuState]:
1328 |         tte_log_dbg(f"Generate full state for light path state {self.state_id}")
1329 |         target_state = FullEmuState(self.state_id,
1330 |                                     self.prev_state_id,
1331 |                                     self.instruction,
1332 |                                     self.instruction_address,
1333 |                                     self.execution_count)
1334 |         accumulated_memory_patches_to_apply = self.memory_patches.copy()
1335 |         assert self.base_full_state_id is not None, "Generate full State: Cannot generate full state: base state id not set."
1336 |         base_full_state: Optional[EmuState] = states_dict.get(self.base_full_state_id)
1337 |         assert base_full_state is not None, f"Generate full State: Cannot generate full state for {self.base_full_state_id}: base state not found."
1338 |         assert isinstance(base_full_state, FullEmuState), f"Generate full State: Cannot generate full state for {self.base_full_state_id}: base state is not a full state."
1339 | 
1340 |         prev_full_state: Optional[EmuState] = None
1341 |         prev_state_id = self.prev_state_id
1342 |         while prev_state_id is not None:
1343 |             prev_state = states_dict.get(prev_state_id)
1344 |             if prev_state is None:
1345 |                 tte_log_warn(f"Generate full State: Cannot generate full state for {prev_state_id}: previous state not found.")
1346 |                 return None
1347 | 
1348 |             if prev_state.type in [EmuState.STATE_TYPE_FULL, EmuState.STATE_TYPE_HEAVY_PATCH]:
1349 |                 prev_full_state = prev_state.generate_full_state(states_dict)
1350 |                 assert prev_full_state is not None, "Generate full State: Cannot generate full state for previous state."
1351 |                 assert isinstance(prev_full_state, FullEmuState), "Generate full State: Previous state is not a full state."
1352 |                 break;
1353 | 
1354 |             elif prev_state.type == EmuState.STATE_TYPE_LIGHT_PATCH:
1355 |                 assert isinstance(prev_state, LightPatchEmuState), "Generate full State: Previous state is not a light patch state."
1356 |                 accumulated_memory_patches_to_apply.extend(prev_state.memory_patches)
1357 |                 prev_state_id = prev_state.prev_state_id
1358 | 
1359 |         # Apply register patches
1360 |         target_state.registers_map = apply_dict_patch(base_full_state.registers_map, self.reg_patches)
1361 | 
1362 |         # Apply memory patches
1363 |         assert prev_full_state is not None, "Generate full State: Cannot apply memory patch: previous state not found."
1364 |         target_state.memory_pages = prev_full_state.memory_pages
1365 |         while accumulated_memory_patches_to_apply:
1366 |             addr, size, value = accumulated_memory_patches_to_apply.pop()
1367 | 
1368 |             align_addr = addr & PAGE_MASK
1369 |             offset = addr & ~PAGE_MASK
1370 | 
1371 |             assert align_addr in target_state.memory_pages, "Generate full State: Cannot apply memory patch: address not found in memory pages."
1372 |             _, page_data = target_state.memory_pages[align_addr]
1373 |             page_data[offset : offset + size] = value
1374 |         return target_state
1375 | 
1376 | 
1377 | class EmuStateManager():
1378 |     """
1379 |     Class responsible for storing and managing all EmuState objects.
1380 |     It handles the creation, extraction, serialization, and deserialization of states.
1381 |     """
1382 | 
1383 |     MAX_PATCH_CHAIN_LENGTH = 100 # The longest patch chain, if it exceeds the full snapshot
1384 |     MAX_HEAVY_PATCH_COUNT = 10
1385 |     MAX_CUMULATIVE_DIFF_THRESHOLD = 0x1000 # Accumulated difference threshold, if it exceeds it, create a full snapshot
1386 |     STATE_ID_FORMAT = "$0x{address:X}#{count}"
1387 | 
1388 |     def __init__(self) -> None:
1389 |         self.arch = get_arch()
1390 | 
1391 |         self.states_dict: Dict[str, EmuState] = {} # dict: {state_id: EmuState}
1392 |         # self.storage_path = storage_path
1393 |         # os.makedirs(self.storage_path, exist_ok=True)
1394 | 
1395 |         self.instruction_execution_counts: Dict[int, int] = defaultdict(int) # {address: count}
1396 |         self.last_state_id: str = "" # Record the status id of the previous created
1397 |         self.last_full_state_id: Optional[str] = None # Record the most recent full status id
1398 | 
1399 |         self.patch_chain_count: int = 0 # Number of statuses in the current patch chain
1400 |         self.heavy_patch_count: int = 0 # Number of heavy patches created
1401 |         self.cumulative_diff_size: int = 0 # Cumulative difference size (used to determine whether to create a complete snapshot)
1402 | 
1403 |         self.has_map_memory = False # Whether the uc instance has mapped new memory page or not
1404 | 
1405 | 
1406 |     def _generate_state_id(self, instruction_address: int) -> str:
1407 |         """
1408 |         Generate a unique state ID based on the instruction address and the number of executions.
1409 |         Format: "<hex_addr><count>"
1410 |         """
1411 |         self.instruction_execution_counts[instruction_address] +=1
1412 |         count = self.instruction_execution_counts[instruction_address]
1413 |         return self.STATE_ID_FORMAT.format(address=instruction_address, count=count)
1414 | 
1415 | 
1416 |     def _read_memory_pages(self, uc, memory_regions: Iterator[Tuple[int, int, int]]) -> MemoryPagesDict:
1417 |         """
1418 |         Read paging memory from Unicorn instance.
1419 | 
1420 |         :param uc: The Unicorn instance.
1421 |         :param memory_regions: A list of memory regions to read.  Iterator(Tuple[start, end, permission])
1422 |         :return: A dictionary of memory pages, where the key is the start address of
1423 |          the page and the value is the page's permission and data.
1424 |         """
1425 |         memory_pages: MemoryPagesDict = {}
1426 |         for start, end, permission in memory_regions:
1427 |             size = end - start + 1
1428 |             try:
1429 |                 # Try to read a memory page
1430 |                 page_data = uc.mem_read(start, size)
1431 |                 memory_pages[start] = (permission, page_data)
1432 |                 tte_log_dbg(f"State Manager: Read memory page at 0x{start:X}, size {size}")
1433 |             except Exception as e:
1434 |                 # If memory area may not be mapped, skipped and log
1435 |                 tte_log_warn(f"Error: Could not read memory page at 0x{start:X}: {e}")
1436 |                 pass
1437 |         return memory_pages
1438 | 
1439 | 
1440 |     def _create_full_state(self,
1441 |                            new_state_id: str,
1442 |                            instruction: bytes,
1443 |                            instruction_address:int,
1444 |                            current_registers_map: RegistersDict,
1445 |                            memory_patches: MemoryPatchList,
1446 |                            current_memory_pages: MemoryPagesDict) -> None:
1447 |         new_state = FullEmuState(new_state_id,
1448 |                                  self.last_state_id,
1449 |                                  instruction,
1450 |                                  instruction_address,
1451 |                                  self.instruction_execution_counts[instruction_address])
1452 |         new_state.set_data(current_registers_map, memory_patches, current_memory_pages)
1453 | 
1454 |         self.states_dict[new_state_id] = new_state
1455 | 
1456 |         self.last_full_state_id = new_state_id
1457 |         self.last_state_id = new_state_id # Update the last status id
1458 |         tte_log_dbg(f"State Manager: Created FULL state: {new_state_id}")
1459 | 
1460 | 
1461 |     def _create_heavy_patch_state(self,
1462 |                                   new_state_id: str,
1463 |                                   instruction: bytes,
1464 |                                   instruction_address:int,
1465 |                                   current_registers_map: RegistersDict,
1466 |                                   memory_patches: MemoryPatchList,
1467 |                                   current_memory_pages: Dict[int, Tuple[int,bytearray]]) -> None:
1468 |         assert self.last_full_state_id is not None, "No full base state available for patch creation."
1469 |         assert self.last_state_id is not None, "No previous state available for patch creation."
1470 | 
1471 |         base_full_state: Optional[FullEmuState] = self.get_state(self.last_full_state_id) # type: ignore
1472 |         assert base_full_state is not None, "No full base state available for patch creation."
1473 |         assert isinstance(base_full_state, FullEmuState), "Base state is not a full state."
1474 | 
1475 |         reg_patches = catch_dict_patch(base_full_state.registers_map, current_registers_map)
1476 |         mem_bsdiff_patches,new_pages = catch_bytes_patch(base_full_state.memory_pages, current_memory_pages) # [ ] TODO support unmap pages in newer state
1477 | 
1478 |         new_state = HeavyPatchEmuState(new_state_id,
1479 |                                        self.last_state_id,
1480 |                                        instruction,
1481 |                                        instruction_address,
1482 |                                        self.instruction_execution_counts[instruction_address])
1483 |         new_state.set_data(self.last_full_state_id, reg_patches, memory_patches, mem_bsdiff_patches, new_pages)
1484 | 
1485 |         self.states_dict[new_state_id] = new_state
1486 | 
1487 |         self.patch_chain_count += 1
1488 |         self.heavy_patch_count += 1
1489 |         self.cumulative_diff_size += len(mem_bsdiff_patches)
1490 |         self.last_state_id = new_state_id # Update the last status id
1491 |         tte_log_dbg(f"State Manager: Created HEAVY PATCH state: {new_state_id}, base: {new_state.base_full_state_id}")
1492 | 
1493 | 
1494 | 
1495 |     def _create_light_patch_state(self,
1496 |                                   new_state_id: str,
1497 |                                   instruction: bytes,
1498 |                                   instruction_address:int,
1499 |                                   current_registers_map: RegistersDict,
1500 |                                   memory_patches: MemoryPatchList) -> None:
1501 |         assert self.last_full_state_id is not None, "No full base state available for patch creation."
1502 |         assert self.last_state_id is not None, "No previous state available for patch creation."
1503 | 
1504 |         base_full_state: Optional[FullEmuState] = self.get_state(self.last_full_state_id) # type: ignore
1505 |         assert base_full_state is not None, "No full base state available for patch creation."
1506 |         assert isinstance(base_full_state, FullEmuState), "Base state is not a full state."
1507 | 
1508 |         reg_patches = catch_dict_patch(base_full_state.registers_map, current_registers_map)
1509 | 
1510 |         new_state = LightPatchEmuState(new_state_id,
1511 |                                        self.last_state_id,
1512 |                                        instruction,
1513 |                                        instruction_address,
1514 |                                        self.instruction_execution_counts[instruction_address])
1515 |         new_state.set_data(self.last_full_state_id, reg_patches, memory_patches)
1516 | 
1517 |         self.states_dict[new_state_id] = new_state
1518 | 
1519 |         self.patch_chain_count += 1
1520 |         self.cumulative_diff_size += len(memory_patches)
1521 |         self.last_state_id = new_state_id # Update the last status id
1522 |         tte_log_dbg(f"State Manager: Created LIGHT PATCH state: {new_state_id}, base:{new_state.base_full_state_id},  prev: {new_state.prev_state_id}")
1523 | 
1524 | 
1525 |     def get_state(self, state_id: Optional[str]) -> Optional[EmuState]:
1526 |         """
1527 |         Get the EmuState object by its ID.
1528 |         :param state_id: The ID of the state to retrieve.
1529 |         :return: The EmuState object if found, otherwise None.
1530 |         """
1531 |         if state_id is None:
1532 |             return None
1533 |         return self.states_dict.get(state_id, None)
1534 | 
1535 | 
1536 |     def get_state_list(self) -> StateList:
1537 |         """
1538 |         Get a list of all EmuState objects with their ID, instruction address, instruction, and execution count.
1539 | 
1540 |         :return: A list of tuples containing the state ID, instruction address, instruction, and execution count.
1541 |         """
1542 |         result: StateList = []
1543 |         for state_id, state in self.states_dict.items():
1544 |             result.append((state_id, state))
1545 |         return result
1546 | 
1547 |     def _determine_next_state_type(self):
1548 |         tte_log_dbg(f"Determine next state type: heavy patch count: {self.heavy_patch_count}, cumulative diff size: {self.cumulative_diff_size}, has_map_memory: {self.has_map_memory}")
1549 | 
1550 |         if self.last_full_state_id is None or self.heavy_patch_count > self.MAX_HEAVY_PATCH_COUNT or \
1551 |             self.patch_chain_count >= self.MAX_PATCH_CHAIN_LENGTH or self.cumulative_diff_size >= self.MAX_CUMULATIVE_DIFF_THRESHOLD :
1552 |             self.has_map_memory = False
1553 |             self.heavy_patch_count = 0
1554 |             self.patch_chain_count = 0
1555 |             self.cumulative_diff_size = 0
1556 | 
1557 |             tte_log_dbg("Create full state")
1558 |             return EmuState.STATE_TYPE_FULL
1559 | 
1560 |         elif self.has_map_memory == True:
1561 |             self.has_map_memory = False
1562 |             tte_log_dbg("Create heavy patch state")
1563 |             return EmuState.STATE_TYPE_HEAVY_PATCH
1564 | 
1565 |         tte_log_dbg("Create light state")
1566 |         return EmuState.STATE_TYPE_LIGHT_PATCH
1567 | 
1568 | 
1569 |     def create_state(self,
1570 |                      uc,
1571 |                      instruction: bytes,
1572 |                      instruction_address: int,
1573 |                      memory_regions: Iterator[Tuple[int, int, int]],
1574 |                      memory_patches: MemoryPatchList) -> None:
1575 |         """
1576 |         Create a new EmuState object based on the current state of the Unicorn instance.
1577 |         Decide whether to create a full state or a patch state based on the difference size or number of steps.
1578 | 
1579 |         :param uc: The Unicorn instance.
1580 |         :param instruction_address: The address of the current instruction.
1581 |         :param memory_regions: A list of memory regions(Iterator[Tuple[start, end, permission]]) to read.
1582 |         :param memory_patches: A list of memory patches(List[Tuple[addr, size, value]]) to apply.
1583 |         """
1584 |         tte_log_dbg(f"State Manager: Create state at 0x{instruction_address:X}")
1585 |         new_state_id = self._generate_state_id(instruction_address)
1586 | 
1587 |         # Read the registers of unicorn instance
1588 |         current_registers_map = {
1589 |             reg_name: uc.reg_read(reg_const)
1590 |             for reg_name,reg_const in UNICORN_REGISTERS_MAP[self.arch].items()
1591 |         }
1592 | 
1593 |         next_state_type = self._determine_next_state_type()
1594 |         if next_state_type == EmuState.STATE_TYPE_FULL:
1595 |             # Create a full status
1596 |             current_memory_pages = self._read_memory_pages(uc, memory_regions)
1597 |             self._create_full_state(new_state_id,
1598 |                                     instruction,
1599 |                                     instruction_address,
1600 |                                     current_registers_map,
1601 |                                     memory_patches,
1602 |                                     current_memory_pages)
1603 | 
1604 |         elif next_state_type == EmuState.STATE_TYPE_HEAVY_PATCH:
1605 |             # Create a heavy patch status
1606 |             current_memory_pages = self._read_memory_pages(uc, memory_regions)
1607 |             self._create_heavy_patch_state(new_state_id,
1608 |                                            instruction,
1609 |                                            instruction_address,
1610 |                                            current_registers_map,
1611 |                                            memory_patches,
1612 |                                            current_memory_pages)
1613 | 
1614 |         else:
1615 |             # Create a light patch status
1616 |             self._create_light_patch_state(new_state_id,
1617 |                                            instruction,
1618 |                                            instruction_address,
1619 |                                            current_registers_map,
1620 |                                            memory_patches)
1621 | 
1622 | 
1623 |     def _get_regs_map(self, state) -> RegistersDict:
1624 |         if state.type in [EmuState.STATE_TYPE_LIGHT_PATCH, EmuState.STATE_TYPE_HEAVY_PATCH]:
1625 |             base_state = self.get_state(state.base_full_state_id)
1626 |             assert base_state, f"Base state not found"
1627 |             assert isinstance(base_state, FullEmuState)
1628 |             return apply_dict_patch(base_state.registers_map, state.reg_patches)
1629 |         else:
1630 |             assert isinstance(state, FullEmuState)
1631 |             return state.registers_map
1632 | 
1633 | 
1634 |     def compare_states(self, state1: EmuState, state2: EmuState):
1635 |         """
1636 |         Compares any two EmuState objects (full or patch) and returns their differences.
1637 |         The first state must be the base state, and the second state is the target state.
1638 | 
1639 |         :param EmuStateManager: An instance of EmuStateManager to extract full states if needed.
1640 |         :param state1: The first EmuState object.
1641 |         :param state2: The second EmuState object.
1642 |         :return: A tuple (regs_diff, mem_diff, pages_diff) representing the differences.
1643 |                  regs_diff: Dict of changed registers in state2 relative to state1.
1644 |                  mem_diff: List of memory different from state1 to state2.
1645 |                  pages_diff: List of page different from state1 to state2, and the type of the page (1: added, 2: removed)
1646 |         """
1647 |         tte_log_dbg(f"\n--- Comparing states Fully: {state1.state_id} vs {state2.state_id} ---")
1648 | 
1649 |         tte_log_dbg(f"Comparison between State '{state1.state_id}' (Type: {state1.type}) and State '{state2.state_id}' (Type: {state2.type}):")
1650 |         tte_log_dbg(f"Instruction Address: 0x{state1.instruction_address:X} -> 0x{state2.instruction_address:X}")
1651 |         tte_log_dbg(f"Execution Count: {state1.execution_count} -> {state2.execution_count}")
1652 | 
1653 |         full_state1 = state1.generate_full_state(self.states_dict)
1654 |         full_state2 = state2.generate_full_state(self.states_dict)
1655 | 
1656 |         if not full_state1:
1657 |             tte_log_warn(f"Could not generate full state for '{state1.state_id}'. Comparison aborted.")
1658 |             return
1659 |         if not full_state2:
1660 |             tte_log_warn(f"Could not generate full state for '{state2.state_id}'. Comparison aborted.")
1661 |             return
1662 | 
1663 |         regs_diff: RegistersDiffsDict = {}
1664 |         def catch_regs_diff(regs_map1, regs_map2):
1665 |             result: RegistersDiffsDict = {}
1666 |             for reg_name in regs_map1:
1667 |                 if regs_map1[reg_name] != regs_map2[reg_name]:
1668 |                     result[reg_name] = (regs_map1[reg_name], regs_map2[reg_name])
1669 |                     tte_log_dbg(f"  {reg_name}: 0x{regs_map1[reg_name]:X} -> 0x{regs_map2[reg_name]:X}")
1670 |             return result
1671 |         regs_diff = catch_regs_diff(full_state1.registers_map, full_state2.registers_map)
1672 | 
1673 |         mem_diff = SortedDict()
1674 |         def catch_mem_diff(page_start_addr, bytes1: bytearray, bytes2: bytearray):
1675 |             assert len(bytes1) == len(bytes2), "Byte arrays must be of equal length."
1676 |             for offset_addr, (b1, b2) in enumerate(zip(bytes1, bytes2)):
1677 |                 if b1 != b2:
1678 |                     tte_log_dbg(f"Memory difference found at 0x{page_start_addr + offset_addr:X}: 0x{b1:X} -> 0x{b2:X}")
1679 |                     mem_diff[page_start_addr + offset_addr] = (b1, b2)
1680 | 
1681 |         for page_start_addr in full_state1.memory_pages.keys() & full_state2.memory_pages.keys():
1682 |             catch_mem_diff(
1683 |                 page_start_addr,
1684 |                 full_state1.memory_pages[page_start_addr][1],
1685 |                 full_state2.memory_pages[page_start_addr][1]
1686 |             )
1687 | 
1688 |         pages_diff = SortedDict()
1689 |         diff_keys_in_state1 = set(full_state1.memory_pages.keys()) - set(full_state2.memory_pages.keys())
1690 |         diff_keys_in_state2 = set(full_state2.memory_pages.keys()) - set(full_state1.memory_pages.keys())
1691 |         for key in diff_keys_in_state1:
1692 |             pages_diff[key] = (1, full_state1.memory_pages[key])
1693 |         for key in diff_keys_in_state2:
1694 |             pages_diff[key] = (2, full_state2.memory_pages[key])
1695 | 
1696 |         return regs_diff, mem_diff, pages_diff
1697 | 
1698 | 
1699 | 
1700 | 
1701 | 
1702 | 
1703 | 
1704 | class EmuTracer():
1705 |     """
1706 |     Class for tracking EmuExecutor and creating state snapshots in EmuStateManager when appropriate.
1707 |     """
1708 |     @dataclass
1709 |     class state_buffer_t(object):
1710 |         instruction: bytes
1711 |         memory_patches: MemoryPatchList
1712 | 
1713 |     def __init__(self, executor:EmuExecutor, state_manager:EmuStateManager) -> None:
1714 |         self.executor: EmuExecutor = executor
1715 |         self.state_manager: EmuStateManager = state_manager
1716 | 
1717 |         self.arch = self.executor.arch
1718 |         self.capstone_arch, self.capstone_mode= CAPSTONE_ARCH_MAP[self.arch]
1719 |         self.md = Cs(self.capstone_arch, self.capstone_mode)
1720 | 
1721 |         self.state_buffer = self.state_buffer_t(b"", [])
1722 | 
1723 | 
1724 |     def init_hook(self) -> None:
1725 |         self._trace_code()
1726 |         self._trace_mem_write()
1727 | 
1728 |         self._trace_custom_operate()
1729 | 
1730 | 
1731 |     def _trace_code(self):
1732 |         # self.executor.add_mu_hook(UC_HOOK_MEM_READ, self.trace_mem_read)
1733 |         self.executor.add_mu_hook(UC_HOOK_CODE, self.cb_catch_insn_execution)
1734 |         self.executor.add_mu_hook(UC_HOOK_CODE, self.cb_create_emu_state)
1735 | 
1736 | 
1737 |     def _trace_mem_write(self):
1738 |         self.executor.add_mu_hook(UC_HOOK_MEM_WRITE, self.cb_log_mem_write)
1739 |         self.executor.add_mu_hook(UC_HOOK_MEM_WRITE, self.cb_catch_mem_write)
1740 | 
1741 | 
1742 |     def _trace_custom_operate(self):
1743 |         self.executor.add_custom_hook(EmuExecutor.CUSTOM_HOOK_EXECUTE_END, self.cb_create_emu_end_state)
1744 |         self.executor.add_custom_hook(EmuExecutor.CUSTOM_HOOK_MEM_MAP, self.cb_mark_mem_map)
1745 | 
1746 | 
1747 | 
1748 |     def cb_catch_insn_execution(self, uc, address, size, user_data) -> bool:
1749 |         tte_log_dbg(message=f"Tracing instruction at 0x{address:X}, instruction size = {size:X}")
1750 |         opcode, operand = InstrctionParser().parse_instruction_to_tuple(uc.mem_read(address, size))
1751 |         tte_log_dbg("Executing instruction:    %s\t%s" % (opcode, operand))
1752 |         self.state_buffer.instruction = bytes(uc.mem_read(address, size))
1753 |         return True
1754 | 
1755 | 
1756 |     def cb_create_emu_state(self, uc, address, size, user_data) -> bool:
1757 |         """
1758 |         @Callback function for hook "UC_HOOK_CODE"
1759 |         Note: UC_HOOK_CODE triggers callback before the instruction is actually executed
1760 | 
1761 |         This function will perform two operations:
1762 |         Create the simulated state saved to the status buffer as a new EmuState and add it to the EmuStateManager
1763 |         Create a new simulated state in the state buffer
1764 |         """
1765 | 
1766 |         # Save the current state to the state buffer
1767 |         self.state_manager.create_state(uc,
1768 |                                         self.state_buffer.instruction,
1769 |                                         address,
1770 |                                         self.executor.get_mem_regions(),
1771 |                                         self.state_buffer.memory_patches)
1772 | 
1773 |         # Create a new state buffer
1774 |         self.state_buffer.memory_patches = []
1775 | 
1776 |         return True
1777 | 
1778 | 
1779 |     def cb_log_mem_write(self, uc, access, address, size, value, user_data):
1780 |         tte_log_dbg(f"Memory write: access={access}, address=0x{address:X}, size={size}, value=0x{value:X}, insn={uc.reg_read(ARCH_TO_INSN_POINTER_MAP[self.arch]):X}")
1781 |         return True
1782 | 
1783 | 
1784 |     def cb_catch_mem_write(self, uc, access, address:int, size:int, value:int, user_data) -> bool:
1785 |         """
1786 |         @Callback function for hook "UC_HOOK_MEM_WRITE"
1787 | 
1788 |         """
1789 |         tte_log_dbg(f"Catch memory write: address={hex(address)}, size={size}, value= {hex(value)}")
1790 |         if value >= 0:
1791 |             value_bytes = value.to_bytes(size, byteorder='big' if get_is_be() else 'little', signed=False)
1792 |         else:
1793 |             value_bytes = value.to_bytes(size, byteorder='big' if get_is_be() else 'little', signed=True)
1794 |         # Save the memory patch to the state buffer
1795 |         self.state_buffer.memory_patches.append((address, size, value_bytes))
1796 |         return True
1797 | 
1798 | 
1799 |     def cb_create_emu_end_state(self, uc):
1800 |         insn_pointer = uc.reg_read(ARCH_TO_INSN_POINTER_MAP[self.arch])
1801 |         self.state_manager.create_state(uc,
1802 |                                         b"",
1803 |                                         insn_pointer,
1804 |                                         self.executor.get_mem_regions(),
1805 |                                         self.state_buffer.memory_patches)
1806 | 
1807 |     def cb_mark_mem_map(self, uc):
1808 |         """
1809 |         Callback function to mark that the binary has been loaded.
1810 |         This is used to determine which type of state to create next time.
1811 |         """
1812 | 
1813 |         self.state_manager.has_map_memory = True
1814 |         tte_log_dbg("Binary has been loaded, will create heavy patch state next time.")
1815 | 
1816 | 
1817 | 
1818 | class InstrctionParser():
1819 |     """
1820 |     Class for parsing instructions and extracting relevant information.
1821 |     """
1822 |     _instance = None
1823 |     _initialized = False
1824 | 
1825 |     def __new__(cls, *args, **kwargs):
1826 |         if cls._instance is None:
1827 |             cls._instance = super(InstrctionParser, cls).__new__(cls)
1828 |         return cls._instance
1829 | 
1830 |     def __init__(self):
1831 |         if self._initialized:
1832 |             return
1833 |         self._initialized = True
1834 | 
1835 |         self.arch = get_arch()
1836 |         self.capstone_arch, self.capstone_mode = CAPSTONE_ARCH_MAP[self.arch]
1837 |         self.md = Cs(self.capstone_arch, self.capstone_mode)
1838 | 
1839 |     def parse_instructions(self, insn_bytes: bytes):
1840 |         return self.md.disasm(insn_bytes, 0)
1841 | 
1842 |     def parse_instruction_to_str(self, insn_bytes: bytes) -> str:
1843 |         """
1844 |         Parses an instruction and returns a string of the form "opcode operand"
1845 |         """
1846 |         assert self.md is not None, "InstrctionParser not initialized"
1847 |         try:
1848 |             _, _, opcode, operand = next(self.md.disasm_lite(insn_bytes, 0))
1849 |             return f"{opcode}\t{operand}"
1850 |         except StopIteration:
1851 |             return ""
1852 | 
1853 |     def parse_instruction_to_tuple(self, insn_bytes: bytes) -> Tuple[str, str]:
1854 |         """
1855 |         Parses an instruction and returns a tuple of (opcode, operand)
1856 |         """
1857 |         assert self.md is not None, "InstrctionParser not initialized"
1858 |         try:
1859 |             _, _, opcode, operand = next(self.md.disasm_lite(insn_bytes, 0))
1860 |             return opcode,operand
1861 |         except StopIteration:
1862 |             return "", ""
1863 | 
1864 | 
1865 | 
1866 | 
1867 | class MenuActionHandler(ida_kernwin.action_handler_t):
1868 |     def __init__(self,
1869 |                  parent_title,
1870 |                  checkable,
1871 |                  action_name,
1872 |                  handler,
1873 |                  title,
1874 |                  shortcut):
1875 |         ida_kernwin.action_handler_t.__init__(self)
1876 |         self.parent_title = parent_title
1877 |         self.checkable: Optional[Callable[[], Optional[bool]]] = checkable
1878 |         self.action_name = action_name
1879 |         self.handler = handler
1880 |         self.title = title
1881 |         self.shortcut = shortcut
1882 |         self.register()
1883 | 
1884 |     def activate(self, ctx):
1885 |         return self.handler()
1886 | 
1887 |     def update(self, ctx):
1888 |         if ctx.widget_title == self.parent_title and self.checkable and self.checkable():
1889 |             return idaapi.AST_ENABLE
1890 |         return idaapi.AST_DISABLE
1891 | 
1892 |     def register(self):
1893 |         actname = self.action_name
1894 |         if ida_kernwin.unregister_action(actname):
1895 |             tte_log_dbg("Unregistered previously-registered action \"%s\"" % actname)
1896 | 
1897 |         desc = ida_kernwin.action_desc_t(actname, self.title , self, self.shortcut)
1898 |         if ida_kernwin.register_action(desc):
1899 |             tte_log_dbg("Registered action \"%s\"" % actname)
1900 | 
1901 |     def unregister(self):
1902 |         actname = self.action_name
1903 |         if ida_kernwin.unregister_action(actname):
1904 |             tte_log_dbg("Unregistered action \"%s\"" % actname)
1905 | 
1906 |     def get_action_name(self) -> str:
1907 |         return self.action_name
1908 | 
1909 | 
1910 | # Line types
1911 | UNKNOW_LINE = 0
1912 | DATA_LINE = 1
1913 | CODE_LINE = 2
1914 | NAME_LINE = 3
1915 | EMPTY_LINE = 4
1916 | SINGLE_DATA_LINE = 5
1917 | ERROR_LINE = 6
1918 | 
1919 | @dataclass
1920 | class address_line_info:
1921 |     address: int
1922 |     address_idx: int = 0
1923 |     type: int = UNKNOW_LINE
1924 |     value: str = ""
1925 |     fgcolor: Optional[int] = 0x0
1926 |     bgcolor: Optional[int] = 0xFFFFFF
1927 | 
1928 | 
1929 | class AddressAwareCustomViewer(ida_kernwin.simplecustviewer_t):
1930 |     """
1931 |     This class extends simplecustviewer_t to manage lines of text associated with binary addresses.
1932 |     It uses a SortedList to maintain lines sorted by their associated address.
1933 |     """
1934 | 
1935 |     def __init__(self):
1936 |         super().__init__()
1937 |         self._lines_data: SortedList = SortedList(key=lambda x: (x.address, x.address_idx)) # SortedList[line_info], using Tuple[address, address_idx] as key
1938 |         self._lines_data_buffer: List[address_line_info] = []
1939 | 
1940 |         self.need_rebuild = False # Flag to indicate if the viewer needs to be refreshed
1941 | 
1942 |         self.keydown_callback_list: List[Tuple[int, bool, Callable]] = [] # List[(ord_key, is_shift, callback_func)], Callback functions list called when corresponding keys is down
1943 |         self.dbl_click_callback: Optional[Callable] = None # Callback function called when double-clicking a line
1944 |         self.menu_action_handlers = []
1945 | 
1946 | 
1947 |     def _get_lineno_left_from_address(self, address):
1948 |         """
1949 |         Gets the line number for a given address using SortedList's bisect_left.
1950 |         Returns -1 if address not found.
1951 |         """
1952 |         idx = self._lines_data.bisect_left(address_line_info(address = address))
1953 |         if idx < len(self._lines_data) and self._lines_data[idx].address == address: # type: ignore
1954 |             return idx
1955 |         return -1
1956 | 
1957 | 
1958 |     def _get_lineno_right_from_address(self, address):
1959 |         """
1960 |         Gets the line number for a given address using SortedList's bisect_right.
1961 |         Returns -1 if address not found.
1962 |         """
1963 |         idx = self._lines_data.bisect_right(address_line_info(address = address))
1964 |         if idx > 0 and self._lines_data[idx - 1].address == address: # type: ignore
1965 |             return idx - 1
1966 |         return -1
1967 | 
1968 | 
1969 |     def _get_address_info_from_lineno(self, lineno) -> Optional[address_line_info]:
1970 |         """
1971 |         Gets the address for a given line number.
1972 |         Returns None if lineno is out of bounds.
1973 |         """
1974 |         if 0 <= lineno < len(self._lines_data):
1975 |             return self._lines_data[lineno] # type: ignore
1976 |         return None
1977 | 
1978 | 
1979 |     def _rebuild_viewer_content(self):
1980 |         """
1981 |         Internal helper to clear the underlying simplecustviewer_t (which is self)
1982 |         and repopulate it from the sorted _lines_data.
1983 |         Notes: It is a performance-consuming function, so call this function as few as possible
1984 |         """
1985 |         super().ClearLines() # Call base ClearLines
1986 |         if self._lines_data_buffer:
1987 |             self._lines_data.update(self._lines_data_buffer)
1988 |             self._lines_data_buffer.clear()
1989 |         for addr_info in self._lines_data:
1990 |             super().AddLine(addr_info.value, addr_info.fgcolor, addr_info.bgcolor) # Call base AddLine
1991 |         super().Refresh() # Refresh the viewer
1992 | 
1993 |     def CheckRebuild(self):
1994 |         if self.need_rebuild:
1995 |             self._rebuild_viewer_content()
1996 |             self.need_rebuild = False
1997 | 
1998 |     def Create(self, title):
1999 |         return super().Create(title)
2000 | 
2001 |     def OnClose(self):
2002 |        self.UnregisterAction()
2003 | 
2004 |     def UnregisterAction(self):
2005 |         for action_handler in self.menu_action_handlers:
2006 |             action_handler.unregister()
2007 |         self.menu_action_handlers.clear()
2008 | 
2009 |     def Show(self):
2010 |         self.CheckRebuild()
2011 |         return super().Show()
2012 | 
2013 |     def Refresh(self):
2014 |         self.CheckRebuild()
2015 |         return super().Refresh()
2016 | 
2017 |     def ClearLines(self):
2018 |         self._lines_data.clear() # Use SortedList's clear method
2019 |         super().ClearLines() # Call base ClearLines
2020 |         super().Refresh() # Refresh the viewer
2021 | 
2022 | 
2023 |     def AddLine(self, address, address_type, line, fgcolor=None, bgcolor=None, lazy = False):
2024 |         """
2025 |         Adds a colored line associated with a binary address to  the last line of viewer.
2026 |         Note: The incoming address parameter must be bigger than all existing addresses in the viewer
2027 | 
2028 |         :param lazy: If True, the viewer will not be refreshed after the deletion.
2029 |         :return: Boolean indicating success.
2030 |         """
2031 | 
2032 |         is_appending = True
2033 |         addr_idx = 0
2034 | 
2035 |         if self._lines_data:
2036 |             last_line_info: address_line_info = self._lines_data[-1] # type: ignore
2037 |             if address < last_line_info.address:
2038 |                 # The address is out of order and cannot be added directly. Forced to go to lazy mode
2039 |                 is_appending = False
2040 |                 tte_log_dbg(
2041 |                     f"AddLine: Incoming address {hex(address)} is out of order "
2042 |                     f"(less than last address {hex(last_line_info.address)}). "
2043 |                     f"Forcing lazy rebuild for correct sorting."
2044 |                 )
2045 |                 idx = self._get_lineno_right_from_address(address)
2046 |                 addr_idx = self._lines_data[idx].address_idx + 1 # type: ignore
2047 |                 lazy = True
2048 | 
2049 |             elif address == last_line_info.address:
2050 |                 # The address is the same, increment the address_idx
2051 |                 addr_idx = last_line_info.address_idx + 1
2052 | 
2053 |             else: # address > last_line_info.address
2054 |                 #Strictly append, new address, address_idx starts from 0
2055 |                 addr_idx = 0
2056 |         else:
2057 |             # The view is empty, the first element
2058 |             addr_idx = 0
2059 | 
2060 |         # Add the new or updated entry. SortedList.add() will place it correctly.
2061 |         new_line_info = address_line_info(address, addr_idx,address_type, line, fgcolor, bgcolor)
2062 |         if lazy:
2063 |             self.need_rebuild = True
2064 |             self._lines_data_buffer.append(new_line_info)
2065 |             return True
2066 |         else:
2067 |             self._lines_data.add(new_line_info)
2068 |             return super().AddLine(line, fgcolor, bgcolor) # Call base InsertLine
2069 | 
2070 | 
2071 |     def InsertLine(self, address, address_type, line, fgcolor=None, bgcolor=None):
2072 |         """
2073 |         Inserts a line at the position determined by the given address.
2074 |         This is equivalent to AddLine as lines are always kept sorted by address.
2075 | 
2076 |         :return: Boolean indicating success.
2077 |         """
2078 |         self.CheckRebuild()
2079 |         # Find if the address already exists
2080 |         idx = self._get_lineno_right_from_address(address)
2081 |         addr_idx = 0
2082 |         if idx != -1 and idx < len(self._lines_data) and self._lines_data[idx].address == address: # type: ignore
2083 |             # Address exists
2084 |             addr_idx = self._lines_data[idx].address_idx + 1 # type: ignore
2085 | 
2086 |         # Add the new or updated entry. SortedList.add() will place it correctly.
2087 |         new_line_info = address_line_info(address, addr_idx,address_type, line, fgcolor, bgcolor)
2088 |         self._lines_data.add(new_line_info)
2089 | 
2090 |         actual_lineno_inserted = self._lines_data.index(new_line_info)
2091 |         if actual_lineno_inserted == len(self._lines_data) - 1:
2092 |             return super().AddLine(line, fgcolor, bgcolor) # Call base InsertLine
2093 |         else:
2094 |             return super().InsertLine(actual_lineno_inserted, line, fgcolor, bgcolor) # Call base InsertLine
2095 | 
2096 | 
2097 |     def EditLine(self, address, address_idx, address_type, line, fgcolor=None, bgcolor=None):
2098 |         """
2099 |         Edits an existing line identified by its binary address.
2100 | 
2101 |         :return: Boolean indicating success.
2102 |         """
2103 |         self.CheckRebuild()
2104 |         target_line_info = address_line_info(address=address, address_idx=address_idx)
2105 |         lineno_to_edit = self._lines_data.bisect_left(target_line_info)
2106 |         if lineno_to_edit < len(self._lines_data):
2107 |             line_info: address_line_info = self._lines_data[lineno_to_edit] # type: ignore
2108 | 
2109 |             if line_info.address == address and \
2110 |             line_info.address_idx == address_idx:
2111 |                 line_info.value = line
2112 |                 line_info.type = address_type # Update type as well
2113 |                 line_info.fgcolor = fgcolor
2114 |                 line_info.bgcolor = bgcolor
2115 | 
2116 |                 return super().EditLine(lineno_to_edit, line, fgcolor, bgcolor) # Call base EditLine
2117 |             return False
2118 |         return False
2119 | 
2120 | 
2121 |     def EditLineColor(self, address, address_idx, fgcolor=None, bgcolor=None):
2122 |         """
2123 |         Edits an existing line's color identified by its binary address.
2124 | 
2125 |         :return: Boolean indicating success.
2126 |         """
2127 |         self.CheckRebuild()
2128 |         target_line_info = address_line_info(address=address, address_idx=address_idx)
2129 |         lineno_to_edit = self._lines_data.bisect_left(target_line_info)
2130 |         if lineno_to_edit < len(self._lines_data):
2131 |             line_info: address_line_info = self._lines_data[lineno_to_edit] # type: ignore
2132 |             if line_info.address == address and \
2133 |             line_info.address_idx == address_idx:
2134 |                 line_info.fgcolor = fgcolor
2135 |                 line_info.bgcolor = bgcolor
2136 |                 return super().EditLine(lineno_to_edit, line_info.value, fgcolor, bgcolor) # Call base EditLine
2137 |             return False
2138 |         return False
2139 | 
2140 | 
2141 |     def UpdateLine(self, address, address_idx, address_type, line, fgcolor=None, bgcolor=None):
2142 |         """
2143 |         Updates an existing line identified by its binary address.
2144 | 
2145 |         :return: Boolean indicating success.
2146 |         """
2147 |         self.CheckRebuild()
2148 |         target_line_info = address_line_info(address=address, address_idx=address_idx)
2149 |         lineno_to_edit = self._lines_data.bisect_left(target_line_info)
2150 |         if lineno_to_edit < len(self._lines_data):
2151 |             line_info: address_line_info = self._lines_data[lineno_to_edit] # type: ignore
2152 | 
2153 |             if line_info.address == address and \
2154 |                 line_info.address_idx == address_idx:
2155 |                 return self.EditLine(address, address_idx, address_type, line, fgcolor, bgcolor)
2156 |             else:
2157 |                 return False
2158 |         return False
2159 | 
2160 | 
2161 |     def UpdateLineRange(self, address, length, address_type, line, fgcolor=None, bgcolor=None):
2162 |         """
2163 |         Deletes an existing line range identified by its binary address range.
2164 |         and inserts a new line.
2165 | 
2166 |         :return: Boolean indicating success.
2167 |         """
2168 |         self.CheckRebuild()
2169 |         for offset in range(length):
2170 |             self.DelLine(address + offset, lazy=False)
2171 |         return self.InsertLine(address, address_type, line, fgcolor, bgcolor)
2172 | 
2173 | 
2174 |     def PatchLine(self, address, address_idx, offs, value):
2175 |         raise AssertionError("PatchLine is Forbidden")
2176 | 
2177 | 
2178 |     def DelLine(self, address, lazy = False):
2179 |         """
2180 |         Deletes an existing line identified by its binary address.
2181 | 
2182 |         :param address: The binary address of the line to delete.
2183 |         :param lazy: If True, the viewer will not be refreshed after the deletion.
2184 |         :return: Boolean indicating success.
2185 |         """
2186 |         self.CheckRebuild()
2187 | 
2188 |         lines_to_delete_indices = []
2189 | 
2190 |         start_idx = self._lines_data.bisect_left(address_line_info(address=address))
2191 |         current_idx = start_idx
2192 |         while current_idx < len(self._lines_data) and self._lines_data[current_idx].address == address: # type: ignore
2193 |             lines_to_delete_indices.append(current_idx)
2194 |             current_idx += 1
2195 | 
2196 |         if not lines_to_delete_indices:
2197 |             return False # No lines found for this address
2198 | 
2199 |         # Delete from _lines_data in reverse order to avoid index shifts
2200 |         for lineno in reversed(lines_to_delete_indices):
2201 |             del self._lines_data[lineno]
2202 |             if not lazy:
2203 |                 # If not lazy, we must also update the base viewer immediately
2204 |                 super().DelLine(lineno)
2205 |         if lazy:
2206 |             self.need_rebuild = True # If any lines were deleted, a rebuild is needed later
2207 |         return True
2208 | 
2209 | 
2210 |     def GetLine(self, address, address_idx):
2211 |         """
2212 |         Returns a line's content and colors identified by its binary address.
2213 | 
2214 |         :param address: The binary address (integer) of the line.
2215 |         :return: Returns a tuple (colored_line, fgcolor, bgcolor) or None if address not found.
2216 |         """
2217 |         self.CheckRebuild()
2218 |         target_line_info = address_line_info(address=address, address_idx=address_idx)
2219 |         lineno_to_get = self._lines_data.bisect_left(target_line_info)
2220 |         if lineno_to_get < len(self._lines_data):
2221 |             # Retrieve the full tuple, then extract line_str, fgcolor, bgcolor
2222 |             addr_info: address_line_info = self._lines_data[lineno_to_get] # type: ignore
2223 | 
2224 |             if addr_info.address == address and \
2225 |             addr_info.address_idx == address_idx:
2226 |                 return addr_info
2227 |         return None
2228 | 
2229 | 
2230 |     def Jump(self, address, x=0, y=0):
2231 |         """
2232 |         Jumps to the line associated with the given binary address.
2233 | 
2234 |         :param address: The binary address to jump to.
2235 |         :param x: Horizontal cursor position (optional).
2236 |         :param y: Vertical cursor position (optional).
2237 |         :return: Boolean indicating success.
2238 |         """
2239 |         self.CheckRebuild()
2240 |         idx = self._lines_data.bisect_left(address_line_info(address = address + 1 ,address_idx = -1))
2241 |         if idx > 0:
2242 |             lineno = idx - 1
2243 |             return super().Jump(lineno, x, y)
2244 |         return False
2245 | 
2246 | 
2247 |     def GetSelection(self):
2248 |         """
2249 |         Returns the selected range in terms of addresses or None.
2250 | 
2251 |         :return:     - tuple(x1, address1, x2, address2)
2252 |                     - None if no selection
2253 |         """
2254 |         self.CheckRebuild()
2255 |         selection = super().GetSelection()
2256 |         if selection:
2257 |             x1, y1, x2, y2 = selection
2258 |             info1 = self._get_address_info_from_lineno(y1)
2259 |             info2 = self._get_address_info_from_lineno(y2)
2260 |             if info1 is not None and info2 is not None:
2261 |                 addr1 = info1.address
2262 |                 addr2 = info2.address
2263 |                 return (x1, addr1, x2, addr2)
2264 |         return None
2265 | 
2266 | 
2267 |     def GetLineNo(self, mouse = 0):
2268 |         """
2269 |         Returns the index of the current line or None if no line is selected.
2270 | 
2271 |         :param mouse: return mouse position.
2272 |         :return: Returns the line number or None on failure.
2273 |         """
2274 |         self.CheckRebuild()
2275 |         lineno = super().GetLineNo(mouse)
2276 |         if lineno != -1:
2277 |             return lineno
2278 |         return None
2279 | 
2280 | 
2281 |     def GetLineNoFromAddress(self, address):
2282 |         """
2283 |         Returns the line number associated with the given binary address.
2284 |         :param address: The binary address to search for.
2285 |         :return: Returns the line number or -1 if address not found.
2286 |         """
2287 |         self.CheckRebuild()
2288 |         lineno = self._get_lineno_left_from_address(address)
2289 |         if lineno != -1:
2290 |             return lineno
2291 |         return -1
2292 | 
2293 | 
2294 |     def GetAddressFromLineNo(self, lineno):
2295 |         """
2296 |         Returns the binary address associated with the given line number.
2297 |         :param lineno: The line number to search for.
2298 |         :return: Returns the binary address or None if lineno is out of bounds.
2299 |         """
2300 |         self.CheckRebuild()
2301 |         info = self._get_address_info_from_lineno(lineno)
2302 |         if info is not None:
2303 |             return info.address
2304 |         return None
2305 | 
2306 | 
2307 |     # def _JumpToWord(self, word):
2308 |     #     ea = None
2309 |     #     if all(c in "x0123456789abcdefABCDEF" for c in word):
2310 |     #         try:
2311 |     #             ea = int(word,16) # word is a address
2312 |     #         except ValueError:
2313 |     #             pass
2314 |     #     else:
2315 |     #         segm = idaapi.get_segm_by_name(word)
2316 |     #         if segm:
2317 |     #             ea = segm.start_ea # word is a segment name
2318 |     #         else:
2319 |     #             ea =  idaapi.str2ea(word) # word is a name
2320 | 
2321 |     #     if ea == idaapi.BADADDR:
2322 |     #         result: List[str] = split(r'\.|:|;|\+|-|\[|\]', word)
2323 |     #         if result:
2324 |     #             for w in reversed(result):
2325 |     #                 ea =  idaapi.str2ea(w)
2326 |     #                 if ea != idaapi.BADADDR:
2327 |     #                     break
2328 | 
2329 |     #     if ea:
2330 |     #         return idaapi.jumpto(ea)
2331 |     #     return False
2332 | 
2333 | 
2334 |     # def OnDblClick(self, shift):
2335 |     #     dblclick_word = self.GetCurrentWord()
2336 |     #     if not dblclick_word:
2337 |     #         return False
2338 |     #     return self._JumpToWord(dblclick_word)
2339 | 
2340 | 
2341 |     def SetDblClickCallback(self, callback):
2342 |         self.dbl_click_callback = callback
2343 | 
2344 | 
2345 |     def OnDblClick(self, shift):
2346 |         if self.dbl_click_callback:
2347 |             return self.dbl_click_callback(self)
2348 |         return False
2349 | 
2350 | 
2351 |     def OnKeydown(self, vkey, shift):
2352 |         for key, need_shift, callback in self.keydown_callback_list:
2353 |             if key == vkey and (not need_shift or shift):
2354 |                 callback(self)
2355 | 
2356 | 
2357 |     def AddKeyDownCallback(self, target_vkey: int, need_shift: bool, callback):
2358 |         """
2359 |         Add keydown callback function for this viewer
2360 | 
2361 |         callback function receive parameters: the customviewer itself.
2362 |         """
2363 |         self.keydown_callback_list.append((target_vkey, need_shift, callback))
2364 | 
2365 | 
2366 |     def AddAction(self, action_handler: MenuActionHandler):
2367 |         actname = action_handler.get_action_name()
2368 |         ida_kernwin.attach_action_to_popup(self.GetWidget(), None, actname)
2369 | 
2370 |         self.menu_action_handlers.append(action_handler)
2371 | 
2372 | 
2373 | 
2374 | 
2375 | 
2376 | 
2377 | class ColorfulLineGenerator():
2378 |     def __init__(self):
2379 |         pass
2380 | 
2381 |     @staticmethod
2382 |     def GenerateRegisterLine(reg_name, reg_value, value_len) -> str:
2383 |         colored_reg_name = ida_lines.COLSTR(f" {reg_name:>6}", ida_lines.SCOLOR_DEFAULT)
2384 |         colored_reg_value = ida_lines.COLSTR(f"0x{reg_value:0{value_len}X}", ida_lines.SCOLOR_SYMBOL)
2385 |         return colored_reg_name + ": " + colored_reg_value
2386 | 
2387 |     @staticmethod
2388 |     def GenerateDisassemblyDataLine(address, address_len, value, value_len) -> str:
2389 |         addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}", ida_lines.SCOLOR_PREFIX)
2390 |         if len(value) <= 8:
2391 |             data_hex_parts = [f"{byte:02x}" for byte in value]
2392 |             full_str = ' '.join(data_hex_parts)
2393 |             data_str =  full_str
2394 |         else:
2395 |             prefix = ' '.join(f"{byte:02x}" for byte in value[:3])
2396 |             suffix = ' '.join(f"{byte:02x}" for byte in value[-3:])
2397 |             data_str =  f"{prefix} ... {suffix}"
2398 |         colored_value = ida_lines.COLSTR(data_str, ida_lines.SCOLOR_BINPREF)
2399 |         return f"     {addr_str}  {colored_value}"
2400 | 
2401 |     @staticmethod
2402 |     def GenerateDisassemblyNameLine(address, address_len) -> str:
2403 |         addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}", ida_lines.SCOLOR_PREFIX)
2404 |         name_str = ida_name.get_nice_colored_name(address)
2405 |         return f"     {addr_str}  {name_str}"
2406 | 
2407 |     @staticmethod
2408 |     def GenerateDisassemblyCodeLine(address, address_len, value, value_len, execution_counts, generate_by_capstone=False) -> str:
2409 | 
2410 |         if execution_counts == 0:
2411 |             execution_counts_str = "    "
2412 |         elif execution_counts == 1:
2413 |             execution_counts_str = ida_lines.COLSTR(f"{execution_counts: 4}", ida_lines.SCOLOR_AUTOCMT)
2414 |         else:
2415 |             execution_counts_str = ida_lines.COLSTR(f"{execution_counts: 4}", ida_lines.SCOLOR_REGCMT)
2416 | 
2417 |         if not generate_by_capstone:
2418 |             addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}", ida_lines.SCOLOR_PREFIX)
2419 |             insn = ida_lines.generate_disasm_line(address)
2420 |         else:
2421 |             addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}", ida_lines.SCOLOR_IMPNAME)
2422 |             opcode, operand = InstrctionParser().parse_instruction_to_tuple(value)
2423 |             insn = ida_lines.COLSTR(opcode, ida_lines.SCOLOR_INSN) + "     " + ida_lines.COLSTR(operand, ida_lines.SCOLOR_DNAME)
2424 |         return f"{execution_counts_str} {addr_str}  {insn}"
2425 | 
2426 |     @staticmethod
2427 |     def GenerateUnknownLine(address, address_len):
2428 |         addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}", ida_lines.SCOLOR_DREFTAIL)
2429 |         return f"     {addr_str}  ??"
2430 | 
2431 |     @staticmethod
2432 |     def GenerateEmptyLine(address, address_len):
2433 |         addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}", ida_lines.SCOLOR_PREFIX)
2434 |         return f"     {addr_str}  "
2435 | 
2436 |     @staticmethod
2437 |     def GenerateErrorLine(address, address_len, error_msg):
2438 |         addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}  {error_msg}", ida_lines.SCOLOR_ERROR)
2439 |         return f"     {addr_str}  "
2440 | 
2441 |     @staticmethod
2442 |     def GenerateMemoryLine(address, address_len, data_bytes, bytes_per_line=16, data_patch_indices: Optional[List[int]] = None) -> Tuple[str, List[Tuple[int, int, int]]]:
2443 |         """
2444 |         Generates a colored memory dump line.
2445 |         Example: 0x00401000  48 83 EC 28 48 8B C4 48 | H. .(.H. .H
2446 | 
2447 |         :return: Returns a tuple (colored_line, hightlight_byte_indices)
2448 |             colored_line: The colored memory dump line.
2449 |             hightlight_byte_indices: A list of tuples (color, start_index, length) indicating the byte ranges to highlight.
2450 |         """
2451 |         addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}", ida_lines.SCOLOR_PREFIX)
2452 | 
2453 |         hex_dump_parts = []
2454 |         ascii_dump_parts = []
2455 | 
2456 |         hightlight_byte_indices: List[Tuple[int, int, int]] = []
2457 |         for i, byte_val in enumerate(data_bytes):
2458 |             hex_part = f"{byte_val:02X}"
2459 | 
2460 |             # Check if this byte is part of a patch
2461 |             if data_patch_indices and i in data_patch_indices:
2462 |                 hex_dump_parts.append(ida_lines.COLSTR(hex_part, ida_lines.SCOLOR_BINPREF)) # Highlight changed bytes
2463 |                 hightlight_byte_indices.append((BYTE_CHANGE_HIGHTLIGHT, 2 + address_len + 2 + 3 * i, 2))
2464 |             else:
2465 |                 hex_dump_parts.append(ida_lines.COLSTR(hex_part, ida_lines.SCOLOR_BINPREF)) # Default color
2466 | 
2467 |             if 0x20 <= byte_val <= 0x7E:
2468 |                 ascii_dump_parts.append(chr(byte_val))
2469 |             else:
2470 |                 ascii_dump_parts.append('.')
2471 | 
2472 |         hex_str = ' '.join(hex_dump_parts).ljust(bytes_per_line * 3 - 1)
2473 |         ascii_str = ''.join(ascii_dump_parts)
2474 | 
2475 |         return f"{addr_str}  {hex_str} | {ascii_str}", hightlight_byte_indices
2476 | 
2477 |     @staticmethod
2478 |     def GenerateEmplyMemoryLine(address, address_len, bytes_per_line=16) -> str:
2479 |         """
2480 |         Generates a colored empty memory dump line.
2481 |         """
2482 |         addr_str = ida_lines.COLSTR(f"0x{address:0{address_len}X}", ida_lines.SCOLOR_DREFTAIL)
2483 |         hex_str = " ".join("??" for _ in range(bytes_per_line))
2484 |         ascii_str = "".join("." for _ in range(bytes_per_line))
2485 | 
2486 |         return f"{addr_str}  {hex_str} | {ascii_str}"
2487 | 
2488 | 
2489 | class TTE_DisassemblyViewer():
2490 |     """
2491 |     a subviewer class of TimeTravelEmuViewer to display disassembly.
2492 | 
2493 |     To use it following the steps:
2494 |     1. Use self.InitViewer() to initializate the viewer.
2495 |     2. Add self.viewer_widget to TimeTravelEmuViewer's layout.
2496 |     3. Use self.LoadListFromESM() to statistics execution counts of code lines.
2497 |     4. Use self.LoadState() to load state data, including self.memory_pages_list and self.execution_counts
2498 |     5. Use self.DisplayMemoryRange() to display memory pages in range.
2499 |     6. Call TimeTravelEmuViewer.Show() to show viewer with the subviewer.
2500 |     """
2501 | 
2502 |     title = "TimeTravelEmuDisassemblyViewer"
2503 | 
2504 |     def __init__(self):
2505 |         self.viewer = AddressAwareCustomViewer()
2506 |         self.bitness = get_bitness()
2507 |         if self.bitness:
2508 |             self.addr_len = self.bitness // 4
2509 |         else:
2510 |             tte_log_err("Failed to get bitness")
2511 | 
2512 |         self.statusbar_state_id_qlabel: Optional[QtWidgets.QLabel] = None
2513 |         self.statusbar_memory_range_qlabel: Optional[QtWidgets.QLabel] = None
2514 | 
2515 |         self.execution_counts = None
2516 |         self.codelines_dict:Dict[int, int] = {}  # {address : address_idx)}
2517 |         self.capstone_lines: Dict[int, Tuple[bytes, int, str]] = {} # {address : (code_size, capstone_insn_str)}
2518 | 
2519 | 
2520 |         self.current_state = None
2521 |         self.memory_pages_list = None
2522 | 
2523 |         self.current_insn_address: int = -1
2524 |         self.current_range_display_start: int = -1;
2525 |         self.current_range_display_end: int = -1;
2526 | 
2527 | 
2528 |         self.hightlighting_lines: List[Tuple[int, int, Optional[int]]] = [] # [(address, lineno, color),...]
2529 |         self.code_lines_comments: List[Tuple[int, bool, str]] = [] # [(address, has_show, comment),...]
2530 | 
2531 |     def InitViewer(self):
2532 |         self.viewer.Create(self.title)
2533 |         self.viewer_widget  = ida_kernwin.PluginForm.FormToPyQtWidget(self.viewer.GetWidget())
2534 |         self._SetCustomViewerStatusBar()
2535 |         self._SetDoubleClickCallback()
2536 |         self._SetMenuActions()
2537 | 
2538 | 
2539 |     def _SetCustomViewerStatusBar(self):
2540 |         # Remove original status bar
2541 |         viewer_status_bar = self.viewer_widget.findChild(QtWidgets.QStatusBar)
2542 |         for widget in viewer_status_bar.findChildren(QtWidgets.QLabel):
2543 |             viewer_status_bar.removeWidget(widget)
2544 | 
2545 |         self.statusbar_state_id_qlabel = QtWidgets.QLabel("[Status: N\\A ]")
2546 |         viewer_status_bar.addWidget(self.statusbar_state_id_qlabel)
2547 |         self.statusbar_memory_range_qlabel = QtWidgets.QLabel("(Memory Range: N\\A )")
2548 |         viewer_status_bar.addWidget(self.statusbar_memory_range_qlabel)
2549 | 
2550 | 
2551 |     def _SetDoubleClickCallback(self):
2552 | 
2553 |         def OnDblClickAction(custom_viewer: AddressAwareCustomViewer):
2554 |             """
2555 |             Action:
2556 |                 If user double click a address, jump to it in IDA View
2557 |                 If user double click a word, jump to it's address in CustomViewer
2558 |             """
2559 |             dblclick_word = custom_viewer.GetCurrentWord()
2560 |             if not dblclick_word:
2561 |                 return False
2562 |             ea = None
2563 |             if all(c in "x0123456789abcdefABCDEF" for c in dblclick_word):
2564 |                 try:
2565 |                     ea = int(dblclick_word,16) # word is a address
2566 |                     return idaapi.jumpto(ea)
2567 |                 except ValueError:
2568 |                     pass
2569 | 
2570 |             if ea is None:
2571 |                 segm = idaapi.get_segm_by_name(dblclick_word)
2572 |                 if segm:
2573 |                     ea = segm.start_ea # word is a segment name
2574 |                 else:
2575 |                     ea =  idaapi.str2ea(dblclick_word) # word is a name
2576 | 
2577 |             if ea == idaapi.BADADDR:
2578 |                 result: List[str] = split(r'\.|:|;|\+|-|\[|\]', dblclick_word)
2579 |                 if result:
2580 |                     for w in reversed(result):
2581 |                         ea =  idaapi.str2ea(w)
2582 |                         if ea != idaapi.BADADDR:
2583 |                             break
2584 |             if ea != idaapi.BADADDR:
2585 |                 self.JumpTo(ea)
2586 |                 return True
2587 |             return False
2588 | 
2589 |         self.viewer.SetDblClickCallback(OnDblClickAction)
2590 | 
2591 | 
2592 |     def GetCursorAddress(self):
2593 |         lineno =  self.viewer.GetLineNo()
2594 |         if lineno is not None:
2595 |             return self.viewer.GetAddressFromLineNo(lineno)
2596 |         return None
2597 | 
2598 | 
2599 |     def _SetMenuActions(self):
2600 |         self.viewer.AddAction(MenuActionHandler(self.title, lambda : True,
2601 |                               f"{self.title}:RefreshAction",
2602 |                               self.RefreshAction, "Refresh", ""))
2603 |         self.viewer.AddAction(MenuActionHandler(self.title, lambda : True,
2604 |                               f"{self.title}:JumpAction",
2605 |                               self.InputJumpAction, "Jump to address", "G"))
2606 |         self.viewer.AddAction(MenuActionHandler(self.title, lambda : True,
2607 |                               f"{self.title}:SetMemoryRangeAction",
2608 |                               self.SetDisplayMemoryRangeAction, "Set memory range", "R"))
2609 | 
2610 | 
2611 |     def UnregisterAction(self):
2612 |         self.viewer.UnregisterAction()
2613 | 
2614 | 
2615 |     def AddMenuActions(self, action_handler: MenuActionHandler):
2616 |         action_handler.parent_title = self.title
2617 |         self.viewer.AddAction(action_handler)
2618 | 
2619 | 
2620 |     def ClearLines(self):
2621 |         self.viewer.ClearLines()
2622 |         self.current_insn_address = -1
2623 |         self.current_range_display_start = -1;
2624 |         self.current_range_display_end = -1;
2625 | 
2626 | 
2627 |     def JumpTo(self, address):
2628 |         assert self.memory_pages_list, "State data not loaded"
2629 | 
2630 |         if self.current_range_display_start < address and self.current_range_display_end > address:
2631 |             self.viewer.Jump(address, 5, 0)
2632 | 
2633 |         else:
2634 |             # page undisplay
2635 |             starts = [entry[0] for entry in self.memory_pages_list]
2636 |             idx = bisect.bisect_right(starts, address) - 1
2637 | 
2638 |             if idx >= 0 and 0 <= address - starts[idx] < PAGE_SIZE:
2639 |                 self.DisplayMemoryRange(starts[idx], starts[idx]+ PAGE_SIZE)
2640 |                 self.viewer.Jump(address, 5, 0)
2641 |             else:
2642 |                 if idaapi.ask_yn(1, "Target address not loaded, Continue?") == 1:
2643 |                     self.DisplayMemoryRange(address, address+ PAGE_SIZE)
2644 |                     self.viewer.Jump(address, 5, 0)
2645 | 
2646 | 
2647 |     def RefreshAction(self):
2648 |         self.ClearHighlightLines()
2649 |         self.ClearCodeLinesComments()
2650 | 
2651 |         self.DisplayMemoryRange(self.current_range_display_start, self.current_range_display_end)
2652 | 
2653 | 
2654 |     def InputJumpAction(self):
2655 |         n = self.viewer.GetLineNo()
2656 |         if n is not None:
2657 |             target_addr = ida_kernwin.ask_addr(self.viewer.GetLineNo(), "Jump to address")
2658 |             if target_addr:
2659 |                 self.JumpTo(target_addr)
2660 |         return None
2661 | 
2662 | 
2663 |     def SetDisplayMemoryRangeAction(self):
2664 |         assert self.memory_pages_list, "State data not loaded"
2665 | 
2666 |         class RangeInputForm(idaapi.Form):
2667 |             def __init__(self, start_addr, end_addr):
2668 |                 self.start_addr = start_addr
2669 |                 self.end_addr = end_addr
2670 | 
2671 |                 self.RangeStart: Optional[ida_kernwin.Form.NumericInput]  = None
2672 |                 self.RangeEnd: Optional[ida_kernwin.Form.NumericInput] = None
2673 |                 super().__init__(
2674 |                 r'''
2675 |                 {FormChangeCb}
2676 |                 <Range Start: {RangeStart}> | <Range End: {RangeEnd}>
2677 |                 ''',
2678 |                 {
2679 |                 "FormChangeCb": self.FormChangeCb(self.OnFormChange),
2680 | 
2681 |                 "RangeStart": self.NumericInput(value = self.start_addr, swidth = 30),
2682 |                 "RangeEnd": self.NumericInput(value = self.end_addr, swidth = 30),
2683 |                 }
2684 |                 )
2685 |                 self.Compile()
2686 | 
2687 |             def OnFormChange(self,fid):
2688 |                 assert self.RangeStart and self.RangeEnd
2689 |                 if fid == self.RangeStart.id or fid == self.RangeEnd.id:
2690 |                     self.start_addr = self.GetControlValue(self.RangeStart)
2691 |                     self.end_addr = self.GetControlValue(self.RangeEnd)
2692 |                 return 1
2693 | 
2694 |         form = RangeInputForm(self.current_range_display_start,self.current_range_display_end)
2695 |         IsSet = form.Execute()
2696 |         if IsSet == 1:
2697 |             range_start: int = form.start_addr # type: ignore
2698 |             range_end: int = form.end_addr # type: ignore
2699 | 
2700 |             if range_start > range_end:
2701 |                 idaapi.warning("Invalid range, start address should be less than end address")
2702 |             elif range_end - range_start > 16 * PAGE_SIZE:
2703 |                 ok = idaapi.ask_yn(0, f"The range is too large, do you want to continue? (Range: {range_start:X} - {range_end:X})")
2704 |                 if ok == 1:
2705 |                     self.DisplayMemoryRange(range_start, range_end)
2706 |             else:
2707 |                 self.DisplayMemoryRange(range_start, range_end)
2708 |         form.Free()
2709 | 
2710 | 
2711 |     def LoadListFromESM(self, state_list: StateList):
2712 |         """
2713 |         Load execution counts from StateList.
2714 | 
2715 |         Here it is believed that the state list is arranged in the order in which it was created,
2716 |         so the execution counts of the later record must be greater than the number of previous records in the same address, and can be directly overwritten.
2717 |         """
2718 |         self.execution_counts= {state.instruction_address: state.execution_count for _, state in state_list}
2719 | 
2720 | 
2721 |     def LoadState(self, state: EmuState, insn_address: int, memory_pages: MemoryPagesDict):
2722 |         assert self.statusbar_state_id_qlabel, "Status bar not initialized"
2723 |         self.ClearHighlightLines()
2724 |         self.ClearCodeLinesComments()
2725 | 
2726 |         self.current_state = state
2727 |         self.current_insn_address = insn_address
2728 |         self.memory_pages_list = sorted(memory_pages.items())
2729 |         self.statusbar_state_id_qlabel.setText(f"[State: {self.current_state.state_id} ]")
2730 | 
2731 | 
2732 |     def AddHightlightLine(self, address, address_idx, color):
2733 |         self.hightlighting_lines.append((address, address_idx, color))
2734 | 
2735 | 
2736 |     def HighlightLines(self):
2737 |         for address, address_idx, color in self.hightlighting_lines:
2738 |             self.viewer.EditLineColor(address, address_idx, 0, color)
2739 | 
2740 | 
2741 |     def ClearHighlightLines(self):
2742 |         while len(self.hightlighting_lines) > 0:
2743 |             address, lineno, color = self.hightlighting_lines.pop()
2744 |             self.viewer.EditLineColor(address, lineno, None)
2745 | 
2746 | 
2747 |     def AddCodeLinesComments(self, address, comment):
2748 |         self.code_lines_comments.append((address, False, comment))
2749 | 
2750 | 
2751 |     def ShowCodeLinesComments(self):
2752 |         for idx, (address, has_show, comment) in enumerate(self.code_lines_comments):
2753 |             if not has_show:
2754 |                 address_idx = self.codelines_dict.get(address, 0)
2755 |                 line = self.viewer.GetLine(address, address_idx)
2756 |                 if line:
2757 |                     line_with_comment = f"{line.value}  {comment}"
2758 |                     self.viewer.EditLine(address,
2759 |                                         address_idx,
2760 |                                         CODE_LINE,
2761 |                                         line_with_comment,
2762 |                                         line.bgcolor,
2763 |                                         line.bgcolor)
2764 |                     # Replace the tuple with a new one with has_show set to True
2765 |                     self.code_lines_comments[idx] = (address, True, comment)
2766 | 
2767 | 
2768 |     def ClearCodeLinesComments(self):
2769 |         while len(self.code_lines_comments) > 0:
2770 |             address, has_show, comment = self.code_lines_comments.pop()
2771 |             address_idx = self.codelines_dict.get(address, 0)
2772 |             if has_show:
2773 |                 line = self.viewer.GetLine(address, address_idx)
2774 |                 if line:
2775 |                     line_without_comment = line.value[:-(len(comment)+2)]
2776 |                     self.viewer.EditLine(address,
2777 |                                         address_idx,
2778 |                                         CODE_LINE,
2779 |                                         line_without_comment,
2780 |                                         line.fgcolor,
2781 |                                         line.bgcolor)
2782 | 
2783 | 
2784 |     def MarkCurrentInsn(self):
2785 |         assert self.current_state is not None, "State not loaded"
2786 |         assert self.memory_pages_list, "State data not loaded"
2787 | 
2788 |         if self.current_insn_address > 0 and self.current_insn_address >= self.current_range_display_start and self.current_insn_address < self.current_range_display_end:
2789 |             # Highlight current instruction in memory range.
2790 |             self.AddHightlightLine(self.current_insn_address, self.codelines_dict.get(self.current_insn_address, 0), EXECUTE_INSN_HILIGHT_COLOR)
2791 | 
2792 |             # If the current instruction is not recognized as a code by IDA, disassemlby it use capstone.
2793 |             if not idc.is_code(ida_bytes.get_flags(self.current_insn_address)):
2794 |                 count = 0
2795 |                 if self.execution_counts and self.current_insn_address in self.execution_counts:
2796 |                     count = self.execution_counts[self.current_insn_address]
2797 |                 if len(self.current_state.instruction) != 0:
2798 |                     code_bytes = self.current_state.instruction
2799 |                     code_size = len(code_bytes)
2800 |                     line_text = ColorfulLineGenerator.GenerateDisassemblyCodeLine(self.current_insn_address,
2801 |                                                                                               self.addr_len,
2802 |                                                                                               code_bytes,
2803 |                                                                                               code_size,
2804 |                                                                                               count,
2805 |                                                                                               True)
2806 |                     self.capstone_lines[self.current_insn_address] = (code_bytes, code_size, line_text)
2807 |                     self.viewer.UpdateLineRange(self.current_insn_address,
2808 |                                                 len(self.current_state.instruction),
2809 |                                                 SINGLE_DATA_LINE,
2810 |                                                 line_text,
2811 |                                                 None,
2812 |                                                 None)
2813 | 
2814 | 
2815 | 
2816 |     def ApplyStatePatchesInViewer(self, mem_patch: Optional[List[Tuple[int, bytes]]], page_diff: Optional[SortedDict]):
2817 |         """
2818 |         Apply memory patches to the viewer and highlight the changed lines.
2819 | 
2820 |         :param mem_patch: A list of tuples (address, value) representing the changed memory.
2821 |         :param page_diff: A sorted dictionary of memory pages, with keys as start addresses and values as tuples (change_mode, data).
2822 |                             change_mode: 1 - removed, 2 - added
2823 |         :return: None
2824 |         """
2825 |         assert self.current_state is not None, "State not loaded"
2826 |         assert self.memory_pages_list, "State data not loaded"
2827 | 
2828 |         if page_diff:
2829 |             need_rebuild = False
2830 |             for start_addr, (change_mode, data) in page_diff.items():
2831 |                 if change_mode == 1:
2832 |                     # removed page
2833 |                     for i in range(start_addr, start_addr + PAGE_SIZE):
2834 |                         if i >= self.current_range_display_start and i < self.current_range_display_end:
2835 |                             need_rebuild = True
2836 | 
2837 |                 elif change_mode == 2:
2838 |                     # added page
2839 |                     for i in range(start_addr, start_addr + PAGE_SIZE):
2840 |                         if i >= self.current_range_display_start and i < self.current_range_display_end:
2841 |                             need_rebuild = True
2842 |             if need_rebuild:
2843 |                 self.DisplayMemoryRange(self.current_range_display_start, self.current_range_display_end)
2844 | 
2845 |         if mem_patch:
2846 |             for addr, value in mem_patch:
2847 |                 if addr > self.current_range_display_start and addr < self.current_range_display_end:
2848 |                     addr_info = self.viewer.GetLine(addr, 0)
2849 |                     if addr_info is not None and addr_info.type == SINGLE_DATA_LINE:
2850 |                         line_text = ColorfulLineGenerator.GenerateDisassemblyDataLine(addr, self.addr_len, value, 1)
2851 |                         self.viewer.EditLine(addr, 0, SINGLE_DATA_LINE, line_text, None, None)
2852 |                     else:
2853 |                         line_text = ColorfulLineGenerator.GenerateDisassemblyDataLine(addr, self.addr_len, value, 1) # TODO: fix this line_text when data size > 1
2854 |                         self.viewer.UpdateLine(addr, 0, SINGLE_DATA_LINE, line_text, None, None)
2855 |                 self.hightlighting_lines.append((addr, 0, CHANGE_HIGHLIGHT_COLOR))
2856 | 
2857 |         self.MarkCurrentInsn()
2858 |         self.HighlightLines()
2859 |         self.ShowCodeLinesComments()
2860 |         self.viewer.Refresh()
2861 | 
2862 | 
2863 |     def DisplayMemoryRange(self, range_start, range_end):
2864 |         """
2865 |         Display memory range in the viewer. This function will clear the previous range and display the new range.
2866 |         By the way, it can also be used to refresh the current range.
2867 | 
2868 |         :param range_start: start address of the range.
2869 |         :param range_end: end address of the range.
2870 |         :return: None
2871 |         """
2872 |         assert self.current_state is not None, "State not loaded"
2873 |         assert self.memory_pages_list and self.execution_counts, "State data not loaded"
2874 |         assert self.statusbar_memory_range_qlabel, "Status bar not initialized"
2875 | 
2876 |         self.viewer.ClearLines()
2877 |         self.codelines_dict.clear()
2878 |         current_addr = range_start
2879 |         for start_addr, (perm, data) in self.memory_pages_list:
2880 |             data_length =  len(data)
2881 | 
2882 |             if start_addr > range_end or start_addr + data_length <= range_start:
2883 |                 continue
2884 | 
2885 |             # Process empty addresses that may exist before the current page
2886 |             while current_addr < start_addr:
2887 |                 line_text = ColorfulLineGenerator.GenerateUnknownLine(current_addr, self.addr_len)
2888 |                 self.viewer.AddLine(current_addr, UNKNOW_LINE, line_text)
2889 |                 self.viewer.AddLine(current_addr, UNKNOW_LINE, line_text)
2890 |                 current_addr += 1
2891 | 
2892 |             while current_addr < start_addr + data_length and current_addr < range_end:
2893 |                 count = 0
2894 |                 address_idx = 0
2895 |                 if current_addr in self.execution_counts:
2896 |                     count = self.execution_counts[current_addr]
2897 | 
2898 |                 current_addr_flag = ida_bytes.get_flags(current_addr)
2899 | 
2900 |                 # If the current address has a name, display it in a separate line.
2901 |                 is_named = ida_bytes.has_any_name(current_addr_flag)
2902 |                 if is_named:
2903 |                     empty_line_text = ColorfulLineGenerator.GenerateEmptyLine(current_addr, self.addr_len)
2904 |                     line_text = ColorfulLineGenerator.GenerateDisassemblyNameLine(current_addr, self.addr_len)
2905 |                     self.viewer.AddLine(current_addr, EMPTY_LINE, empty_line_text)
2906 |                     self.viewer.AddLine(current_addr, NAME_LINE, line_text)
2907 |                     address_idx += 2
2908 | 
2909 |                 # If the data at the current address is recognized as code.
2910 |                 if idc.is_code(current_addr_flag):
2911 |                     code_size = idc.get_item_size(current_addr)
2912 |                     offset = current_addr - start_addr
2913 |                     line_text = ColorfulLineGenerator.GenerateDisassemblyCodeLine(current_addr,
2914 |                                                                                        self.addr_len,
2915 |                                                                                        data[offset : offset + code_size],
2916 |                                                                                        code_size,
2917 |                                                                                        count)
2918 |                     self.viewer.AddLine(current_addr, CODE_LINE, line_text)
2919 |                     self.codelines_dict[current_addr] = address_idx
2920 |                     current_addr += code_size
2921 |                     continue
2922 | 
2923 |                 # If the data at the current address isn't recognized as code, but cached at capstone line dict and has the same bytes.
2924 |                 if current_addr in self.capstone_lines.keys():
2925 |                     code_bytes, code_size, line_text = self.capstone_lines[current_addr]
2926 |                     if code_bytes == data[current_addr - start_addr : current_addr - start_addr + code_size]:
2927 |                         self.viewer.AddLine(current_addr, CODE_LINE, line_text)
2928 |                         self.codelines_dict[current_addr] = address_idx
2929 |                         current_addr += code_size
2930 |                         continue
2931 | 
2932 |                 # If the data at the current address is recognized as data.
2933 |                 if idc.is_data(current_addr_flag):
2934 |                     data_size = idc.get_item_size(current_addr)
2935 |                     offset = current_addr - start_addr
2936 |                     line_text = ColorfulLineGenerator.GenerateDisassemblyDataLine(current_addr,
2937 |                                                                                   self.addr_len,
2938 |                                                                                   data[offset : offset + data_size],
2939 |                                                                                   data_size)
2940 |                     if data_size == 1:
2941 |                         self.viewer.AddLine(current_addr, SINGLE_DATA_LINE, line_text)
2942 |                     else:
2943 |                         self.viewer.AddLine(current_addr, CODE_LINE, line_text)
2944 | 
2945 |                     current_addr += data_size
2946 |                     continue
2947 | 
2948 |                 # Unknown address type.
2949 |                 else:
2950 |                     line_text = ColorfulLineGenerator.GenerateDisassemblyDataLine(current_addr,
2951 |                                                                                   self.addr_len,
2952 |                                                                                   data[current_addr - start_addr : current_addr - start_addr + 1],
2953 |                                                                                   1)
2954 |                     self.viewer.AddLine(current_addr, SINGLE_DATA_LINE, line_text)
2955 | 
2956 |                     current_addr += 1
2957 |                     continue
2958 | 
2959 |         # After processing all memory pages, fill in the possible empty addresses.
2960 |         while current_addr < range_end:
2961 |             line_text = ColorfulLineGenerator.GenerateUnknownLine(current_addr, self.addr_len)
2962 |             self.viewer.AddLine(current_addr, UNKNOW_LINE, line_text)
2963 |             current_addr += 1
2964 | 
2965 |         self.current_range_display_start = range_start
2966 |         self.current_range_display_end = range_end
2967 | 
2968 |         self.MarkCurrentInsn()
2969 |         self.HighlightLines()
2970 |         self.ShowCodeLinesComments()
2971 |         self.viewer.Refresh()
2972 | 
2973 |         # Set status bar memory range label.
2974 |         self.statusbar_memory_range_qlabel.setText(f"(Mem: 0x{range_start:0{self.addr_len}X} ~ 0x{range_end:0{self.addr_len}X})")
2975 | 
2976 | 
2977 | class TTE_RegistersViewer:
2978 |     """
2979 |     a subviewer class of TimeTravelEmuViewer to display registers.
2980 | 
2981 |     To use it following the steps:
2982 |     1. Use self.InitViewer() to initializate the viewer.
2983 |     2. Add self.viewer_widget to TimeTravelEmuViewer's layout.
2984 |     3. Use self.SetRegisters() to load registers values for a specific state.
2985 |     4. Use self.DisplayRegisters() to display the registers values.
2986 |     5. Call TimeTravelEmuViewer.Show() to show the viewer with this subviewer.
2987 |     """
2988 | 
2989 |     title = "TimeTravelEmuRegistersViewer"
2990 | 
2991 |     class custviewer(ida_kernwin.simplecustviewer_t):
2992 |         def __init__(self):
2993 |             super().__init__()
2994 |             self.dbl_click_callback = None
2995 | 
2996 |         def SetDblClickCallback(self, callback):
2997 |             self.dbl_click_callback = callback
2998 | 
2999 |         def OnDblClick(self, shift):
3000 |             if self.dbl_click_callback:
3001 |                 return self.dbl_click_callback(self)
3002 |             return False
3003 | 
3004 | 
3005 |     def __init__(self):
3006 |         self.viewer = self.custviewer()
3007 |         self.bitness = get_bitness()
3008 |         if self.bitness:
3009 |             self.hex_len = self.bitness // 4
3010 |         else:
3011 |             tte_log_err("Failed to get bitness")
3012 | 
3013 |         self.current_state_id: Optional[str] = None # Only SetRegisters() change this value.
3014 | 
3015 |         self.regs_values: Optional[RegistersDict] = None # Only SetRegisters() change this value.
3016 |         self.regs_patch: Optional[List[str]] = None # Only SetRegsPatch() change this value.
3017 | 
3018 | 
3019 |     def InitViewer(self):
3020 |         self.viewer.Create(self.title)
3021 |         self.viewer_widget  = ida_kernwin.PluginForm.FormToPyQtWidget(self.viewer.GetWidget())
3022 |         self.statusbar_label = QtWidgets.QLabel("State ID: N\\A")
3023 | 
3024 |         self._SetCustomViewerStatusBar()
3025 | 
3026 | 
3027 |     def SetDblClickCallback(self, callback):
3028 |         self.viewer.SetDblClickCallback(callback)
3029 | 
3030 | 
3031 |     def _SetCustomViewerStatusBar(self):
3032 |         # Remove original status bar
3033 |         viewer_status_bar = self.viewer_widget.findChild(QtWidgets.QStatusBar)
3034 |         for widget in viewer_status_bar.findChildren(QtWidgets.QLabel):
3035 |             viewer_status_bar.removeWidget(widget)
3036 | 
3037 |         viewer_status_bar.addWidget(self.statusbar_label)
3038 | 
3039 | 
3040 |     def _RefreshStatusBar(self, state_id: str):
3041 |         self.statusbar_label.setText(f"State ID: {state_id}")
3042 | 
3043 | 
3044 |     def SetRegisters(self, state_id: str, regs_values: RegistersDict):
3045 |         self.current_state_id = state_id
3046 |         self.regs_values = regs_values
3047 |         self._RefreshStatusBar(state_id)
3048 | 
3049 | 
3050 |     def SetRegsPatch(self, state_id: str, regs_patch: Optional[RegistersDict]):
3051 |         assert state_id == self.current_state_id
3052 |         if regs_patch is not None:
3053 |             self.regs_patch = list(regs_patch.keys())
3054 | 
3055 | 
3056 |     def DisplayRegisters(self):
3057 |         assert self.regs_values is not None, "No state data loaded"
3058 | 
3059 |         self.viewer.ClearLines()
3060 |         changed_bgcolor = CHANGE_HIGHLIGHT_COLOR
3061 |         for reg_name, reg_value in self.regs_values.items():
3062 | 
3063 |             line_bgcolor = None
3064 | 
3065 |             if self.regs_patch is not None and reg_name in self.regs_patch:
3066 |                 line_bgcolor = changed_bgcolor
3067 | 
3068 |             line_text = ColorfulLineGenerator.GenerateRegisterLine(reg_name, reg_value, self.hex_len)
3069 |             self.viewer.AddLine(line_text, bgcolor=line_bgcolor)
3070 | 
3071 |         self.viewer.Refresh()
3072 | 
3073 | 
3074 | 
3075 | class TTE_MemoryViewer:
3076 |     """
3077 |     A subviewer class of TimeTravelEmuViewer to display memory.
3078 | 
3079 |     To use it, follow these steps:
3080 |     1. Use self.InitViewer() to initialize the viewer.
3081 |     2. Add self.viewer_widget to TimeTravelEmuViewer's layout.
3082 |     3. Use self.LoadState() to load memory pages for a specific state.
3083 |     4. Use self.DisplayMemoryRange() to display a specific memory range.
3084 |     5. Call TimeTravelEmuViewer.Show() to show the viewer with this subviewer.
3085 |     """
3086 | 
3087 |     title = "TimeTravelEmuMemoryViewer"
3088 |     BYTES_PER_LINE = 16 # How many bytes to display per line in the memory dump
3089 | 
3090 |     class MemoryViewerHooks(ida_kernwin.UI_Hooks):
3091 |         def __init__(self, viewer):
3092 |             ida_kernwin.UI_Hooks.__init__(self)
3093 |             self.viewer: AddressAwareCustomViewer = viewer
3094 | 
3095 |             # Key: Line number Value: List, each element is (color key, character start index, number of characters)
3096 |             self.marked_bytes = {}
3097 |             self.hook()
3098 | 
3099 |         def get_lines_rendering_info(self, out, widget, rin):
3100 |             # Apply coloring only if widget is customview
3101 |             if widget == self.viewer.GetWidget():
3102 |                 for section_lines in rin.sections_lines:
3103 |                     for line_info in section_lines:
3104 |                         current_line_nr = ida_kernwin.place_t.as_simpleline_place_t(line_info.at).n
3105 |                         current_address = self.viewer.GetAddressFromLineNo(current_line_nr)
3106 |                         if current_address in self.marked_bytes:
3107 |                             directives = self.marked_bytes[current_address]
3108 | 
3109 |                             for directive in directives:
3110 |                                 e = ida_kernwin.line_rendering_output_entry_t(line_info)
3111 | 
3112 |                                 color, char_start_idx, char_count = directive
3113 |                                 e.bg_color = color
3114 |                                 e.cpx = char_start_idx
3115 |                                 e.nchars = char_count
3116 |                                 e.flags |= ida_kernwin.LROEF_CPS_RANGE
3117 | 
3118 |                                 out.entries.push_back(e)
3119 | 
3120 | 
3121 |     def __init__(self):
3122 |         self.viewer = AddressAwareCustomViewer()
3123 |         self.hook: TTE_MemoryViewer.MemoryViewerHooks = self.MemoryViewerHooks(self.viewer)
3124 |         self.hightlighting_bytes = self.hook.marked_bytes
3125 | 
3126 | 
3127 |         self.bitness = get_bitness()
3128 |         if self.bitness:
3129 |             self.addr_len = self.bitness // 4
3130 |         else:
3131 |             tte_log_err("Failed to get bitness")
3132 |         self.is_be = get_is_be()
3133 | 
3134 |         self.statusbar_state_id_qlabel: Optional[QtWidgets.QLabel] = None
3135 |         self.statusbar_memory_range_qlabel: Optional[QtWidgets.QLabel] = None
3136 | 
3137 |         self.current_state_id: Optional[str] = None
3138 |         self.memory_pages_list: Optional[List[Tuple[int, Tuple[int, bytearray]]]] = None # Sorted list of (start_addr, (perm, data))
3139 | 
3140 |         self.current_range_display_start: int = -1
3141 |         self.current_range_display_end: int = -1
3142 | 
3143 | 
3144 |     def InitViewer(self):
3145 |         self.viewer.Create(self.title)
3146 |         self.viewer_widget  = ida_kernwin.PluginForm.FormToPyQtWidget(self.viewer.GetWidget())
3147 |         self._SetCustomViewerStatusBar()
3148 |         self._SetDoubleClickCallback()
3149 |         self._SetMenuActions()
3150 | 
3151 | 
3152 |     def _SetCustomViewerStatusBar(self):
3153 |         viewer_status_bar = self.viewer_widget.findChild(QtWidgets.QStatusBar)
3154 |         if not viewer_status_bar:
3155 |             # Create a new status bar if it doesn't exist (e.g., if simplecustviewer_t doesn't create one by default)
3156 |             viewer_status_bar = QtWidgets.QStatusBar()
3157 |             self.viewer_widget.layout().addWidget(viewer_status_bar) # Assuming a layout exists
3158 | 
3159 |         # Clear existing widgets from status bar
3160 |         for widget in viewer_status_bar.findChildren(QtWidgets.QLabel):
3161 |             viewer_status_bar.removeWidget(widget)
3162 | 
3163 |         self.statusbar_state_id_qlabel = QtWidgets.QLabel("[State: N\\A ]")
3164 |         viewer_status_bar.addWidget(self.statusbar_state_id_qlabel)
3165 |         self.statusbar_memory_range_qlabel = QtWidgets.QLabel("(Mem: N\\A )")
3166 |         viewer_status_bar.addWidget(self.statusbar_memory_range_qlabel)
3167 | 
3168 | 
3169 |     def _SetDoubleClickCallback(self):
3170 |         def OnDblClickAction(custom_viewer: AddressAwareCustomViewer):
3171 |             """
3172 |             Action: If user double-clicks an address, jump to it in IDA View.
3173 |             """
3174 |             lineno = custom_viewer.GetLineNo() # Get the line number of the clicked line
3175 |             if lineno is not None:
3176 |                 addr = custom_viewer.GetAddressFromLineNo(lineno) # Get the address of the clicked line
3177 |                 if addr is not None:
3178 |                     return idaapi.jumpto(addr)
3179 |             return False
3180 | 
3181 |         self.viewer.SetDblClickCallback(OnDblClickAction)
3182 | 
3183 | 
3184 |     def _SetMenuActions(self):
3185 |         self.viewer.AddAction(MenuActionHandler(self.title, lambda : True,
3186 |                               f"{self.title}:RefreshAction",
3187 |                               self.RefreshAction, "Refresh", ""))
3188 |         self.viewer.AddAction(MenuActionHandler(self.title, lambda : True,
3189 |                               f"{self.title}:JumpAction",
3190 |                               self.InputJumpAction, "Jump to address", "G"))
3191 |         self.viewer.AddAction(MenuActionHandler(self.title, lambda : True,
3192 |                               f"{self.title}:SetMemoryRangeAction",
3193 |                               self.SetDisplayMemoryRangeAction, "Set memory range", "R"))
3194 |         self.viewer.AddAction(MenuActionHandler(self.title, lambda : True,
3195 |                               f"{self.title}:ExportSelectedMemoryAction",
3196 |                               self.ExportSelectedMemoryAction, "Export selected memory", "E"))
3197 | 
3198 | 
3199 |     def AddMenuActions(self, action_handler: MenuActionHandler):
3200 |         action_handler.parent_title = self.title
3201 |         self.viewer.AddAction(action_handler)
3202 | 
3203 | 
3204 |     def UnregisterAction(self):
3205 |         self.viewer.UnregisterAction()
3206 | 
3207 | 
3208 |     def ClearLines(self):
3209 |         self.viewer.ClearLines()
3210 |         self.current_range_display_start = -1
3211 |         self.current_range_display_end = -1
3212 | 
3213 | 
3214 |     def JumpTo(self, address):
3215 |         """
3216 |         Jumps to the given address in the memory viewer.
3217 |         If the address is not within the current display range, it will try to load a page around it.
3218 |         """
3219 |         assert self.memory_pages_list, "Memory pages data not loaded"
3220 | 
3221 |         # Check if the address is already within the currently displayed range
3222 |         if self.current_range_display_start <= address < self.current_range_display_end:
3223 |             self.viewer.Jump(address)
3224 |             return
3225 | 
3226 |         # Find the page containing the address
3227 |         idx = bisect.bisect_right([entry[0] for entry in self.memory_pages_list], address) - 1
3228 | 
3229 |         if idx >= 0:
3230 |             page_start_addr, _ = self.memory_pages_list[idx]
3231 |             if page_start_addr <= address < page_start_addr + PAGE_SIZE:
3232 |                 # Address is within a known page, display that page
3233 |                 self.DisplayMemoryRange(page_start_addr, page_start_addr + PAGE_SIZE)
3234 |                 self.viewer.Jump(address)
3235 |                 return
3236 | 
3237 |         # If the address is not in any known page, or idx is -1,
3238 |         # try to display a default PAGE_SIZE block centered/starting at the address.
3239 |         # Align to page boundary for display consistency.
3240 |         start_display_addr = (address // self.BYTES_PER_LINE) * self.BYTES_PER_LINE
3241 |         end_display_addr = start_display_addr + PAGE_SIZE # Display one page size for context
3242 |         idaapi.warning(f"Target address 0x{address:X} not loaded in any known memory page. Displaying a default range 0x{start_display_addr:X}-0x{end_display_addr:X}.")
3243 |         self.DisplayMemoryRange(start_display_addr, end_display_addr)
3244 |         self.viewer.Jump(address)
3245 | 
3246 | 
3247 |     def RefreshAction(self):
3248 |         self.DisplayMemoryRange(self.current_range_display_start, self.current_range_display_end)
3249 | 
3250 | 
3251 |     def InputJumpAction(self):
3252 |         n = self.viewer.GetLineNo()
3253 |         current_addr = n if n is not None else 0
3254 |         target_addr = ida_kernwin.ask_addr(current_addr, "Jump to address in Memory View")
3255 |         if target_addr is not None and target_addr != idaapi.BADADDR:
3256 |             self.JumpTo(target_addr)
3257 |         return None
3258 | 
3259 | 
3260 |     def SetDisplayMemoryRangeAction(self):
3261 |         assert self.memory_pages_list, "Memory pages data not loaded"
3262 | 
3263 |         class RangeInputForm(idaapi.Form):
3264 |             def __init__(self, start_addr, end_addr):
3265 |                 self.start_addr = start_addr
3266 |                 self.end_addr = end_addr
3267 | 
3268 |                 self.RangeStart: Optional[ida_kernwin.Form.NumericInput]  = None
3269 |                 self.RangeEnd: Optional[ida_kernwin.Form.NumericInput] = None
3270 |                 super().__init__(
3271 |                 r'''
3272 |                 {FormChangeCb}
3273 |                 <Range Start: {RangeStart}> | <Range End: {RangeEnd}>
3274 |                 ''',
3275 |                 {
3276 |                 "FormChangeCb": self.FormChangeCb(self.OnFormChange),
3277 |                 "RangeStart": self.NumericInput(value = self.start_addr, swidth = 30),
3278 |                 "RangeEnd": self.NumericInput(value = self.end_addr, swidth = 30),
3279 |                 }
3280 |                 )
3281 |                 self.Compile()
3282 | 
3283 |             def OnFormChange(self,fid):
3284 |                 assert self.RangeStart and self.RangeEnd
3285 |                 if fid == self.RangeStart.id or fid == self.RangeEnd.id:
3286 |                     self.start_addr = self.GetControlValue(self.RangeStart)
3287 |                     self.end_addr = self.GetControlValue(self.RangeEnd)
3288 |                 return 1
3289 | 
3290 |         # Use current display range as default for the input form
3291 |         form = RangeInputForm(self.current_range_display_start, self.current_range_display_end)
3292 |         IsSet = form.Execute()
3293 |         if IsSet == 1:
3294 |             range_start: int = form.start_addr # type: ignore
3295 |             range_end: int = form.end_addr # type: ignore
3296 | 
3297 |             if range_start > range_end:
3298 |                 idaapi.warning("Invalid range, start address should be less than end address")
3299 |             elif range_end - range_start > 16 * PAGE_SIZE:
3300 |                 ok = idaapi.ask_yn(0, f"The range is too large, do you want to continue? (Range: {range_start:X} - {range_end:X})")
3301 |                 if ok == 1:
3302 |                     self.DisplayMemoryRange(range_start, range_end)
3303 |             else:
3304 |                 self.DisplayMemoryRange(range_start, range_end)
3305 |         form.Free()
3306 | 
3307 | 
3308 |     def ExportSelectedMemoryAction(self):
3309 |         """
3310 |         Exports the currently selected memory data as a byte array to idaapi.msg.
3311 |         """
3312 |         if not self.memory_pages_list:
3313 |             idaapi.warning("No memory data loaded to export.")
3314 |             return
3315 | 
3316 |         selection = self.viewer.GetSelection()
3317 |         if not selection:
3318 |             return
3319 | 
3320 |         # GetSelection returns (x1, address1, x2, address2)
3321 |         x, start_addr, y, end_addr = selection
3322 | 
3323 |         x_offset = max(min((x - self.addr_len - 3) // 3, self.BYTES_PER_LINE ), 0)
3324 |         y_offset = max(min((y - self.addr_len - 2) // 3, self.BYTES_PER_LINE ), 0)
3325 | 
3326 |         start_addr += x_offset
3327 |         end_addr += y_offset
3328 | 
3329 |         if start_addr > end_addr:
3330 |             start_addr, end_addr = end_addr, start_addr # Swap if selected backwards
3331 | 
3332 |         exported_bytes = bytearray()
3333 | 
3334 |         # Iterate through the selected range
3335 |         current_export_addr = start_addr
3336 |         while current_export_addr < end_addr: # <= because end_addr is inclusive from GetSelection
3337 |             found_byte = False
3338 |             for page_start_addr, (perm, page_content) in self.memory_pages_list:
3339 |                 if page_start_addr <= current_export_addr < page_start_addr + PAGE_SIZE:
3340 |                     offset_in_page = current_export_addr - page_start_addr
3341 |                     exported_bytes.append(page_content[offset_in_page])
3342 |                     found_byte = True
3343 |                     break
3344 | 
3345 |             if not found_byte:
3346 |                 exported_bytes.append(0x00) # Append a null byte for unmapped or unavailable memory
3347 | 
3348 |             current_export_addr += 1
3349 | 
3350 |         if not exported_bytes:
3351 |             idaapi.msg("No data found in the selected range.\n")
3352 |             return
3353 | 
3354 |         # Format the bytearray for output via idaapi.msg
3355 |         hex_dump = ', '.join(f"0x{byte:02x}" for byte in exported_bytes)
3356 |         int_dump = int.from_bytes(exported_bytes, byteorder='big' if self.is_be else 'little',  signed=False)
3357 |         str_dump = exported_bytes.decode('utf-8', 'ignore')
3358 | 
3359 |         output_msg = f"--- Exported Memory Data (0x{start_addr:X} - 0x{end_addr:X}) ---\n"
3360 |         output_msg += f"Hex: {hex_dump}\n"
3361 |         output_msg += f"Int: {int_dump:X}\n"
3362 |         output_msg += f"Str: {str_dump}\n"
3363 |         output_msg += f"Length: {len(exported_bytes)} bytes\n"
3364 |         output_msg += "----------------------------------------\n"
3365 | 
3366 |         idaapi.msg(output_msg)
3367 | 
3368 |     def LoadState(self, state_id: str, memory_pages: MemoryPagesDict):
3369 |         """
3370 |         Loads the memory state for display.
3371 |         """
3372 |         assert self.statusbar_state_id_qlabel, "Status bar not initialized"
3373 |         self.current_state_id = state_id
3374 |         # Convert the dictionary to a sorted list of (address, (perm, data)) for efficient iteration
3375 |         self.memory_pages_list = sorted(memory_pages.items())
3376 |         self.statusbar_state_id_qlabel.setText(f"[State: {self.current_state_id} ]")
3377 | 
3378 | 
3379 |     def ClearHighlightBytes(self):
3380 |         self.hightlighting_bytes.clear()
3381 | 
3382 | 
3383 |     def ApplyStatePatchesInViewer(self, mem_patch: Optional[List[Tuple[int, bytes]]], page_diff: Optional[SortedDict]):
3384 |         """
3385 |         Applies memory patches and highlights changed memory in the viewer.
3386 | 
3387 |         :param mem_patch: A list of tuples (address, value) representing the changed memory bytes.
3388 |         :param page_diff: A sorted dictionary of memory pages, with keys as start addresses and values as tuples (change_mode, data).
3389 |                             change_mode: 1 - removed, 2 - added
3390 |         """
3391 |         assert self.memory_pages_list, "Memory pages data not loaded"
3392 | 
3393 |         # Check if a full rebuild is needed due to page additions/removals
3394 |         need_rebuild_for_page_changes = False
3395 |         if page_diff:
3396 |             for start_addr, (change_mode, data) in page_diff.items():
3397 |                 if (start_addr < self.current_range_display_end and
3398 |                     start_addr + PAGE_SIZE > self.current_range_display_start):
3399 |                     need_rebuild_for_page_changes = True
3400 |                     break
3401 | 
3402 |         if need_rebuild_for_page_changes:
3403 |             # Re-display the entire range if pages were added or removed within it
3404 |             self.DisplayMemoryRange(self.current_range_display_start, self.current_range_display_end)
3405 | 
3406 |         # Clear previous highlights
3407 |         self.ClearHighlightBytes()
3408 | 
3409 |         # Apply mem_patch and add new highlights
3410 |         if mem_patch:
3411 |             # Group patches by the start of their display lines
3412 |             patches_by_line_start: Dict[int, Dict[int, bytes]] = defaultdict(dict)
3413 |             for addr, value_bytes in mem_patch:
3414 |                 line_start_addr = (addr // self.BYTES_PER_LINE) * self.BYTES_PER_LINE
3415 |                 offset_within_line = addr % self.BYTES_PER_LINE
3416 |                 # Store the change, potentially as a single byte
3417 |                 for i, byte_val in enumerate(value_bytes):
3418 |                     patches_by_line_start[line_start_addr][offset_within_line + i] = byte_val.to_bytes(1, byteorder='big' if self.is_be else 'little')
3419 | 
3420 |             for line_start_addr, line_patches in patches_by_line_start.items():
3421 |                 line_info = self.viewer.GetLine(line_start_addr, 0) # Assuming address_idx 0 for memory lines
3422 | 
3423 |                 if line_info:
3424 |                     # Get the raw bytes for the line from the current memory state
3425 |                     # This requires searching the memory_pages_list for the relevant page
3426 |                     line_data_bytes = bytearray(self.BYTES_PER_LINE) # Default empty bytes
3427 |                     for page_addr, (perm, page_data) in self.memory_pages_list:
3428 |                         if page_addr <= line_start_addr < page_addr + PAGE_SIZE:
3429 |                             # Calculate offset into page_data
3430 |                             offset_in_page = line_start_addr - page_addr
3431 |                             # Get the relevant chunk, clamp to page boundary
3432 |                             chunk_len = min(self.BYTES_PER_LINE, PAGE_SIZE - offset_in_page)
3433 |                             line_data_bytes = bytearray(page_data[offset_in_page : offset_in_page + chunk_len])
3434 |                             break
3435 | 
3436 |                     # Identify indices within line_data_bytes that were patched
3437 |                     patch_indices_within_line = sorted(line_patches.keys())
3438 | 
3439 |                     # Regenerate the line with highlight colors
3440 |                     new_line_text, hightlight_bytes = ColorfulLineGenerator.GenerateMemoryLine(
3441 |                         line_start_addr,
3442 |                         self.addr_len,
3443 |                         line_data_bytes,
3444 |                         self.BYTES_PER_LINE,
3445 |                         patch_indices_within_line
3446 |                     )
3447 |                     self.viewer.EditLine(line_start_addr, 0, DATA_LINE, new_line_text, None, None) # Background is handled by COLSTR now
3448 |                     self.hightlighting_bytes[line_start_addr] = hightlight_bytes
3449 | 
3450 | 
3451 |                 elif self.current_range_display_start <= line_start_addr < self.current_range_display_end:
3452 |                     temp_line_data = bytearray([0] * self.BYTES_PER_LINE)
3453 |                     for offset, byte_val in line_patches.items():
3454 |                         if offset < self.BYTES_PER_LINE:
3455 |                             temp_line_data[offset : offset + len(byte_val)] = byte_val
3456 | 
3457 |                     new_line_text, hightlight_bytes = ColorfulLineGenerator.GenerateMemoryLine(
3458 |                         line_start_addr,
3459 |                         self.addr_len,
3460 |                         temp_line_data,
3461 |                         self.BYTES_PER_LINE,
3462 |                         sorted(line_patches.keys())
3463 |                     )
3464 |                     self.viewer.AddLine(line_start_addr, DATA_LINE, new_line_text, bgcolor=CHANGE_HIGHLIGHT_COLOR, lazy=True)
3465 |                     self.hightlighting_bytes[line_start_addr] = hightlight_bytes
3466 | 
3467 |         self.viewer.Refresh()
3468 | 
3469 | 
3470 |     def DisplayMemoryRange(self, range_start, range_end):
3471 |         """
3472 |         Displays a specified memory range in the viewer.
3473 |         This function clears the current view and repopulates it.
3474 | 
3475 |         :param range_start: The starting address of the range to display.
3476 |         :param range_end: The ending address (exclusive) of the range to display.
3477 |         """
3478 |         assert self.statusbar_memory_range_qlabel, "Status bar not initialized"
3479 |         assert self.memory_pages_list, "Memory pages data not loaded"
3480 | 
3481 |         self.viewer.ClearLines()
3482 | 
3483 |         current_addr = (range_start // self.BYTES_PER_LINE) * self.BYTES_PER_LINE # Align to line start
3484 | 
3485 |         # Track the actual displayed range for the status bar
3486 |         actual_display_start = -1
3487 |         actual_display_end = -1
3488 | 
3489 |         # Iterate through the requested display range, line by line
3490 |         while current_addr < range_end:
3491 |             line_data = bytearray(self.BYTES_PER_LINE)
3492 |             has_data = False
3493 | 
3494 |             # Find which memory page(s) contain the current line
3495 |             page_found_for_line = False
3496 |             for page_start_addr, (perm, page_content) in self.memory_pages_list:
3497 |                 if page_start_addr <= current_addr < page_start_addr + PAGE_SIZE:
3498 |                     # This page contains part of the current line
3499 |                     page_found_for_line = True
3500 | 
3501 |                     # Calculate relevant offsets and lengths
3502 |                     offset_in_page = current_addr - page_start_addr
3503 |                     bytes_to_copy = min(self.BYTES_PER_LINE, PAGE_SIZE - offset_in_page)
3504 | 
3505 |                     # Copy data from the memory page into line_data
3506 |                     line_data[0:bytes_to_copy] = page_content[offset_in_page : offset_in_page + bytes_to_copy]
3507 |                     has_data = True
3508 | 
3509 |                     # Update actual displayed range
3510 |                     if actual_display_start == -1:
3511 |                         actual_display_start = current_addr
3512 |                     actual_display_end = current_addr + self.BYTES_PER_LINE
3513 |                     break # Assuming one page fully covers or starts the line
3514 | 
3515 |             if not page_found_for_line:
3516 |                 line_text = ColorfulLineGenerator.GenerateEmplyMemoryLine(
3517 |                     current_addr,
3518 |                     self.addr_len,
3519 |                     self.BYTES_PER_LINE
3520 |                 )
3521 |             else:
3522 |                 line_text, _ = ColorfulLineGenerator.GenerateMemoryLine(
3523 |                     current_addr,
3524 |                     self.addr_len,
3525 |                     line_data,
3526 |                     self.BYTES_PER_LINE
3527 |                 )
3528 |             self.viewer.AddLine(current_addr, DATA_LINE, line_text) # Use DATA_LINE for generic memory view
3529 | 
3530 |             current_addr += self.BYTES_PER_LINE
3531 | 
3532 |         self.current_range_display_start = range_start
3533 |         self.current_range_display_end = range_end
3534 | 
3535 |         # Update status bar
3536 |         display_range_text = f"(Mem: 0x{range_start:0{self.addr_len}X} ~ 0x{range_end:0{self.addr_len}X})"
3537 |         if actual_display_start != -1:
3538 |             display_range_text += f" (Loaded: 0x{actual_display_start:0{self.addr_len}X} ~ 0x{actual_display_end:0{self.addr_len}X})"
3539 |         self.statusbar_memory_range_qlabel.setText(display_range_text)
3540 | 
3541 |         self.viewer.Refresh()
3542 | 
3543 | 
3544 | 
3545 | class TimeTravelEmuViewer(ida_kernwin.PluginForm):
3546 |     """
3547 |     The main TimeTravelEmuViewer class.
3548 | 
3549 |     To use it, you need to do the following steps:
3550 |     1. use self.Init() to initialize the viewer.
3551 |     2. use self.LoadESM(state_manager) to load the EmuStateManager into the viewer.
3552 |     3. use self.Show(“title”) to show the viewer with a title.
3553 | 
3554 |     """
3555 | 
3556 |     title = "TimeTravelEmuViewer"
3557 | 
3558 |     def __init__(self):
3559 |         super().__init__()
3560 |         self.is_be = get_is_be()
3561 | 
3562 | 
3563 |         self.disassembly_viewer: TTE_DisassemblyViewer =  TTE_DisassemblyViewer()
3564 |         self.registers_viewer: TTE_RegistersViewer = TTE_RegistersViewer()
3565 |         self.mempages_viewer: TTE_MemoryViewer = TTE_MemoryViewer()
3566 | 
3567 |         self.subchooser_list: List[ida_kernwin.Choose] = []
3568 | 
3569 |         self.state_manager: Optional[EmuStateManager] = None
3570 |         self.state_list: Optional[StateList] = None # List of (state_id, state) formed in order of generation
3571 | 
3572 |         self.current_state_idx: int = 0 # Index of current state in state_list.
3573 |         self.current_state_id: Optional[str] = None # Only SwitchStateDisplay() can change this value.
3574 |         self.current_full_state: Optional[FullEmuState] = None # Only SwitchStateDisplay() can change this value.
3575 | 
3576 |         self.current_diffs: Optional[Tuple[Tuple[Optional[str], str],  Optional[RegistersDiffsDict], Optional[SortedDict], Optional[SortedDict]]] = None
3577 | 
3578 |         # configs
3579 |         self.follow_current_instruction = False # True # Whether to follow the current instruction when switching states.
3580 | 
3581 | 
3582 |     def Init(self):
3583 |         self.disassembly_viewer.InitViewer()
3584 |         self.registers_viewer.InitViewer()
3585 |         self.mempages_viewer.InitViewer()
3586 | 
3587 |         self.SetRegsViewerDoubleClickCallback()
3588 | 
3589 |         self.disassembly_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3590 |                               f"{self.disassembly_viewer.title}:NextStateAction",
3591 |                               self.SwitchNextStateAction, "Next state", NEXT_STATE_ACTION_SHORTCUT))
3592 | 
3593 |         self.disassembly_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3594 |                               f"{self.disassembly_viewer.title}:PrevStateAction",
3595 |                               self.SwitchPrevStateAction, "Previous state", PREV_STATE_ACTION_SHORTCUT))
3596 | 
3597 |         self.disassembly_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3598 |                               f"{self.disassembly_viewer.title}:CursorStateAction",
3599 |                               self.SwitchCursorStateAction, "Switch to state at cursor address", CURSOR_STATE_ACTION_SHORTCUT))
3600 | 
3601 |         self.disassembly_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3602 |                               f"{self.disassembly_viewer.title}:StageInputAction",
3603 |                               self.SwitchInputStateAction, "Switch to input state", "I"))
3604 | 
3605 |         self.disassembly_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3606 |                               f"{self.disassembly_viewer.title}:ChooseStatesAction",
3607 |                               self.ChooseStatesAction,   "Choose states", "C"))
3608 | 
3609 |         self.disassembly_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3610 |                               f"{self.disassembly_viewer.title}:ChooseMemoryPagesAction",
3611 |                               lambda : self.ChooseMemoryPagesAction(self.disassembly_viewer.title, self.disassembly_viewer.DisplayMemoryRange),   "Choose memory pages", "M"))
3612 | 
3613 |         self.disassembly_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3614 |                               f"{self.disassembly_viewer.title}:ShowDiffAction",
3615 |                               lambda : self.ShowDiffsAciton(self.disassembly_viewer.title, self.disassembly_viewer.JumpTo),   "Show diff", "D"))
3616 | 
3617 |         self.disassembly_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3618 |                               f"{self.disassembly_viewer.title}:ToggleFollowCurrentInstructionAction",
3619 |                               self.ToggleFollowCurrentInstructionAction, "Toggle follow current instruction", "F"))
3620 | 
3621 |         self.mempages_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3622 |                               f"{self.mempages_viewer.title}:ChooseMemoryPagesAction",
3623 |                               lambda : self.ChooseMemoryPagesAction(self.mempages_viewer.title, self.mempages_viewer.DisplayMemoryRange),   "Choose memory pages", "M"))
3624 | 
3625 |         self.mempages_viewer.AddMenuActions(MenuActionHandler(None, lambda : True,
3626 |                               f"{self.mempages_viewer.title}:ShowDiffAction",
3627 |                               lambda : self.ShowDiffsAciton(self.mempages_viewer.title, self.mempages_viewer.JumpTo),   "Show diff", "D"))
3628 | 
3629 | 
3630 |     def SetRegsViewerDoubleClickCallback(self):
3631 |         """
3632 |         Set a double click callback for mempages_viewer to allow user to jump to address in disassembly_viewer.
3633 | 
3634 |         """
3635 | 
3636 |         def OnDblClickAction(custom_viewer: AddressAwareCustomViewer):
3637 |             """
3638 |             Action:
3639 |                 If user double click a address which has execute permission, jump to it in DisassemblyViewer
3640 |                 If user double click a address which only has read permission, jump to it in MemoryViewer
3641 |             """
3642 |             dblclick_word = custom_viewer.GetCurrentWord()
3643 |             if not dblclick_word:
3644 |                 return False
3645 |             ea = None
3646 |             if all(c in "x0123456789abcdefABCDEF" for c in dblclick_word):
3647 |                 try:
3648 |                     ea = int(dblclick_word,16) # word is a address
3649 |                     if self.current_full_state:
3650 |                         memory_pages_list: List[Tuple[int, int, bytearray]] = [(addr, perm, data) for addr, (perm, data) in self.current_full_state.memory_pages.items()]
3651 |                         for start_addr, perm, data in memory_pages_list:
3652 |                             if start_addr <= ea < start_addr + len(data) and perm & UC_PROT_EXEC:
3653 |                                 return self.disassembly_viewer.JumpTo(ea)
3654 |                             elif start_addr <= ea < start_addr + len(data) and perm & UC_PROT_READ:
3655 |                                 return self.mempages_viewer.JumpTo(ea)
3656 |                 except ValueError:
3657 |                     pass
3658 |             return False
3659 | 
3660 |         self.registers_viewer.SetDblClickCallback(OnDblClickAction)
3661 | 
3662 | 
3663 |     def OnCreate(self, form):
3664 |         self.parent = self.FormToPyQtWidget(form)
3665 | 
3666 |         self.PopulateForm()
3667 | 
3668 | 
3669 |     def OnClose(self, form):
3670 |         self.disassembly_viewer.UnregisterAction()
3671 |         self.mempages_viewer.UnregisterAction()
3672 | 
3673 |         del self.disassembly_viewer
3674 |         del self.registers_viewer
3675 |         del self.mempages_viewer
3676 | 
3677 |         for sub_chooser in self.subchooser_list:
3678 |             if sub_chooser:
3679 |                 sub_chooser.Close()
3680 | 
3681 |         del self.subchooser_list
3682 |         del self.state_manager
3683 |         del self.state_list
3684 | 
3685 | 
3686 | 
3687 |     def LoadESM(self, state_manager: EmuStateManager):
3688 |         """
3689 |         Load the EmuStateManager into the viewer.
3690 | 
3691 |         """
3692 |         self.state_manager = state_manager
3693 | 
3694 |         self.state_list = self.state_manager.get_state_list()
3695 | 
3696 |         self.disassembly_viewer.LoadListFromESM(self.state_list)
3697 | 
3698 |         self.current_state_idx = 0
3699 |         if self.state_list:
3700 |             self.SwitchStateDisplay(self.state_list[self.current_state_idx][0])
3701 |         else:
3702 |             idaapi.warning("No states loaded from EmuStateManager.")
3703 | 
3704 | 
3705 |     def PopulateForm(self):
3706 |         main_splitter = QtWidgets.QSplitter(QtCore.Qt.Horizontal) # type: ignore
3707 |         main_splitter.setChildrenCollapsible(True)
3708 | 
3709 | 
3710 |         main_splitter.addWidget(self.disassembly_viewer.viewer_widget)
3711 | 
3712 | 
3713 |         right_splitter = QtWidgets.QSplitter(QtCore.Qt.Vertical) # type: ignore
3714 |         right_splitter.setChildrenCollapsible(True)
3715 | 
3716 |         right_splitter.addWidget(self.registers_viewer.viewer_widget)
3717 |         right_splitter.addWidget(self.mempages_viewer.viewer_widget)
3718 | 
3719 |         main_splitter.addWidget(right_splitter)
3720 | 
3721 |         main_layout = QtWidgets.QVBoxLayout(self.parent)
3722 |         main_layout.addWidget(main_splitter)
3723 |         main_layout.setContentsMargins(0, 0, 0, 0)
3724 | 
3725 | 
3726 |     def SwitchNextStateAction(self):
3727 |         assert self.state_list is not None, "No state_list loaded"
3728 |         if self.current_state_idx < len(self.state_list) - 1:
3729 |             self.current_state_idx += 1
3730 |             self.SwitchStateDisplay(self.state_list[self.current_state_idx][0])
3731 | 
3732 | 
3733 |     def SwitchPrevStateAction(self):
3734 |         assert self.state_list is not None, "No state_list loaded"
3735 |         if self.current_state_idx > 0:
3736 |             self.current_state_idx -= 1
3737 |             self.SwitchStateDisplay(self.state_list[self.current_state_idx][0])
3738 | 
3739 | 
3740 |     def SwitchCursorStateAction(self):
3741 |         assert self.state_manager is not None, "No state_manager loaded"
3742 |         assert self.state_list is not None, "No state_list loaded"
3743 | 
3744 |         address = self.disassembly_viewer.GetCursorAddress()
3745 |         if address:
3746 |             address_str = f"0x{address:X}"
3747 |             target_state_id = next((state_id for state_id, state in self.state_list if address_str in state_id), None)
3748 |             if target_state_id:
3749 |                 self.SwitchState(target_state_id)
3750 | 
3751 | 
3752 |     def SwitchInputStateAction(self):
3753 |         assert self.state_manager is not None, "No state_manager loaded"
3754 |         assert self.state_list is not None, "No state_list loaded"
3755 | 
3756 |         input_str = ida_kernwin.ask_str(self.current_state_id if self.current_state_id else "", 0, "Input state ID:")
3757 |         if input_str:
3758 |             self.SwitchState(input_str)
3759 | 
3760 | 
3761 |     def ChooseStatesAction(self):
3762 | 
3763 |         class StateChooser(ida_kernwin.Choose):
3764 |             def __init__(self, title, state_list: StateList, switch_state_func):
3765 |                 ida_kernwin.Choose.__init__(
3766 |                     self,
3767 |                     title,
3768 | 
3769 |                     [
3770 |                         ["Idx", 10 | ida_kernwin.Choose.CHCOL_DEC],
3771 |                         ["State ID", 10 | ida_kernwin.Choose.CHCOL_PLAIN],
3772 |                         ["Instruction", 10 | ida_kernwin.Choose.CHCOL_HEX]
3773 |                      ])
3774 |                 self.is_available = True
3775 | 
3776 |                 self.state_list = state_list
3777 |                 self.switch_state_display_func = switch_state_func
3778 | 
3779 | 
3780 |             def OnInit(self):
3781 |                 self.items = [[str(i),
3782 |                         state_id,
3783 |                         InstrctionParser().parse_instruction_to_str(state.instruction)]
3784 |                         for i, (state_id, state) in enumerate(self.state_list)]
3785 |                 return True
3786 | 
3787 |             def OnClose(self):
3788 |                 self.is_available = False
3789 | 
3790 |             def OnGetSize(self):
3791 |                 return len(self.items)
3792 | 
3793 |             def OnGetLine(self, n):
3794 |                 return self.items[n]
3795 | 
3796 |             def OnDeleteLine(self, n):
3797 |                 pass
3798 | 
3799 |             def OnSelectLine(self, n):
3800 |                 ida_kernwin.msg(f"Switch to state {self.items[n][1]}\n")
3801 |                 self.switch_state_display_func(self.items[n][1])
3802 |                 return ida_kernwin.Choose.NOTHING_CHANGED
3803 | 
3804 |             def OnRefresh(self, n):
3805 |                 self.OnInit()
3806 |                 return [ida_kernwin.Choose.ALL_CHANGED] + self.adjust_last_item(n)
3807 | 
3808 | 
3809 |         if self.state_list is None:
3810 |             tte_log_dbg("No state_list loaded")
3811 |             return
3812 | 
3813 |         states_chooser = StateChooser("State Chooser", self.state_list, self.SwitchState)
3814 |         states_chooser.Show()
3815 |         self.subchooser_list.append(states_chooser)
3816 | 
3817 | 
3818 |     def ChooseMemoryPagesAction(self, parent_title: str,  display_memory_range_func):
3819 | 
3820 |         class MemPageChooser(ida_kernwin.Choose):
3821 |             def __init__(self, title, get_current_full_state_func, display_memory_range_func):
3822 |                 ida_kernwin.Choose.__init__(
3823 |                     self,
3824 |                     title,
3825 |                     [
3826 |                         ["address", 10 | ida_kernwin.Choose.CHCOL_HEX],
3827 |                         ["R", 3 | ida_kernwin.Choose.CHCOL_DEC],
3828 |                         ["W", 3 | ida_kernwin.Choose.CHCOL_DEC],
3829 |                         ["X", 3 | ida_kernwin.Choose.CHCOL_DEC],
3830 |                         ["size", 10 | ida_kernwin.Choose.CHCOL_HEX],
3831 |                     ])
3832 |                 self.is_available = True
3833 | 
3834 |                 self.get_current_full_state_func = get_current_full_state_func
3835 |                 self.display_memory_range_func = display_memory_range_func
3836 |                 self.bitness = get_bitness()
3837 |                 if self.bitness:
3838 |                     self.hex_len = self.bitness // 4
3839 |                 else:
3840 |                     tte_log_err("Failed to get bitness")
3841 |                 self.items = [] # Initialize items list
3842 | 
3843 |             def OnInit(self):
3844 |                 full_state = self.get_current_full_state_func()
3845 |                 if full_state:
3846 |                     memory_pages_list: List[Tuple[int, int, bytearray]] = [(addr, perm, data) for addr, (perm, data) in full_state.memory_pages.items()]
3847 |                     self.items = [[
3848 |                                     f"0x{addr:0{self.hex_len}X}",
3849 |                                     "R" if perm & UC_PROT_READ else " ",
3850 |                                     "W" if perm & UC_PROT_WRITE else " ",
3851 |                                     "X" if perm & UC_PROT_EXEC else " ",
3852 |                                     f"0x{len(data):X}"]
3853 |                                     for addr, perm, data in memory_pages_list]
3854 |                 else:
3855 |                     self.items = [] # No state, no items
3856 |                 return True
3857 | 
3858 |             def OnClose(self):
3859 |                 self.is_available = False
3860 | 
3861 |             def OnGetSize(self):
3862 |                 return len(self.items)
3863 | 
3864 |             def OnGetLine(self, n):
3865 |                 return self.items[n]
3866 | 
3867 |             def OnDeleteLine(self, n):
3868 |                 pass
3869 | 
3870 |             def OnSelectLine(self, n):
3871 |                 addr = int(self.items[n][0], 16)
3872 |                 size = int(self.items[n][4], 16)
3873 |                 self.display_memory_range_func(addr, addr + size)
3874 | 
3875 |                 return ida_kernwin.Choose.NOTHING_CHANGED
3876 | 
3877 |             def OnRefresh(self, n):
3878 |                 self.OnInit()
3879 |                 return [ida_kernwin.Choose.ALL_CHANGED] + self.adjust_last_item(n)
3880 | 
3881 |         if self.current_full_state is None:
3882 |             tte_log_dbg("No current full state")
3883 |             return
3884 | 
3885 |         memory_pages_chooser = MemPageChooser(f"Memory Chooser for {parent_title}", lambda : self.current_full_state, display_memory_range_func)
3886 |         memory_pages_chooser.Show()
3887 |         self.subchooser_list.append(memory_pages_chooser)
3888 | 
3889 | 
3890 |     def ShowDiffsAciton(self, parent_title: str, jumpto_address_func):
3891 | 
3892 |         class DiffsChooser(ida_kernwin.Choose):
3893 |             def __init__(self, title, get_current_diff_func, jumpto_address_func):
3894 |                 ida_kernwin.Choose.__init__(
3895 |                     self,
3896 |                     title,
3897 | 
3898 |                     [
3899 |                         ["Diff Type", 10 | ida_kernwin.Choose.CHCOL_PLAIN],
3900 |                         ['Diff Location', 10 | ida_kernwin.Choose.CHCOL_HEX],
3901 |                         ["Diff Value", 10 | ida_kernwin.Choose.CHCOL_PLAIN],
3902 |                      ])
3903 |                 self.is_available = True
3904 | 
3905 |                 self.get_current_diff_func = get_current_diff_func
3906 |                 self.jumpto_address_func = jumpto_address_func
3907 | 
3908 |                 self.bitness = get_bitness()
3909 |                 if self.bitness:
3910 |                     self.hex_len = self.bitness // 4
3911 |                 else:
3912 |                     tte_log_err("Failed to get bitness")
3913 |                 self.items = [] # Initialize items list
3914 | 
3915 |             def OnInit(self):
3916 |                 (prev_state_id, curr_state_id), regs_diff, mem_diff, page_diff = self.get_current_diff_func() # Optional[Tuple[Optional[RegistersPatchDict], Optional[SortedDict], Optional[SortedDict]]]
3917 |                 tte_log_info(f"Showing differents from {prev_state_id} to {curr_state_id}")
3918 |                 self.items = []
3919 |                 if regs_diff:
3920 |                     for reg_name, (prev_reg_val, reg_val) in regs_diff.items():
3921 |                         self.items.append(["reg diffs",
3922 |                             reg_name,
3923 |                             f"0x{prev_reg_val:0{self.hex_len}X} -> 0x{reg_val:0{self.hex_len}X}"])
3924 |                 if mem_diff:
3925 |                     for addr, (prev_data, data) in mem_diff.items():
3926 |                         # mem_diff stores single byte changes
3927 |                         self.items.append(["mem diffs",
3928 |                             f"0x{addr:0{self.hex_len}X}",
3929 |                             f"0x{prev_data:02X} -> 0x{data:02X}"]) # prev_data and data are int for single bytes
3930 |                 if page_diff:
3931 |                     for addr, (mode, (perm, data)) in page_diff.items(): # data here is bytearray for the whole page
3932 |                         if mode == 2: # Added page
3933 |                             self.items.append(["add page", f"0x{addr:0{self.hex_len}X}", f"perm: 0x{perm:X}, size: 0x{len(data):X}"])
3934 |                         elif mode == 1: # Removed page
3935 |                             self.items.append(["del page", f"0x{addr:0{self.hex_len}X}", f"perm: 0x{perm:X}, size: 0x{len(data):X}"])
3936 |                 return True
3937 | 
3938 |             def OnClose(self):
3939 |                 self.is_available = False
3940 | 
3941 |             def OnGetSize(self):
3942 |                 return len(self.items)
3943 | 
3944 |             def OnDeleteLine(self, n):
3945 |                 pass
3946 | 
3947 |             def OnGetLine(self, n):
3948 |                 return self.items[n]
3949 | 
3950 |             def OnSelectLine(self, n):
3951 |                 locate = self.items[n][1]
3952 |                 if locate.startswith("0x"):
3953 |                     try:
3954 |                         addr = int(locate, 16)
3955 |                         self.jumpto_address_func(addr)
3956 |                     except ValueError:
3957 |                         pass # Not a valid address
3958 |                 return ida_kernwin.Choose.NOTHING_CHANGED
3959 | 
3960 |             def OnRefresh(self, n):
3961 |                 self.OnInit()
3962 |                 return [ida_kernwin.Choose.ALL_CHANGED] + self.adjust_last_item(n)
3963 | 
3964 |         if self.state_list is None:
3965 |             tte_log_dbg("No state_list loaded")
3966 |             return
3967 | 
3968 |         diff_chooser = DiffsChooser(
3969 |             f"Diffs Chooser for {parent_title}",
3970 |             lambda : self.current_diffs, jumpto_address_func)
3971 |         diff_chooser.Show()
3972 |         self.subchooser_list.append(diff_chooser)
3973 | 
3974 | 
3975 |     def ToggleFollowCurrentInstructionAction(self):
3976 |         self.follow_current_instruction = not self.follow_current_instruction
3977 |         idaapi.msg(f"Follow current instruction: {self.follow_current_instruction}\n")
3978 |         # If toggled to True, jump to current instruction
3979 |         if self.follow_current_instruction and self.current_full_state:
3980 |             self.disassembly_viewer.JumpTo(self.current_full_state.instruction_address)
3981 | 
3982 | 
3983 |     def RefreshSubviewers(self):
3984 |         self.subchooser_list = [subchooser for subchooser in self.subchooser_list if hasattr(subchooser, "is_available") and subchooser.is_available] # type: ignore
3985 |         for chooser in self.subchooser_list:
3986 |             chooser.Refresh()
3987 | 
3988 | 
3989 |     def SwitchState(self, target_state_id: str):
3990 |         assert self.state_list is not None, "No state_list loaded"
3991 |         target_state_idx = next((i for i, (x, y) in enumerate(self.state_list) if x == target_state_id), -1)
3992 |         if target_state_idx >= 0:
3993 |             self.current_state_idx = target_state_idx
3994 |             self.SwitchStateDisplay(self.state_list[self.current_state_idx][0])
3995 |         self.SwitchStateDisplay(target_state_id)
3996 | 
3997 | 
3998 |     def SwitchStateDisplay(self, target_state_id: str):
3999 |         assert self.state_manager is not None, "No state_manager loaded"
4000 | 
4001 |         # Avoid unnecessary re-rendering if already on this state
4002 |         if target_state_id == self.current_state_id and self.current_full_state is not None:
4003 |             return
4004 | 
4005 |         target_state = self.state_manager.get_state(target_state_id)
4006 |         if target_state is None:
4007 |             idaapi.warning(f"State {target_state_id} not found in state_manager.")
4008 |             return
4009 | 
4010 |         target_full_state = target_state.generate_full_state(self.state_manager.states_dict)
4011 |         if target_full_state is None:
4012 |             idaapi.warning(f"Failed to generate full state for state {target_state_id}.")
4013 |             return
4014 | 
4015 |         # Catch up all different information between current and target state
4016 |         regs_diff: Optional[RegistersDiffsDict] = None
4017 |         mem_diff: Optional[SortedDict] = None # SortedDict[address, (prev_value, new_value)] (single byte)
4018 |         page_diff: Optional[SortedDict] = None # SortedDict[page_start_addr, (change_mode, (perm, data))]
4019 | 
4020 |         regs_patch: Optional[RegistersPatchDict] = None
4021 |         mem_patch: Optional[List[Tuple[int, bytes]]] = None # List of (address, value_bytes) for changed memory regions
4022 | 
4023 |         if self.current_full_state is not None:
4024 |             entry = self.state_manager.compare_states(self.current_full_state, target_full_state)
4025 |             if entry:
4026 |                 regs_diff = entry[0]
4027 |                 mem_diff_raw = entry[1] # This is SortedDict[address, (prev_int, new_int)]
4028 |                 page_diff = entry[2]
4029 | 
4030 |                 regs_patch = {reg_name : reg_value for reg_name, (_, reg_value) in regs_diff.items()}
4031 | 
4032 |                 # Convert raw mem_diff (single int bytes) to mem_patch (addr, bytes)
4033 |                 if mem_diff_raw:
4034 |                     mem_patch = [(addr, new_val.to_bytes(1, byteorder='big' if get_is_be() else 'little')) for addr, (_, new_val) in mem_diff_raw.items()]
4035 |                     # Reconstruct mem_diff for DiffsChooser as SortedDict[addr, (prev_int, new_int)]
4036 |                     mem_diff = mem_diff_raw
4037 |                 else:
4038 |                     mem_patch = None
4039 |                     mem_diff = None
4040 | 
4041 | 
4042 |         self.current_diffs = ((self.current_state_id, target_state_id), regs_diff, mem_diff, page_diff)
4043 | 
4044 | 
4045 |         # Update Disassembly Viewer
4046 |         self.disassembly_viewer.LoadState(target_state, target_full_state.instruction_address, target_full_state.memory_pages)
4047 |         self.disassembly_viewer.ApplyStatePatchesInViewer(mem_patch, page_diff)
4048 |         if self.follow_current_instruction or self.current_state_id is None: # Only jump if it's the first load or follow is enabled
4049 |             self.disassembly_viewer.JumpTo(target_full_state.instruction_address)
4050 | 
4051 |         # Update Registers Viewer
4052 |         self.registers_viewer.SetRegisters(target_full_state.state_id, target_full_state.registers_map)
4053 |         self.registers_viewer.SetRegsPatch(target_full_state.state_id, regs_patch)
4054 |         self.registers_viewer.DisplayRegisters()
4055 | 
4056 |         # Update Memory Viewer
4057 |         self.mempages_viewer.LoadState(target_state_id, target_full_state.memory_pages)
4058 |         # If no previous range, default to displaying around the instruction address.
4059 |         if self.mempages_viewer.current_range_display_start == -1:
4060 |             default_mem_start = (target_full_state.instruction_address // PAGE_SIZE) * PAGE_SIZE
4061 |             self.mempages_viewer.DisplayMemoryRange(default_mem_start, default_mem_start + PAGE_SIZE)
4062 | 
4063 |         self.mempages_viewer.ApplyStatePatchesInViewer(mem_patch, page_diff)
4064 | 
4065 | 
4066 |         self.current_full_state = target_full_state
4067 |         self.current_state_id = target_state_id
4068 | 
4069 |         self.RefreshSubviewers() # Refresh Choose dialogs if open
4070 | 
4071 | 
4072 | 
4073 | def StartTimeTravelEmulator(settings: EmuSettings) -> None:
4074 |     """
4075 |     Start the Time Travel Emulator with the given settings and display the Time Travel Emu Viewer.
4076 |     :param setting: EmuSettings object containing the emulator settings.
4077 |     """
4078 |     TTE_Logger().start(settings.log_level, settings.log_file_path)
4079 | 
4080 |     emu_executor = EmuExecutor(settings)
4081 |     emu_executor.init()
4082 | 
4083 |     emu_statesmanager = EmuStateManager()
4084 |     emu_tracer = EmuTracer(emu_executor, emu_statesmanager)
4085 |     emu_tracer.init_hook()
4086 | 
4087 |     continues = emu_executor.execute_preprocessing_code(settings.preprocessing_code, emu_executor)
4088 |     if continues != 1:
4089 |         return
4090 | 
4091 |     emu_executor.run()
4092 |     emu_executor.destroy()
4093 | 
4094 |     time_travel_emulator_viewer = TimeTravelEmuViewer()
4095 |     time_travel_emulator_viewer.Init()
4096 |     time_travel_emulator_viewer.LoadESM(emu_statesmanager)
4097 |     time_travel_emulator_viewer.Show(time_travel_emulator_viewer.title)
4098 | 
4099 |     TTE_Logger().stop()
4100 | 
4101 | 
4102 | 
4103 | 
4104 | TTE_RUN_ACTION_NAME =  "TimeTravelEmulator:Run"
4105 | 
4106 | class action_handler(ida_kernwin.action_handler_t):
4107 |     def __init__(self, handler):
4108 |         super().__init__()
4109 |         self.handler = handler
4110 | 
4111 |     def activate(self, ctx):
4112 |         self.handler()
4113 | 
4114 |     def update(self, ctx):
4115 |         return ida_kernwin.AST_ENABLE_FOR_WIDGET if ctx.widget_type == ida_kernwin.BWN_DISASM else ida_kernwin.AST_DISABLE_FOR_WIDGET
4116 | 
4117 | class Hooks(ida_kernwin.UI_Hooks):
4118 |     def finish_populating_widget_popup(self, widget, popup):
4119 |         if ida_kernwin.get_widget_type(widget) == ida_kernwin.BWN_DISASM:
4120 |             ida_kernwin.attach_action_to_popup(widget, popup, TTE_RUN_ACTION_NAME, None)
4121 | tte_run_menu_hook = None
4122 | 
4123 | 
4124 | class TimeTravelEmulator(idaapi.plugin_t):
4125 |     flags = idaapi.PLUGIN_DRAW
4126 |     comment = "Time Travel Emulator"
4127 |     wanted_name = "Time Travel Emulator"
4128 |     wanted_hotkey = ""
4129 | 
4130 |     def __init__(self):
4131 |         super().__init__()
4132 |         self._data = None
4133 | 
4134 |     def term(self):
4135 |         pass
4136 | 
4137 |     def init(self):
4138 |         idaapi.msg(f"[TimeTravelEmulator] Init (version {VERSION})\n")
4139 |         arch = get_arch()
4140 |         if arch != "":
4141 |             # Supported architecture found.
4142 |             idaapi.msg(f"[TimeTravelEmulator] CPU architecture: {arch}\n")
4143 | 
4144 |             # Register actions
4145 |             ida_kernwin.register_action(ida_kernwin.action_desc_t(
4146 |                     TTE_RUN_ACTION_NAME,
4147 |                     "Run Time Travel Emulator",
4148 |                     action_handler(lambda: self.run(None)),
4149 |                     PLUGIN_HOTKEY))
4150 |             global tte_run_menu_hook
4151 |             tte_run_menu_hook = Hooks()
4152 |             tte_run_menu_hook.hook()
4153 | 
4154 | 
4155 |             idaapi.msg("[TimeTravelEmulator] Ready.\n")
4156 |             return idaapi.PLUGIN_KEEP
4157 |         else:
4158 |             # No supported architecture found.
4159 |             idaapi.msg("[TimeTravelEmulator] No supported architecture found.\n")
4160 |             idaapi.msg("[TimeTravelEmulator] Disabled.\n")
4161 |             return idaapi.PLUGIN_SKIP
4162 | 
4163 |     def run(self, arg):
4164 |         if ida_kernwin.find_widget(TimeTravelEmuViewer.title):
4165 |             idaapi.msg("[TimeTravelEmulator] Already running.\n")
4166 |             return
4167 | 
4168 |         start, end = self.get_ea_range()
4169 |         F = EmuSettingsForm(start, end)
4170 |         IsEmulate = F.Execute()
4171 |         if IsEmulate:
4172 |             setting = F.GetSetting()
4173 |             if setting is not None:
4174 |                 StartTimeTravelEmulator(setting)
4175 |         F.Free()
4176 | 
4177 |     def get_ea_range(self):
4178 |         selection, ea_addr, ea_addr_end = idaapi.read_range_selection(None)
4179 | 
4180 |         if selection:
4181 |             return ea_addr, ea_addr_end
4182 |         else:
4183 |             ea = idaapi.get_screen_ea()
4184 |             if ida_dbg.is_debugger_on():
4185 |                 start = ida_dbg.get_reg_val(IP_REG_NAME_MAP[get_arch()])
4186 |                 end =  idc.get_fchunk_attr(ea,8)
4187 |             else:
4188 |                 start = idc.get_func_attr(ea,0)
4189 |                 end = idc.get_fchunk_attr(ea,8)
4190 |             return start, end
4191 | 
4192 | 
4193 | def PLUGIN_ENTRY():
4194 |     return TimeTravelEmulator()


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | bsdiff4
2 | capstone
3 | sortedcontainers
4 | unicorn==2.1.3


--------------------------------------------------------------------------------