9 Star 28 Fork 7

Gitee 极速下载/pe-tree

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
此仓库是为了提升国内下载速度的镜像仓库,每日同步一次。 原始仓库: https://github.com/blackberry/pe_tree
克隆/下载
pe_tree_ida.py 14.26 KB
一键复制 编辑 原始数据 按行查看 历史
Tom Bonner 提交于 2021-05-17 20:55 +08:00 . v1.0.30 (#12)
#
# Copyright (c) 2020 BlackBerry Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""PE Tree IDAPython plugin"""
# Standard imports
import os
import sys
# Qt imports
from PyQt5 import QtCore, QtWidgets
# IDAPython imports
# pylint: disable=import-error
import idaapi
import ida_idaapi
import idc
import ida_bytes
import ida_name
import ida_kernwin
import idautils
import sip
# PE Tree imports
if "pe_tree" in sys.modules:
# Reload all PE Tree modules as IDA will keep them loaded
try:
from importlib import reload
except ImportError:
from imp import reload
try:
reload(pe_tree)
reload(pe_tree.runtime)
reload(pe_tree.form)
reload(pe_tree.info)
reload(pe_tree.contextmenu)
reload(pe_tree.dump_pe)
reload(pe_tree.exceptions)
reload(pe_tree.hash_pe)
reload(pe_tree.map)
reload(pe_tree.qstandarditems)
reload(pe_tree.tree)
reload(pe_tree.utils)
reload(pe_tree.dialogs)
except NameError:
pass
import pe_tree.runtime
import pe_tree.form
import pe_tree.info
class IDARuntime(pe_tree.runtime.Runtime):
"""IDA runtime callbacks"""
def __init__(self, widget, args):
# Load configuration
self.config_file = os.path.join(self.get_temp_dir(), "pe_tree.ini")
super(IDARuntime, self).__init__(widget, args)
def get_temp_dir(self):
"""Get path to temporary directory"""
self.ret = idaapi.get_user_idadir()
return self.ret
def ask_file(self, filename, caption, filter="", save=False):
"""Open/save file dialog"""
if not save:
# Open file dialog
self.ret = ida_kernwin.ask_file(0, os.path.basename(str(filename)), str(caption))
else:
# Save file dialog
if filename[0] == ".":
# Remove leading dot from section names
filename = filename[1:]
self.ret = ida_kernwin.ask_file(1, os.path.basename(str(filename)), str(caption))
return self.ret
def show_widget(self):
"""Display the PE Tree widget"""
self.ret = idaapi.display_widget(self.widget, (idaapi.PluginForm.WOPN_TAB | idaapi.PluginForm.WOPN_MENU | idaapi.PluginForm.WOPN_RESTORE | idaapi.PluginForm.WOPN_PERSIST))
return self.ret
@QtCore.pyqtSlot(object, int)
def jumpto(self, item, offset):
"""Jump to offset in IDA view"""
try:
self.ret = idc.jumpto(offset)
except:
self.ret = False
return self.ret
@QtCore.pyqtSlot(str)
def log(self, output):
"""Print to output"""
print(output)
@QtCore.pyqtSlot(object, object)
def get_bytes(self, start, size):
"""Read bytes from IDB"""
self.ret = idc.get_bytes(start, size)
return self.ret
@QtCore.pyqtSlot(object)
def get_byte(self, offset):
"""Read byte from IDB"""
self.ret = ida_bytes.get_byte(offset)
return self.ret
@QtCore.pyqtSlot(object)
def get_word(self, offset):
"""Read word from IDB"""
self.ret = ida_bytes.get_word(offset)
return self.ret
@QtCore.pyqtSlot(object)
def get_dword(self, offset):
"""Read dword from IDB"""
self.ret = ida_bytes.get_dword(offset)
return self.ret
@QtCore.pyqtSlot(object)
def get_qword(self, offset):
"""Read qword from IDB"""
self.ret = ida_bytes.get_qword(offset)
return self.ret
@QtCore.pyqtSlot(object)
def get_name(self, offset):
"""Get name for offset from IDB"""
self.ret = idc.get_name(offset, ida_name.GN_VISIBLE)
return self.ret
@QtCore.pyqtSlot(object)
def get_segment_name(self, offset):
"""Get segment name for offset from IDB"""
self.ret = idc.get_segm_name(offset)
return self.ret
@QtCore.pyqtSlot(object)
def is_writable(self, offset):
"""Determine if memory address is writable"""
self.ret = bool(idaapi.getseg(offset).perm & idaapi.SEGPERM_WRITE)
return self.ret
@QtCore.pyqtSlot(object)
def resolve_address(self, offset):
"""Get module/symbol name for address"""
symbol = self.get_name(offset)
module = self.get_segment_name(offset)
if not module and "_" in symbol:
# No module name for the segment, try to determine from the symbol name
symbol_split = symbol.split("_")
# Given a symbol, i.e. ws2_32_WSAStartup, can we find ws2_32.dll in the list of segments?
for segment in idautils.Segments():
segment_name = idc.get_segm_name(segment).lower()
if segment_name.startswith(symbol_split[0].lower()):
new_name = ""
for i in range(0, len(symbol_split)):
new_name = "{}.dll".format("_".join(symbol_split[0:i]))
if new_name == segment_name:
break
if new_name == segment_name:
module = new_name
break
# Still nothing?!
if not module and "_" in symbol:
symbol_split = symbol.split("_")
j = 1
if symbol_split[0] == "ws2":
j += 1
module = "{}.dll".format("_".join(symbol_split[0:j]))
else:
module = "{}.dll".format(symbol_split[0])
# Strip module name from symbol name
if module:
module_name = module.split(".")[0].lower()
if symbol[:len(module_name)].lower().startswith(module_name):
symbol = symbol[len(module_name) + 1:]
if not symbol:
symbol = "{:x}".format(offset)
self.ret = (module, symbol)
return self.ret
@QtCore.pyqtSlot(object)
def get_label(self, offset):
"""Get label for offest from IDB"""
self.ret = idc.GetDisasm(offset).replace("extrn ", "").split(":")[0]
return self.ret
@QtCore.pyqtSlot(object, object)
def make_string(self, offset, size):
"""Convert range to byte string in IDB"""
self.ret = idc.create_strlit(offset, offset + size)
return self.ret
@QtCore.pyqtSlot(object, object, str, str, bytes)
def make_segment(self, offset, size, class_name="DATA", name="pe_map", data=None):
"""Create a new section in the IDB"""
idc.AddSeg(offset, offset + size, 0, 1, 0, idaapi.scPub)
idc.RenameSeg(offset, str(name))
idc.SetSegClass(offset, str(class_name))
#idc.SegAlign(offset, idc.saRelPara)
if data:
idaapi.patch_many_bytes(offset, bytes(data))
self.ret = None
return self.ret
@QtCore.pyqtSlot(object)
def make_qword(self, offset):
"""Create a qword at the given offset in the IDB"""
self.ret = idaapi.create_qword(offset, 8)
return self.ret
@QtCore.pyqtSlot(object)
def make_dword(self, offset):
"""Create a dword at the given offset in the IDB"""
self.ret = idaapi.create_dword(offset, 4)
return self.ret
@QtCore.pyqtSlot(object)
def make_word(self, offset):
"""Create a word at the given offset in the IDB"""
self.ret = idaapi.create_word(offset, 2)
return self.ret
@QtCore.pyqtSlot(object)
def make_byte(self, offset, size=1):
"""Create a byte at the given offset in the IDB"""
self.ret = idaapi.create_byte(offset, size)
return self.ret
@QtCore.pyqtSlot(object, str)
def make_comment(self, offset, comment):
"""Create a comment at the given offset in the IDB"""
#self.ret = idc.MakeComm(offset, comment)
return self.ret
@QtCore.pyqtSlot(object, str, int)
def make_name(self, offset, name, flags=0):
"""Name the given offset in the IDB"""
self.ret = idc.set_name(offset, str(name), idc.SN_NOCHECK | idc.SN_NOWARN | 0x800)
return self.ret
@QtCore.pyqtSlot()
def get_names(self):
self.ret = list(idautils.Names())
return self.ret
@QtCore.pyqtSlot(object, object, object, object)
def find_iat_ptrs(self, pe, image_base, size, get_word):
"""Find all likely IAT pointers"""
iat_ptrs = []
next_offset = image_base
while next_offset < image_base + size:
offset = next_offset
next_offset = ida_bytes.next_addr(offset)
# Attempt to read the current instruction's effective memory address operand (if present)
mnem = idc.print_insn_mnem(offset).lower()
ptr = 0
if mnem in ["call", "push", "jmp"]:
if idc.get_operand_type(offset, 0) == idc.o_mem:
# Get memory offset for branch instructions
ptr = idc.get_operand_value(offset, 0)
elif mnem in ["mov", "lea"]:
if idc.get_operand_type(offset, 0) == idc.o_reg and idc.get_operand_type(offset, 1) == idc.o_mem:
# Get memory offset for mov/lea instructions
ptr = idc.get_operand_value(offset, 1)
# Does the instruction's memory address operand seem somewhat valid?!
if ptr < 0x1000:
continue
# Resolve pointer from memory operand
iat_offset = get_word(ptr)
# Ignore offset if it is in our image
if image_base <= iat_offset <= image_base + size:
continue
# Get module and API name for offset
module, api = self.resolve_address(iat_offset)
# Ignore the offset if it is in a debug segment or stack etc
if api and module and module.endswith(".dll"):
if not iat_offset in iat_ptrs:
# Add IAT offset, address to patch, module name and API name to list
iat_ptrs.append((iat_offset, offset + idc.get_item_size(offset) - 4, module, api))
self.ret = iat_ptrs
return self.ret
@QtCore.pyqtSlot(object)
def find_pe(self, cursor=False):
"""Search IDB for possible MZ/PE headers"""
info = idaapi.get_inf_structure()
mz_headers = []
# Check current cursor for MZ/PE?
if cursor:
# Get IDA cursor address
addr = idc.here()
# Get segment and end address
s = idaapi.getseg(addr)
e = idc.get_segm_end(addr)
# Check for MZ magic
if ida_bytes.get_word(addr) == 0x5a4d:
# Ensure the PE header is in the segment
e_lfanew = ida_bytes.get_dword(addr + 0x3c)
if addr + e_lfanew + 1 < e:
# Check for PE magic
if ida_bytes.get_word(addr + e_lfanew) == 0x4550:
# Found possible MZ/PE header
mz_headers.append([addr, idc.get_segm_name(addr), info.is_64bit()])
self.ret = mz_headers
return self.ret
# Search all segments
for seg in idautils.Segments():
s = idc.get_segm_start(seg)
e = idc.get_segm_end(seg)
addr = s
while True:
# Find first byte of MZ header
addr = ida_bytes.find_byte(addr, e-addr, 0x4d, 0)
if addr == ida_idaapi.BADADDR or addr >= e:
break
# Check for MZ magic
if ida_bytes.get_word(addr) == 0x5a4d:
# Ensure the PE header is in the segment
e_lfanew = ida_bytes.get_dword(addr + 0x3c)
if addr + e_lfanew + 1 < e:
# Check for PE magic
if ida_bytes.get_word(addr + e_lfanew) == 0x4550:
# Found possible MZ/PE header
mz_headers.append([addr, idc.get_segm_name(s), info.is_64bit()])
# Resume search from next address
addr += 1
self.ret = mz_headers
return self.ret
class pe_tree_plugin_t(idaapi.plugin_t):
"""Main IDAPro plugin class for PE Tree"""
wanted_name = pe_tree.info.__title__
comment = pe_tree.info.__title__
version = pe_tree.info.__version__
website = pe_tree.info.__url__
help = pe_tree.info.__title__
wanted_hotkey = ""
flags = 0
def __init__(self):
self.pe_tree_form = None
def init(self):
"""Initialise PE Tree IDA plugin"""
return ida_idaapi.PLUGIN_KEEP
def term(self):
"""Terminate PE Tree IDA plugin"""
# Wait for plugin threads to complete
if self.pe_tree_form:
self.pe_tree_form.wait_for_threads()
def run(self, arg):
"""Run PE Tree IDA plugin"""
# Get form and widget
form = idaapi.create_empty_widget(self.comment)
widget = sip.wrapinstance(int(form), QtWidgets.QWidget)
runtime = IDARuntime(form, {})
self.pe_tree_form = pe_tree.form.PETreeForm(widget, None, runtime)
# Try to find the IDA input file
filename = idaapi.get_input_file_path()
if filename:
if not os.path.isfile(filename):
filename = os.path.basename(filename)
if not os.path.isfile(filename):
filename = runtime.ask_file(os.path.basename(filename), "Where is the input file?")
if os.path.isfile(filename):
# Map input file
self.pe_tree_form.map_pe(image_base=idaapi.get_imagebase(), filename=filename)
self.pe_tree_form.treeview.setVisible(True)
self.pe_tree_form.show()
def PLUGIN_ENTRY():
"""PE Tree IDAPython plugin entry-point"""
return pe_tree_plugin_t()
def main():
"""PE Tree IDAPython script entry-point"""
PLUGIN_ENTRY().run("")
if __name__ == "__main__":
main()
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Python
1
https://gitee.com/mirrors/pe-tree.git
git@gitee.com:mirrors/pe-tree.git
mirrors
pe-tree
pe-tree
master

搜索帮助