Skip to content

修复支持IDA pro 9.0 #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sabason wants to merge 1 commit into
base: main
Choose a base branch
from
Open

修复支持IDA pro 9.0 #5

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
305 changes: 235 additions & 70 deletions UserComment.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
import ida_hexrays
import ida_nalt
import pickle
import ida_segment
import ida_funcs

title = "Comments"

Expand All @@ -17,29 +19,111 @@
class UserAddedComments():
def __init__(self):
self.netnode = ida_netnode.netnode()
self.netnode.create("$ UserAddedComments")
node_name = "$ user_comments"
# 尝试获取现有的 netnode
self.netnode = ida_netnode.netnode(node_name)
if not self.netnode:
# 如果不存在则创建新的
self.netnode = ida_netnode.netnode()
self.netnode.create(node_name)
self.imagebase = ida_nalt.get_imagebase()
if self.imagebase == ida_idaapi.BADADDR:
self.imagebase = 0
self.comments = {}
self.load_comments()

def save_comments(self):
blob = pickle.dumps(self.comments)
self.netnode.setblob(blob, 0, 'C')
try:
# 清理无效地址的注释
valid_comments = {}
for (offset, cmt_type, line_num), comment in self.comments.items():
addr = offset + self.imagebase
if addr != ida_idaapi.BADADDR and addr < 0xFFFFFFFFFFFFFFFF:
valid_comments[(offset, cmt_type, line_num)] = comment

# 更新注释字典
self.comments = valid_comments

# 保存到数据库
blob = pickle.dumps(self.comments)
self.netnode.setblob(blob, 0, 'C')
print(f"Saved {len(self.comments)} comments")
except Exception as e:
print(f"Save comments error: {e}")

def load_comments(self):
blob = self.netnode.getblob(0, 'C')
if blob is not None:
self.comments = pickle.loads(blob)
else:
try:
blob = self.netnode.getblob(0, 'C')
if blob is not None:
self.comments = pickle.loads(blob)
print("加载的注释详情：")
for (offset, cmt_type, line_num), comment in self.comments.items():
print(f"地址: {hex(offset + self.imagebase)}, 类型: {cmt_type}, 行号: {line_num}, 内容: {comment}")
else:
self.comments = {}
print(f"总共加载了 {len(self.comments)} 条注释")
except Exception as e:
print(f"加载注释时出错: {e}")
self.comments = {}

def add_comment(self, ea, cmt_type, comment, line_num=None):
offset = ea - self.imagebase
key = (offset, cmt_type, line_num)
if not comment:
self.comments.pop(key, 0)
else:
self.comments[key] = comment
self.save_comments()
try:
# 确保地址有效
if ea == ida_idaapi.BADADDR:
print(f"Invalid address: {hex(ea)}")
return

offset = ea - self.imagebase
key = (offset, cmt_type, line_num)

# 确保注释是字符串类型
if comment is not None:
comment = str(comment)

print(f"Adding comment at address {hex(ea)} (offset {hex(offset)})")
print(f"Comment type: {cmt_type}, line: {line_num}, content: {comment}")

if not comment:
self.comments.pop(key, 0)
else:
self.comments[key] = comment
self.save_comments()
except Exception as e:
print(f"Error in add_comment: {e}")

def load_all_pseudocode_comments(self):
try:
# 遍历所有函数
for func in ida_funcs.get_func_ranges():
# 尝试反编译每个函数
cfunc = ida_hexrays.decompile(func.start_ea)
if cfunc:
# 获取该函数的所有用户注释
user_cmts = ida_hexrays.restore_user_cmts(cfunc.entry_ea)
if user_cmts:
for tl, cmt in user_cmts.items():
loc = ida_hexrays.treeloc_t()
loc.ea = tl.ea
loc.itp = tl.itp
# 保存注释
self.add_comment(loc.ea, 'pseudocode', cmt)
except Exception as e:
print(f"Error loading all pseudocode comments: {e}")

def clear_invalid_comments(self):
"""清理所有无效地址的注释"""
try:
valid_comments = {}
for (offset, cmt_type, line_num), comment in self.comments.items():
addr = offset + self.imagebase
if addr != ida_idaapi.BADADDR and addr < 0xFFFFFFFFFFFFFFFF:
valid_comments[(offset, cmt_type, line_num)] = comment

self.comments = valid_comments
self.save_comments()
print(f"Cleaned up comments, {len(self.comments)} valid comments remaining")
except Exception as e:
print(f"Error clearing invalid comments: {e}")


