9 Star 28 Fork 7

Gitee 极速下载/pe-tree

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
此仓库是为了提升国内下载速度的镜像仓库,每日同步一次。 原始仓库: https://github.com/blackberry/pe_tree
克隆/下载
pe_tree_rekall.py 11.67 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 Rekall plugin"""
# Standard imports
import os
import io
import sys
# Qt imports
from PyQt5 import QtCore, QtWidgets, QtGui
# Rekall imports
from rekall_lib import utils
# PE Tree imports
import pe_tree.window
import pe_tree.runtime
import pe_tree.form
import pe_tree.info
# pylint: disable=undefined-variable
class RekallRuntime(pe_tree.runtime.Runtime):
"""Rekall runtime callbacks"""
def __init__(self, widget, args):
# Load configuration
self.config_file = os.path.join(self.get_temp_dir(), "pe_tree_rekall.ini")
super(RekallRuntime, self).__init__(widget, args)
# Initialise Rekall
self.cc = session.plugins.cc()
self.pedump = session.plugins.pedump()
@QtCore.pyqtSlot(object, object)
def read_pe(self, image_base, size=0):
"""Read PE image from memory"""
outfd = io.BytesIO()
self.pedump.WritePEFile(outfd, self.opaque["address_space"], image_base)
self.ret = bytearray(outfd.getbuffer())
return self.ret
@QtCore.pyqtSlot(object, object)
def get_bytes(self, start, size):
"""Read bytes from memory"""
task = self.opaque["task"]
if task:
self.cc.SwitchProcessContext(self.opaque["task"])
self.ret = self.opaque["address_space"].read(start, size)
self.ret = b"" if self.ret is None else self.ret
return self.ret
@QtCore.pyqtSlot(object)
def is_writable(self, offset):
"""Determine if memory address is writable"""
self.ret = False
task = self.opaque["task"]
if task is not None:
self.cc.SwitchProcessContext(task)
# Traverse VADs
for vad in task.RealVadRoot.traverse():
# Is address within the VAD region?
if vad.Start <= offset <= vad.End:
# Is the VAD region writable?
if "WRITE" in str(vad.u.VadFlags.ProtectionEnum):
self.ret = True
return self.ret
@QtCore.pyqtSlot(object)
def resolve_address(self, offset):
"""Get module/symbol name for address"""
module = ""
api = ""
task = self.opaque["task"]
if task is not None:
self.cc.SwitchProcessContext(task)
symbols = session.address_resolver.format_address(offset)
if len(symbols) > 0:
symbol = symbols[0]
if "!" in symbol:
split_symbol = symbol.split("!")
module = "{}.dll".format(split_symbol[0])
api = "".join(split_symbol[1:])
self.ret = (module, api)
return self.ret
class ExpandableStandardItemModel(QtGui.QStandardItemModel):
"""Set expandable role on process items to force the drawing of expand/collapse arrows without child items present (allows for lazy loading of modules)"""
ExpandableRole = QtCore.Qt.UserRole + 500
def hasChildren(self, index):
"""Determine if the item has children or is expandable
Args:
index (QModelIndex): Index of item
Returns:
bool: True if the item has children or is expandable
"""
if self.data(index, ExpandableStandardItemModel.ExpandableRole):
return True
return super(ExpandableStandardItemModel, self).hasChildren(index)
class SelectProcess(pe_tree.dialogs.ModulePicker):
"""Process/module picker dialog"""
def __init__(self, window, parent=None):
super(SelectProcess, self).__init__(window, "Select process...", ["Name", "PID", "Image base"], model=ExpandableStandardItemModel, parent=parent)
self.treeview.expanded.connect(self.populate_modules)
def populate_processes(self):
"""Populate the tree view with processes"""
root = self.model.invisibleRootItem()
# Iterate over all processes
for task in session.plugins.pslist(proc_regex="").filter_processes():
# Determine image base (ignore if null)
image_base = int(task.Peb.ImageBaseAddress) if int(task.Peb.ImageBaseAddress) else 0
if image_base == 0:
continue
# Create internal filename
filename = "{} ({}) - {:#08x}".format(str(task.name), task.pid, image_base)
# Create process tree view item
process_item = QtGui.QStandardItem("{}".format(str(task.name)))
process_item.setData({"filename": filename, "address_space": task.get_process_address_space(), "task": task, "image_base": image_base}, QtCore.Qt.UserRole)
process_item.setUserTristate(True)
process_item.setCheckable(True)
process_item.setData(True, ExpandableStandardItemModel.ExpandableRole)
# Determine architecture specifics
if task.address_mode == "AMD64":
ptr_mask = 0xffffffffffffffff
ptr_width = 16
else:
ptr_mask = 0xffffffff
ptr_width = 8
root.appendRow([process_item, QtGui.QStandardItem("{}".format(task.pid)), QtGui.QStandardItem("0x{:0{w}x}".format(image_base & ptr_mask, w=ptr_width))])
self.treeview.resizeColumnToContents(0)
def populate_drivers(self, parent):
"""Populate the tree view with drivers"""
# Build a list of all address spaces
address_spaces = [session.kernel_address_space]
for task in session.plugins.pslist().filter_processes():
address_spaces.append(task.get_process_address_space())
# Determine architecture specifics
if session.profile.metadata("arch") == "AMD64":
ptr_mask = 0xffffffffffffffff
ptr_width = 16
else:
ptr_mask = 0xffffffff
ptr_width = 8
for module in session.plugins.modules().lsmod():
# Ensure the driver's image base is non-null
image_base = int(module.DllBase) if int(module.DllBase) else 0
if image_base == 0:
continue
# Locate the address space
address_space = None
for a in address_spaces:
if a.is_valid_address(image_base):
address_space = a
if not address_space:
continue
# Create internal filename
base_name = os.path.basename(utils.SmartUnicode(module.BaseDllName))
filename = "{} - 0x{:0x}".format(utils.EscapeForFilesystem(base_name), int(module.DllBase))
# Create driver tree view item
driver_item = QtGui.QStandardItem("{}".format(utils.EscapeForFilesystem(base_name)))
driver_item.setData({"filename": filename, "task": None, "address_space": address_space, "image_base": image_base}, QtCore.Qt.UserRole)
driver_item.setCheckable(True)
driver_item.setCheckState(QtCore.Qt.Checked if parent.checkState() == QtCore.Qt.Checked else QtCore.Qt.Unchecked)
if parent.checkState() == QtCore.Qt.Checked:
self.items.add(pe_tree.qstandarditems.HashableQStandardItem(driver_item))
parent.appendRow([driver_item, QtGui.QStandardItem("N/A"), QtGui.QStandardItem("0x{:0{w}x}".format(image_base & ptr_mask, w=ptr_width))])
def populate_dlls(self, parent, task):
"""Populate the tree view with DLLs"""
task_as = task.get_process_address_space()
# Determine architecture specifics
if task.address_mode == "AMD64":
ptr_mask = 0xffffffffffffffff
ptr_width = 16
else:
ptr_mask = 0xffffffff
ptr_width = 8
# Iterate over loaded modules (use ldrmodules plugin)
for module in session.plugins.ldrmodules(pids=[task.pid]):
if module.get("in_mem") and module.get("base") > 0:
# Ensure the image base is non-null and not the parent process
image_base = int(module.get("base"))
if image_base == 0 or image_base == int(task.Peb.ImageBaseAddress):
continue
# Create internal filename
base_name = os.path.basename(utils.SmartUnicode(module.get("in_mem_path")))
filename = "{} ({} - {}) - 0x{:0x}".format(utils.EscapeForFilesystem(base_name), task.name, task.pid, image_base)
# Create DLL tree view item
dll_item = QtGui.QStandardItem("{}".format(utils.EscapeForFilesystem(base_name)))
dll_item.setData({"filename": filename, "task": task, "address_space": task_as, "image_base": image_base}, QtCore.Qt.UserRole)
dll_item.setCheckable(True)
dll_item.setCheckState(QtCore.Qt.Checked if parent.checkState() == QtCore.Qt.Checked else QtCore.Qt.Unchecked)
if parent.checkState() == QtCore.Qt.Checked:
self.items.add(pe_tree.qstandarditems.HashableQStandardItem(dll_item))
parent.appendRow([dll_item, QtGui.QStandardItem("{}".format(task.pid)), QtGui.QStandardItem("0x{:0{w}x}".format(image_base & ptr_mask, w=ptr_width))])
def populate_modules(self, index):
"""Lazy load modules to the tree view when the user expands a process"""
# Find the parent item
parent = self.model.itemFromIndex(index)
args = parent.data(QtCore.Qt.UserRole)
# Ensure the parent is a root item without any children
if parent.parent() is not None or parent.rowCount() > 0:
return
# Get the associated task and address space
task = args["task"]
if task.pid == 4:
# Load drivers under the system process
self.populate_drivers(parent)
else:
# Load DLLs under the parent process
self.populate_dlls(parent, task)
# Expand the tree view and resize the filename column
self.treeview.expand(index)
self.treeview.resizeColumnToContents(0)
def invoke(self):
"""Display the select process/modules dialog"""
super(SelectProcess, self).invoke()
# Has the user accepted, i.e. pressed "Dump"?
if self.exec_() != 0:
# Dump selected processes/modules
for item in self.items:
args = item.item.data(QtCore.Qt.UserRole)
self.window.pe_tree_form.map_pe(filename=args["filename"], image_base=args["image_base"], opaque={"address_space": args["address_space"], "task": args["task"]})
self.close()
def main():
"""PE Tree Rekall script entry-point"""
# Create PE Tree Qt application
application = QtWidgets.QApplication(sys.argv)
window = pe_tree.window.PETreeWindow(application, RekallRuntime, {}, open_file=False)
# Extend menu to include open process
select_process = SelectProcess(window)
open_process_action = QtWidgets.QAction("Process", window)
open_process_action.setShortcut("Ctrl+Shift+P")
open_process_action.setStatusTip("Open process")
open_process_action.triggered.connect(select_process.invoke)
window.open_menu.addAction(open_process_action)
# Invoke the process/module picker dialog
select_process.invoke()
sys.exit(application.exec_())
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

搜索帮助