class UIHooks(ida_kernwin.UI_Hooks):
Expand All @@ -57,9 +141,28 @@ def __init__(self, usr_cmt):
ida_hexrays.Hexrays_Hooks.__init__(self)
self.usr_cmt = usr_cmt

def func_printed(self, cfunc):
try:
# 当伪代码生成时加载注释
user_cmts = ida_hexrays.restore_user_cmts(cfunc.entry_ea)
if user_cmts:
for tl, cmt in user_cmts.items():
loc = ida_hexrays.treeloc_t()
loc.ea = tl.ea
loc.itp = tl.itp
self.usr_cmt.add_comment(loc.ea, 'pseudocode', cmt)
return 0
except Exception as e:
print(f"Error in func_printed: {e}")
return 0

def cmt_changed(self, cfunc, loc, cmt):
self.usr_cmt.add_comment(loc.ea, 'pseudocode', cmt)
return 0
try:
self.usr_cmt.add_comment(loc.ea, 'pseudocode', cmt)
return 0
except Exception as e:
print(f"Error in cmt_changed: {e}")
return 0


class DisasmHooks(ida_idp.IDB_Hooks):
Expand All @@ -68,42 +171,37 @@ def __init__(self, usr_cmt):
self.usr_cmt = usr_cmt
self.rebased = False

# hook common and repeatable cmts
# 常规注释和可重复注释
def changing_cmt(self, ea, is_repeatable, new_comment):
cur_ea = idc.here()
if cur_ea == ea:
# solve start_ea problems
cur = ida_kernwin.get_cursor()
if (cur != (True, 0, 0)):
# Fix rebasing bug: Rebase pragram will trigger 'changing_cmt', causing to capture auto cmts at ea.
if self.rebased:
self.rebased = False
return 0
if is_repeatable:
self.usr_cmt.add_comment(ea, 'repeatable', new_comment)
else:
self.usr_cmt.add_comment(ea, 'common', new_comment)
return 0
try:
cmt_type = 'repeatable' if is_repeatable else 'common'
self.usr_cmt.add_comment(ea, cmt_type, new_comment)
return 0
except Exception as e:
print(f"Error in changing_cmt: {e}")
return 0

# hook anterior and posterior cmts
# 前置注释和后置注释
def extra_cmt_changed(self, ea, line_idx, cmt):
cur_ea = idc.here()
if cur_ea == ea:
cur = ida_kernwin.get_cursor()
if (cur != (True, 0, 0)):
if line_idx // 1000 == 1: # line_idx = 1xxx
self.usr_cmt.add_comment(ea, 'anterior', cmt, line_num=line_idx % 1000)
if line_idx // 1000 == 2: # line_idx = 2xxx
self.usr_cmt.add_comment(ea, 'posterior', cmt, line_num=line_idx % 1000)
return 0
try:
if line_idx // 1000 == 1: # 前置注释 (line_idx = 1xxx)
self.usr_cmt.add_comment(ea, 'anterior', cmt, line_num=line_idx % 1000)
elif line_idx // 1000 == 2: # 后置注释 (line_idx = 2xxx)
self.usr_cmt.add_comment(ea, 'posterior', cmt, line_num=line_idx % 1000)
return 0
except Exception as e:
print(f"Error in extra_cmt_changed: {e}")
return 0

# hook Function cmts and repeatable Function cmts
def changing_range_cmt(self, kind, a, cmt, is_repeatable):
if is_repeatable:
self.usr_cmt.add_comment(a.start_ea, 'func_repeatable', cmt)
else:
self.usr_cmt.add_comment(a.start_ea, 'func_common', cmt)
return 0
# 函数注释和可重复函数注释
def changing_range_cmt(self, kind, a, cmt, repeatable):
try:
cmt_type = 'func_repeatable' if repeatable else 'func_common'
self.usr_cmt.add_comment(a.start_ea, cmt_type, cmt)
return 0
except Exception as e:
print(f"Error in changing_range_cmt: {e}")
return 0

# program image rebased
def allsegs_moved(self, info):
Expand All @@ -119,14 +217,34 @@ def __init__(self, usr_cmt):
[ ["Address", 10 | ida_kernwin.Choose.CHCOL_HEX],
["Type", 20 | ida_kernwin.Choose.CHCOL_PLAIN],
["Comments", 30 | ida_kernwin.Choose.CHCOL_PLAIN]],
flags = ida_kernwin.Choose.CH_CAN_REFRESH)
flags = ida_kernwin.Choose.CH_CAN_REFRESH | ida_kernwin.Choose.CH_CAN_DEL)
self.usr_cmt = usr_cmt
self.items = []
self.OnInit()

def OnInit(self):
self.usr_cmt.load_comments() # load comments again
self.items = [ [hex(k[0] + self.usr_cmt.imagebase), k[1], v] for k, v in self.usr_cmt.comments.items() ]
return True
try:
self.items = []
# 强制重新加载注释
self.usr_cmt.load_comments()

print(f"Comments count: {len(self.usr_cmt.comments)}") # 调试信息

for (offset, cmt_type, line_num), comment in self.usr_cmt.comments.items():
if comment: # 只添加非空注释
addr = offset + self.usr_cmt.imagebase
type_str = cmt_type
if line_num is not None:
type_str = f"{cmt_type}:{line_num}"
self.items.append([hex(addr), type_str, str(comment)])

# 按地址排序
self.items.sort(key=lambda x: int(x[0], 16))
print(f"Loaded items: {len(self.items)}") # 调试信息
return True
except Exception as e:
print(f"OnInit error: {e}")
return False

def OnGetSize(self):
return len(self.items)
Expand All @@ -145,6 +263,28 @@ def OnSelectLine(self, n):
addr = int(selected_item[0], 16)
ida_kernwin.jumpto(addr)

def OnDeleteLine(self, n):
try:
selected_item = self.items[n]
addr = int(selected_item[0], 16)
type_str = selected_item[1]

# 解析注释类型和行号
cmt_type = type_str
line_num = None
if ':' in type_str:
cmt_type, line_num = type_str.split(':')
line_num = int(line_num)

# 调用 add_comment 传入空注释来删除
self.usr_cmt.add_comment(addr, cmt_type, None, line_num)

print(f"Deleted comment at {hex(addr)}, type: {cmt_type}")
return [ida_kernwin.Choose.ALL_CHANGED]
except Exception as e:
print(f"Error deleting comment: {e}")
return [ida_kernwin.Choose.NOTHING_CHANGED]


def register_open_action(cmt_view):
"""
Expand Down Expand Up @@ -180,30 +320,55 @@ class my_plugin_t(ida_idaapi.plugin_t):
help = ""

def init(self):
self.usr_cmt = UserAddedComments() # Create Comments instance here

self.idb_hook = DisasmHooks(self.usr_cmt) # Hook disassembly comments(common, repeatable, anterior, posterior cmts)
self.idb_hook.hook()

self.ray_hook = PseudoHooks(self.usr_cmt) # Hook pseudo-code comments(cmts in "F5" pseudo-code view)
self.ray_hook.hook()

self.cmt_view = CommentViewer(self.usr_cmt) # Create comment viewer instance

register_open_action(self.cmt_view) # Register to desktop widget and bind shortcut

self.ui_hook = UIHooks(self.cmt_view) # Refresh commnets viewer in real time
self.ui_hook.hook()
return ida_idaapi.PLUGIN_KEEP # Keep us in the memory
try:
self.usr_cmt = UserAddedComments()
# 清理无效注释
self.usr_cmt.clear_invalid_comments()

# 获取当前段的基址作为测试地址
first_seg = ida_segment.get_first_seg()
if first_seg:
test_ea = first_seg.start_ea
print(f"Testing with address: {hex(test_ea)}")
#self.usr_cmt.add_comment(test_ea, 'common', "Test Comment 1")
#self.usr_cmt.add_comment(test_ea + 4, 'repeatable', "Test Comment 2")

# 确保钩子正确安装
self.idb_hooks = DisasmHooks(self.usr_cmt)
if not self.idb_hooks.hook():
print("Failed to install IDB hooks")

self.ray_hooks = PseudoHooks(self.usr_cmt)
if not self.ray_hooks.hook():
print("Failed to install Hexrays hooks")

self.cmt_view = CommentViewer(self.usr_cmt)
register_open_action(self.cmt_view)

self.ui_hooks = UIHooks(self.cmt_view)
if not self.ui_hooks.hook():
print("Failed to install UI hooks")

print("Plugin initialized successfully")
return ida_idaapi.PLUGIN_KEEP
except Exception as e:
print(f"Plugin initialization error: {e}")
return ida_idaapi.PLUGIN_SKIP

def run(self, arg):
#self.cmt_view.Show()
pass

def term(self):
self.ui_hook.unhook()
self.ray_hook.unhook()
self.idb_hook.unhook()
try:
if hasattr(self, 'ui_hooks'):
self.ui_hooks.unhook()
if hasattr(self, 'ray_hooks'):
self.ray_hooks.unhook()
if hasattr(self, 'idb_hooks'):
self.idb_hooks.unhook()
except Exception as e:
print(f"Plugin termination error: {e}")
return


Expand Down