├── .gitignore ├── CHANGELOG.md ├── CONTRIBUTING.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── doc ├── images │ ├── pe_tree.png │ ├── pe_tree_ghidra.png │ ├── pe_tree_ida.png │ ├── pe_tree_minidump.png │ ├── pe_tree_rekall.png │ └── pe_tree_volatility.png └── source │ ├── conf.py │ ├── index.rst │ └── pe_tree.rst ├── pe_tree ├── __init__.py ├── __main__.py ├── carve.py ├── contextmenu.py ├── dialogs.py ├── dump_pe.py ├── exceptions.py ├── form.py ├── ghidra.py ├── hash_pe.py ├── info.py ├── map.py ├── minidump.py ├── qstandarditems.py ├── resources │ ├── cyberchef.ico │ ├── spear.png │ └── virustotal.png ├── runtime.py ├── scandir.py ├── tree.py ├── utils.py ├── volatility.py └── window.py ├── pe_tree_ida.py ├── pe_tree_rekall.py ├── requirements.txt └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | */env* 3 | */__pycache* 4 | *.egg-info* 5 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | All notable changes to this project will be documented in this file. 4 | 5 | ## [1.0.30](/../../releases/tag/1.0.30) - 2021-04-08 6 | 7 | ### Added 8 | 9 | - Plugin for [Ghidra](https://ghidra-sre.org/) using [Ghidra Bridge](https://github.com/justfoxing/ghidra_bridge) 10 | - Plugin for [Volatility](https://github.com/volatilityfoundation/volatility3) 11 | - Application for extracting PE images from Minidumps using [minidump](https://github.com/skelsec/minidump) 12 | - Application for carving PE files from binary files 13 | - Support for unpacking PE files from zip files (including password protected) 14 | - Additional information in PE summary view 15 | - Resource MIME type detection 16 | 17 | ### Changed 18 | - Moved code for finding PE files in an IDB to the IDA runtime 19 | - Switched to non-native file open dialogs as accept was slow on Windows 20 | - Improved IDAPro in-memory PE handling 21 | - Consolidated code in runtime 22 | - Hide the output log view if it is empty 23 | 24 | ### Fixed 25 | 26 | - Bug loading minidump files under IDA when no input file is present 27 | - Added drop shadow to PE map labels to improve text readability 28 | - Bug loading capstone disassembler 29 | - Dumping without touching the IAT works properly now 30 | 31 | ## [1.0.29](/../../releases/tag/1.0.29) - 2020-10-05 32 | 33 | ### Fixed 34 | 35 | - Fixed section MD5 VT search query 36 | 37 | ## [1.0.28](/../../releases/tag/1.0.28) - 2020-07-30 38 | 39 | ### Added 40 | 41 | - [Rekall](http://www.rekall-forensic.com/) integration 42 | - [.gitignore](.gitignore) file 43 | 44 | ### Changed 45 | 46 | - Tidied runtime interface 47 | - Renamed reconstructed imports section from .idata to .pe_tree 48 | 49 | ### Fixed 50 | 51 | - Certificate parsing 52 | - Improved PE dumping/import reconstruction 53 | - pylint updates 54 | 55 | ## [1.0.27](/../../releases/tag/1.0.27) - 2020-05-20 56 | 57 | ### Added 58 | 59 | - Added [LICENSE](LICENSE) 60 | 61 | ### Changed 62 | 63 | - [setup.py](setup.py) now reads requirements from [requirements.txt](requirements.txt) 64 | - Updated IDAPython installation documentation in [README.md](README.md) 65 | 66 | ### Fixed 67 | 68 | - The PE no-overlay hash was not calculated correctly 69 | - Ensure the correct tree item is removed from the view via right click -> remove 70 | 71 | ## [1.0.26](/../../releases/tag/1.0.26) - 2020-05-19 72 | 73 | ### Changed 74 | 75 | - Removed *package_data* from [setup.py](setup.py) 76 | 77 | ### Added 78 | 79 | - Added [MANIFEST.in](MANIFEST.in) to include package resources 80 | 81 | ## [1.0.25](/../../releases/tag/1.0.25) - 2020-05-19 82 | 83 | ### Changed 84 | 85 | - Updated copyright 86 | - Updated VirusTotal logo 87 | - Updated about box information 88 | 89 | ## [1.0.24](/../../releases/tag/1.0.24) - 2020-05-19 90 | 91 | ### Added 92 | 93 | - VirusTotal and CyberChef URLs moved to configuration file 94 | - Improved developer documentation 95 | 96 | ### Changed 97 | 98 | - Updated about box information 99 | 100 | ### Fixed 101 | 102 | - Fixed fatal exception when starting pe_tree application with partial configuration 103 | - Fixed several pylint warnings 104 | 105 | ## [1.0.23](/../../releases/tag/1.0.23) - 2020-05-15 106 | 107 | ### Changed 108 | 109 | - Updated [setup.py](setup.py) to better support installing as either standalone application or IDA plugin 110 | 111 | ## [1.0.22](/../../releases/tag/1.0.22) - 2020-05-14 112 | 113 | ### Added 114 | 115 | - Changelog markdown 116 | - Contributing markdown 117 | 118 | ### Changed 119 | 120 | - Updated [README.md](README.md) 121 | 122 | ### Fixed 123 | 124 | - Fixed bug when loading as IDAPython plugin 125 | 126 | ## [1.0.21](/../../releases/tag/1.0.21) - 2020-05-12 127 | 128 | ### Added 129 | 130 | - IAT reconstruction 131 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | See [README](README.md#for-developers). -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | 203 | -------------------------------------------------- 204 | 205 | pe_tree/map.py contains some code from: 206 | https://github.com/bsouthga/blog/blob/master/public/posts/color-gradients-with-python.md 207 | licensed as follows: 208 | 209 | Copyright 2017 Ben Southgate 210 | Permission is hereby granted, free of charge, to any person obtaining a copy 211 | of this software and associated documentation files (the "Software"), to deal 212 | in the Software without restriction, including without limitation the rights 213 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 214 | copies of the Software, and to permit persons to whom the Software is 215 | furnished to do so, subject to the following conditions: 216 | 217 | The above copyright notice and this permission notice shall be included in 218 | all copies or substantial portions of the Software. 219 | 220 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 221 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 222 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 223 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 224 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 225 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 226 | THE SOFTWARE 227 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include *.md LICENSE 2 | recursive-include pe_tree *.py *.ico *.png 3 | include requirements.txt -------------------------------------------------------------------------------- /doc/images/pe_tree.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/doc/images/pe_tree.png -------------------------------------------------------------------------------- /doc/images/pe_tree_ghidra.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/doc/images/pe_tree_ghidra.png -------------------------------------------------------------------------------- /doc/images/pe_tree_ida.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/doc/images/pe_tree_ida.png -------------------------------------------------------------------------------- /doc/images/pe_tree_minidump.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/doc/images/pe_tree_minidump.png -------------------------------------------------------------------------------- /doc/images/pe_tree_rekall.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/doc/images/pe_tree_rekall.png -------------------------------------------------------------------------------- /doc/images/pe_tree_volatility.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/doc/images/pe_tree_volatility.png -------------------------------------------------------------------------------- /doc/source/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # https://www.sphinx-doc.org/en/master/usage/configuration.html 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | # import os 14 | # import sys 15 | # sys.path.insert(0, os.path.abspath('.')) 16 | 17 | 18 | # -- Project information ----------------------------------------------------- 19 | 20 | import pe_tree.info 21 | 22 | project = 'pe_tree' 23 | copyright = '2020, BlackBerry Limited' 24 | author = 'BlackBerry Research and Intelligence Team' 25 | 26 | # The full version, including alpha/beta/rc tags 27 | release = pe_tree.info.__version__ 28 | 29 | 30 | # -- General configuration --------------------------------------------------- 31 | 32 | # Add any Sphinx extension module names here, as strings. They can be 33 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 34 | # ones. 35 | extensions = ['sphinx.ext.todo', 'sphinx.ext.autodoc', 'sphinx.ext.napoleon'] 36 | 37 | # Add any paths that contain templates here, relative to this directory. 38 | templates_path = ['_templates'] 39 | 40 | # List of patterns, relative to source directory, that match files and 41 | # directories to ignore when looking for source files. 42 | # This pattern also affects html_static_path and html_extra_path. 43 | exclude_patterns = [] 44 | 45 | 46 | # -- Options for HTML output ------------------------------------------------- 47 | 48 | # The theme to use for HTML and HTML Help pages. See the documentation for 49 | # a list of builtin themes. 50 | # 51 | html_theme = 'classic' 52 | 53 | # Add any paths that contain custom static files (such as style sheets) here, 54 | # relative to this directory. They are copied after the builtin static files, 55 | # so a file named "default.css" will overwrite the builtin "default.css". 56 | html_static_path = ['_static'] -------------------------------------------------------------------------------- /doc/source/index.rst: -------------------------------------------------------------------------------- 1 | .. pe_tree documentation master file, created by 2 | sphinx-quickstart on Mon May 18 13:00:58 2020. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | Welcome to pe_tree's documentation! 7 | =================================== 8 | 9 | .. toctree:: 10 | :maxdepth: 2 11 | :caption: Contents: 12 | 13 | .. automodule:: pe_tree 14 | :members: 15 | 16 | Indices and tables 17 | ================== 18 | 19 | * :ref:`genindex` 20 | * :ref:`modindex` 21 | * :ref:`search` 22 | -------------------------------------------------------------------------------- /doc/source/pe_tree.rst: -------------------------------------------------------------------------------- 1 | pe\_tree package 2 | ================ 3 | 4 | Submodules 5 | ---------- 6 | 7 | pe\_tree.contextmenu module 8 | --------------------------- 9 | 10 | .. automodule:: pe_tree.contextmenu 11 | :members: 12 | :undoc-members: 13 | :show-inheritance: 14 | 15 | pe\_tree.dump\_pe module 16 | ------------------------ 17 | 18 | .. automodule:: pe_tree.dump_pe 19 | :members: 20 | :undoc-members: 21 | :show-inheritance: 22 | 23 | pe\_tree.exceptions module 24 | -------------------------- 25 | 26 | .. automodule:: pe_tree.exceptions 27 | :members: 28 | :undoc-members: 29 | :show-inheritance: 30 | 31 | pe\_tree.form module 32 | -------------------- 33 | 34 | .. automodule:: pe_tree.form 35 | :members: 36 | :undoc-members: 37 | :show-inheritance: 38 | 39 | pe\_tree.hash\_pe module 40 | ------------------------ 41 | 42 | .. automodule:: pe_tree.hash_pe 43 | :members: 44 | :undoc-members: 45 | :show-inheritance: 46 | 47 | pe\_tree.info module 48 | -------------------- 49 | 50 | .. automodule:: pe_tree.info 51 | :members: 52 | :undoc-members: 53 | :show-inheritance: 54 | 55 | pe\_tree.map module 56 | ------------------- 57 | 58 | .. automodule:: pe_tree.map 59 | :members: 60 | :undoc-members: 61 | :show-inheritance: 62 | 63 | pe\_tree.qstandarditems module 64 | ------------------------------ 65 | 66 | .. automodule:: pe_tree.qstandarditems 67 | :members: 68 | :undoc-members: 69 | :show-inheritance: 70 | 71 | pe\_tree.runtime module 72 | ----------------------- 73 | 74 | .. automodule:: pe_tree.runtime 75 | :members: 76 | :undoc-members: 77 | :show-inheritance: 78 | 79 | pe\_tree.tree module 80 | -------------------- 81 | 82 | .. automodule:: pe_tree.tree 83 | :members: 84 | :undoc-members: 85 | :show-inheritance: 86 | 87 | pe\_tree.utils module 88 | --------------------- 89 | 90 | .. automodule:: pe_tree.utils 91 | :members: 92 | :undoc-members: 93 | :show-inheritance: 94 | 95 | 96 | Module contents 97 | --------------- 98 | 99 | .. automodule:: pe_tree 100 | :members: 101 | :undoc-members: 102 | :show-inheritance: 103 | -------------------------------------------------------------------------------- /pe_tree/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | -------------------------------------------------------------------------------- /pe_tree/__main__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree standalone application""" 18 | 19 | # Standard imports 20 | import os 21 | import sys 22 | from argparse import ArgumentParser 23 | 24 | # Qt imports 25 | from PyQt5 import QtWidgets 26 | 27 | # QDarkStyle 28 | try: 29 | import qdarkstyle 30 | 31 | HAVE_DARKSTYLE = True 32 | except: 33 | HAVE_DARKSTYLE = False 34 | 35 | # PE Tree imports 36 | import pe_tree.window 37 | import pe_tree.runtime 38 | 39 | class ApplicationRuntime(pe_tree.runtime.Runtime): 40 | """Standalone application runtime""" 41 | def __init__(self, widget, args): 42 | # Load configuration 43 | self.config_file = os.path.join(self.get_temp_dir(), "pe_tree.ini") 44 | 45 | super(ApplicationRuntime, self).__init__(widget, args) 46 | 47 | self.pe_tree_form = None 48 | 49 | def main(args=None): 50 | """Create PE Tree Qt standalone application""" 51 | # Check command line arguments 52 | parser = ArgumentParser(description="PE-Tree") 53 | parser.add_argument("filenames", nargs="*", help="Path(s) to file/folder/zip") 54 | args = parser.parse_args() 55 | 56 | # Create PE Tree application 57 | application = QtWidgets.QApplication(sys.argv) 58 | window = pe_tree.window.PETreeWindow(application, ApplicationRuntime, args) 59 | sys.exit(application.exec_()) 60 | 61 | if __name__ == "__main__": 62 | main() 63 | -------------------------------------------------------------------------------- /pe_tree/carve.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2021 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree Carve application""" 18 | 19 | # Standard imports 20 | import os 21 | import sys 22 | import re 23 | from argparse import ArgumentParser, FileType 24 | 25 | # Qt imports 26 | from PyQt5 import QtCore, QtWidgets 27 | 28 | # PE Tree imports 29 | import pe_tree.window 30 | import pe_tree.runtime 31 | 32 | class CarveRuntime(pe_tree.runtime.Runtime): 33 | """Carve runtime callbacks""" 34 | def __init__(self, widget, args): 35 | # Load configuration 36 | self.config_file = os.path.join(self.get_temp_dir(), "pe_tree_carve.ini") 37 | super(CarveRuntime, self).__init__(widget, args) 38 | 39 | # Read input file 40 | self.data = self.args.filename.read() 41 | 42 | @QtCore.pyqtSlot(object, object) 43 | def get_bytes(self, start, size): 44 | """Read bytes from memory""" 45 | try: 46 | self.ret = self.data[start:start+size] 47 | except: 48 | self.ret = None 49 | 50 | return self.ret 51 | 52 | def main(): 53 | """PE Tree Carve script entry-point""" 54 | # Check command line arguments 55 | parser = ArgumentParser(description="PE-Tree (Carve)") 56 | parser.add_argument("filename", help="Path to file to carve", type=FileType("rb")) 57 | args = parser.parse_args() 58 | 59 | # Create PE Tree Qt application 60 | application = QtWidgets.QApplication(sys.argv) 61 | window = pe_tree.window.PETreeWindow(application, CarveRuntime, args, open_file=False) 62 | 63 | # Determine architecture specifics 64 | ptr_width = 16 65 | 66 | # Iterate over all MZ bytes in the input file 67 | for match in re.compile(b"MZ").finditer(window.runtime.data): 68 | # Determine image base 69 | image_base = int(match.start()) 70 | 71 | # Attempt to map PE 72 | window.pe_tree_form.map_pe(filename="Offset {} ({:#08x})".format(image_base, image_base), image_base=image_base, disable_dump=True) 73 | 74 | sys.exit(application.exec_()) 75 | 76 | if __name__ == "__main__": 77 | main() 78 | -------------------------------------------------------------------------------- /pe_tree/contextmenu.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree right-click context menu""" 18 | 19 | # Standard imports 20 | import sys 21 | import base64 22 | import binascii 23 | import webbrowser 24 | 25 | try: 26 | from urllib import quote 27 | except ImportError: 28 | from urllib.parse import quote 29 | 30 | # pefile 31 | import pefile 32 | 33 | # IDA imports 34 | try: 35 | import idaapi 36 | 37 | HAVE_IDA = True 38 | except ImportError: 39 | HAVE_IDA = False 40 | 41 | # Qt imports 42 | from PyQt5 import QtWidgets, Qt 43 | 44 | # PE Tree imports 45 | import pe_tree.form 46 | import pe_tree.utils 47 | import pe_tree.dialogs 48 | 49 | class CommonStandardItemContextMenu(): 50 | """Context menu actions for pe_tree.qstandarditems.CommonStandardItem 51 | 52 | Instantiate once and re-use for each tree item right-click action by calling new_menu(). 53 | 54 | Args: 55 | form (pe_tree.form): PE Tree form 56 | 57 | """ 58 | 59 | def __init__(self, form): 60 | self.form = form 61 | self.point = None 62 | self.item = None 63 | self.index = None 64 | self.menu = None 65 | self.actions = [] 66 | self.altitudes = {} 67 | 68 | style = form.widget.style() 69 | 70 | # Create actions/menus 71 | self.save_dump_action = QtWidgets.QAction(style.standardIcon(QtWidgets.QStyle.SP_DialogSaveButton), "Dump...", form.widget) 72 | self.expand_all_action = QtWidgets.QAction(style.standardIcon(QtWidgets.QStyle.SP_DirOpenIcon), "Expand all", form.widget) 73 | self.remove_action = QtWidgets.QAction(style.standardIcon(QtWidgets.QStyle.SP_DialogCancelButton), "Remove", form.widget) 74 | self.copy_action = QtWidgets.QAction("Copy", form.widget) 75 | self.hexdump_action = QtWidgets.QAction("Hex-dump", form.widget) 76 | self.disassemble_action = QtWidgets.QAction("Disassemble", form.widget) 77 | self.search_idb_action = QtWidgets.QAction(style.standardIcon(QtWidgets.QStyle.SP_FileDialogContentsView), "Search IDB", form.widget) 78 | self.map_pe_from_cursor_action = QtWidgets.QAction(style.standardIcon(QtWidgets.QStyle.SP_ArrowRight), "From cursor", form.widget) 79 | self.map_pe_from_file_action = QtWidgets.QAction("From file", form.widget) 80 | self.add_pe_menu = QtWidgets.QMenu("Add PE") 81 | self.add_pe_menu.addAction(self.search_idb_action) 82 | self.add_pe_menu.addAction(self.map_pe_from_cursor_action) 83 | self.add_pe_menu.addAction(self.map_pe_from_file_action) 84 | self.rva_action = QtWidgets.QAction("RVA", form.widget) 85 | self.rva_action.setCheckable(True) 86 | self.va_action = QtWidgets.QAction("VA", form.widget) 87 | self.va_action.setCheckable(True) 88 | self.addressing_menu = QtWidgets.QMenu("Addressing") 89 | self.addressing_menu.addAction(self.rva_action) 90 | self.addressing_menu.addAction(self.va_action) 91 | self.about_action = QtWidgets.QAction("About", form.widget) 92 | 93 | if not HAVE_IDA: 94 | self.add_pe_menu.menuAction().setVisible(False) 95 | 96 | self.save_action = QtWidgets.QAction(style.standardIcon(QtWidgets.QStyle.SP_DialogSaveButton), "Save", form.widget) 97 | self.search_vt_action = QtWidgets.QAction(self.form.vt_icon, "VirusTotal", form.widget) 98 | self.cyberchef_action = QtWidgets.QAction(self.form.cyberchef_icon, "CyberChef", form.widget) 99 | 100 | self.toggle_debug_action = QtWidgets.QAction("Debug", form.widget) 101 | self.toggle_debug_action.setCheckable(True) 102 | 103 | self.options_menu = QtWidgets.QMenu("Options") 104 | self.options_menu.addAction(self.toggle_debug_action) 105 | 106 | if not HAVE_IDA: 107 | self.options_menu.menuAction().setVisible(False) 108 | 109 | # Connect triggers 110 | self.save_dump_action.triggered.connect(self.save_dump) 111 | self.expand_all_action.triggered.connect(self.expand_all) 112 | self.remove_action.triggered.connect(self.remove_root_item) 113 | self.copy_action.triggered.connect(self.copy_to_clipboard) 114 | self.hexdump_action.triggered.connect(self.hexdump) 115 | self.disassemble_action.triggered.connect(self.disassemble) 116 | self.search_idb_action.triggered.connect(self.search_idb) 117 | self.map_pe_from_cursor_action.triggered.connect(self.map_pe_from_cursor) 118 | self.map_pe_from_file_action.triggered.connect(self.map_pe_from_file) 119 | self.save_action.triggered.connect(self.save_to_file) 120 | self.search_vt_action.triggered.connect(self.search_vt) 121 | self.cyberchef_action.triggered.connect(self.cyberchef) 122 | self.toggle_debug_action.triggered.connect(self.toggle_debug) 123 | self.rva_action.triggered.connect(self.toggle_rva) 124 | self.va_action.triggered.connect(self.toggle_va) 125 | self.about_action.triggered.connect(self.about_box) 126 | 127 | # Set altitudes 128 | self.altitudes[self.expand_all_action] = 0 129 | self.altitudes[self.save_dump_action] = 1 130 | self.altitudes[self.save_action] = 2 131 | self.altitudes[self.remove_action] = 3 132 | self.altitudes[self.copy_action] = 4 133 | self.altitudes[self.hexdump_action] = 5 134 | self.altitudes[self.disassemble_action] = 6 135 | self.altitudes[self.search_vt_action] = 7 136 | self.altitudes[self.cyberchef_action] = 8 137 | self.altitudes[self.add_pe_menu] = 9 138 | self.altitudes[self.addressing_menu] = 10 139 | self.altitudes[self.options_menu] = 11 140 | 141 | def show_menu(self): 142 | """Display context menu""" 143 | actions_by_altitude = {} 144 | 145 | # Determine altitude for actions 146 | for action in self.actions: 147 | actions_by_altitude[action] = self.altitudes[action] 148 | 149 | # Add actions/sub-menus to the menu by altitude 150 | visible = 0 151 | 152 | for action in sorted(actions_by_altitude, key=actions_by_altitude.get): 153 | if isinstance(action, QtWidgets.QAction): 154 | self.menu.addAction(action) 155 | visible += 1 if action.isVisible() else 0 156 | else: 157 | self.menu.addMenu(action) 158 | visible += 1 if action.menuAction().isVisible() else 0 159 | 160 | if HAVE_IDA: 161 | # Add about to IDA context menu 162 | self.menu.addSeparator() 163 | self.menu.addAction(self.about_action) 164 | visible += 1 165 | 166 | if visible == 0: 167 | return 168 | 169 | # Read the config file again 170 | self.form.runtime.read_config() 171 | 172 | # Set config checkboxes 173 | self.toggle_debug_action.setChecked(self.form.runtime.get_config_option("config", "debug", False)) 174 | 175 | # Set addressing checkboxes 176 | self.rva_action.setChecked(self.item.tree.show_rva) 177 | self.va_action.setChecked(not self.item.tree.show_rva) 178 | 179 | # Show the menu 180 | self.menu.exec_(self.point) 181 | 182 | def new_menu(self, form, point, item, index): 183 | """Initialise a new context menu upon right clicking the tree view 184 | 185 | Args: 186 | form (pe_tree.form.PETreeForm): PE Tree form 187 | point (QPoint): Mouse click co-ordinates 188 | item (pe_tree.tree.PETree): Item that was right-clicked 189 | index (QModelIndex): Item model index 190 | 191 | Returns: 192 | CommonStandardItemContextMenu: self 193 | 194 | """ 195 | self.menu = QtWidgets.QMenu() 196 | 197 | self.form = form 198 | self.point = point 199 | self.item = item 200 | self.index = index 201 | self.actions = [] 202 | 203 | return self 204 | 205 | def toggle_option(self, section, option): 206 | """Toggle boolean config option 207 | 208 | Args: 209 | section (str): Section name 210 | option (str): Option name 211 | 212 | """ 213 | self.form.runtime.set_config_option(section, option, str(not self.form.runtime.get_config_option(section, option, False))) 214 | 215 | def toggle_debug(self): 216 | """Enable/disable debug config option""" 217 | self.toggle_option("config", "debug") 218 | 219 | def toggle_rva(self): 220 | """Enable RVA addressing""" 221 | self.item.tree.show_rva = True 222 | 223 | def toggle_va(self): 224 | """Enable VA addressing""" 225 | self.item.tree.show_rva = False 226 | 227 | def save_dump(self): 228 | """Dump/save PE""" 229 | tree = self.item.tree 230 | pe_tree.dialogs.DumpPEForm(pefile.PE(data=tree.org_data), tree.image_base, tree.size, tree.filename, tree.ptr_size, self.form).invoke() 231 | 232 | def expand_all(self): 233 | """Expand all nodes beneath the selected node""" 234 | self.form.treeview.setExpanded(self.index, True) 235 | self.form.expanding = True 236 | self.form.expand_items(self.index) 237 | self.form.expanding = False 238 | self.form.treeview.resizeColumnToContents(0) 239 | self.form.treeview.resizeColumnToContents(1) 240 | 241 | def remove_root_item(self): 242 | """Remove PE file from tree""" 243 | tree = self.item.tree 244 | self.form.tree_roots.remove(tree) 245 | 246 | widget = self.form.map_stack.currentWidget() 247 | self.form.map_stack.removeWidget(tree.map) 248 | widget.deleteLater() 249 | 250 | if self.form.map_stack.count() == 0: 251 | self.form.map_stack.hide() 252 | 253 | if hasattr(self.form, "output_stack"): 254 | self.form.output_stack.removeWidget(tree.output_view) 255 | 256 | self.form.model.removeRow(self.index.row()) 257 | 258 | def copy_to_clipboard(self): 259 | """Copy item text to clipboard and print to IDA output""" 260 | Qt.QApplication.clipboard().setText(self.item.text()) 261 | self.item.tree.form.runtime.log(self.item.text()) 262 | 263 | def hexdump(self): 264 | """Print hexdump to IDA output""" 265 | self.item.tree.form.runtime.log(pe_tree.utils.hexdump(self.item.get_data(), self.item.tree.image_base + self.item.offset)) 266 | 267 | def disassemble(self): 268 | """Disassemble using capstone""" 269 | if self.item.tree.disasm: 270 | for i in self.item.tree.disasm.disasm(self.item.get_data(size=max(self.item.size, 0x100)), self.item.offset): 271 | self.item.tree.form.runtime.log("0x{:x}:\t{}\t{}".format(i.address, i.mnemonic, i.op_str)) 272 | 273 | def search_idb(self): 274 | """Search IDB for possible MZ/PE headers""" 275 | # Find all PE files in memory 276 | for image_base, section_name, is_64 in self.item.tree.form.runtime.find_pe(): 277 | # Map PE file 278 | self.form.runtime.log("0x{:0{w}x} - {}".format(image_base, section_name, w=16 if is_64 else 8)) 279 | self.form.map_pe(image_base=image_base, filename="{} - 0x{:0{w}x}".format(section_name, image_base, w=16 if is_64 else 8)) 280 | 281 | def map_pe_from_cursor(self): 282 | """Map PE file from current cursor position""" 283 | # Is the cursor on a PE file? 284 | for image_base, section_name, is_64 in self.item.tree.form.runtime.find_pe(cursor=True): 285 | # Map PE file 286 | self.form.map_pe(image_base=image_base, filename="{} - 0x{:0{w}x}".format(section_name, image_base, w=16 if is_64 else 8), priority=1) 287 | 288 | def map_pe_from_file(self): 289 | """Prompt user to select a file and attempt to map""" 290 | self.form.map_pe(filename=self.form.runtime.ask_file("*.*", "Open file")) 291 | 292 | def save_to_file(self): 293 | """Ask user where to save data""" 294 | filename = self.item.tree.form.runtime.ask_file(self.item.filename, "Save to file", save=True) 295 | if filename: 296 | with open(filename, "wb") as ofile: 297 | # Write data to file 298 | ofile.write(self.item.get_data()) 299 | 300 | def search_vt(self): 301 | """Search value using VirusTotal""" 302 | try: 303 | url = self.form.runtime.get_config_option("config", "virustotal_url", None) 304 | 305 | if url: 306 | webbrowser.open_new_tab("{}/{}".format(url, quote(quote(self.item.vt_query)))) 307 | except: 308 | pass 309 | 310 | def cyberchef(self): 311 | """Open data in CyberChef""" 312 | if sys.version_info > (3,): 313 | equals = b"=" 314 | else: 315 | equals = "=" 316 | 317 | try: 318 | url = self.form.runtime.get_config_option("config", "cyberchef_url", None) 319 | 320 | if url: 321 | webbrowser.open_new_tab("{}/#recipe=From_Hex('Auto'){}&input={}".format(url, self.item.cyberchef_recipe, base64.b64encode(binascii.hexlify(self.item.get_data())).rstrip(equals).decode("ascii"))) 322 | except: 323 | pass 324 | 325 | def about_box(self): 326 | """Display application about box""" 327 | self.item.tree.form.runtime.about_box() 328 | -------------------------------------------------------------------------------- /pe_tree/dialogs.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree application dialogs""" 18 | 19 | # pefile 20 | import pefile 21 | 22 | # Qt imports 23 | from PyQt5 import QtGui, QtWidgets, QtCore 24 | 25 | # PE Tree imports 26 | import pe_tree.form 27 | import pe_tree.dump_pe 28 | 29 | class DumpPEForm(QtWidgets.QDialog): 30 | """Display simple dialog for dumping PE files 31 | 32 | Args: 33 | pe (pefile.PE): Parsed PE file 34 | image_base (int): Base address of PE file in-memory 35 | size (int): Size of PE file in-memory 36 | filename (str): Name of PE file/memory segment 37 | ptr_size (int): Width of pointer in characters, 8 = 32-bit, 16 = 64-bit 38 | form (pe_tree.form.PETreeForm): PE Tree main form 39 | parent (QWidget, optional): Dialog parent widget 40 | 41 | """ 42 | def __init__(self, pe, image_base, size, filename, ptr_size, form, parent=None): 43 | super(DumpPEForm, self).__init__(parent) 44 | 45 | self.pe = pe 46 | self.size = size 47 | self.image_base = image_base 48 | self.filename = filename 49 | self.form = form 50 | self.runtime = form.runtime 51 | self.ptr_size = ptr_size 52 | 53 | self.iat_ptrs = [] 54 | 55 | self.setWindowTitle("Dump - {}".format(filename)) 56 | 57 | # Construct widgets 58 | self.image_base_edit = QtWidgets.QLineEdit("{:0{w}x}".format(image_base, w=self.ptr_size)) 59 | self.image_base_edit.setInputMask("H" * self.ptr_size) 60 | self.image_base_edit.setWhatsThis("New image-base for the dumped PE file") 61 | 62 | self.entry_point_edit = QtWidgets.QLineEdit("{:0{w}x}".format(pe.OPTIONAL_HEADER.AddressOfEntryPoint, w=self.ptr_size)) 63 | self.entry_point_edit.setInputMask("H" * self.ptr_size) 64 | self.entry_point_edit.setWhatsThis("New entry-point for the dumped PE file") 65 | 66 | self.size_of_optional_header_edit = QtWidgets.QLineEdit("{:0{w}x}".format(pe.FILE_HEADER.SizeOfOptionalHeader, w=self.ptr_size)) 67 | self.size_of_optional_header_edit.setInputMask("H" * self.ptr_size) 68 | self.size_of_optional_header_edit.setWhatsThis("FILE_HEADER.SizeOfOptionalHeader to use when parsing the PE file") 69 | 70 | self.iat_rva_edit = QtWidgets.QLineEdit("{:0{w}x}".format(pe.OPTIONAL_HEADER.DATA_DIRECTORY[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IAT"]].VirtualAddress, w=self.ptr_size)) 71 | self.iat_rva_edit.setInputMask("H" * self.ptr_size) 72 | self.iat_rva_edit.setWhatsThis("RVA of IAT to rebuild imports from.") 73 | self.iat_rva_edit.setDisabled(True) 74 | 75 | self.iat_size_edit = QtWidgets.QLineEdit("{:0{w}x}".format(pe.OPTIONAL_HEADER.DATA_DIRECTORY[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IAT"]].Size, w=self.ptr_size)) 76 | self.iat_size_edit.setInputMask("H" * self.ptr_size) 77 | self.iat_size_edit.setWhatsThis("Size of IAT to rebuild imports from.") 78 | self.iat_size_edit.setDisabled(True) 79 | 80 | self.use_existing_imports_radio = QtWidgets.QRadioButton("Rebuild IDT from existing IAT:") 81 | self.use_existing_imports_radio.setWhatsThis("Use existing IAT to construct a new IDT") 82 | self.use_existing_imports_radio.toggled.connect(self.toggle_iat_radios) 83 | 84 | self.no_imports_radio = QtWidgets.QRadioButton("Keep existing IDT/IAT") 85 | self.no_imports_radio.setWhatsThis("Do not rebuild IDT/IAT") 86 | self.no_imports_radio.toggled.connect(self.toggle_iat_radios) 87 | 88 | self.rebuild_imports_radio = QtWidgets.QRadioButton("Rebuild IDT/IAT") 89 | self.rebuild_imports_radio.setWhatsThis("Disassemble PE, search for IAT pointers and construct a new IDT/IAT") 90 | self.rebuild_imports_radio.toggled.connect(self.toggle_iat_radios) 91 | self.rebuild_imports_radio.setChecked(True) 92 | 93 | self.recalculate_pe_checksum = QtWidgets.QCheckBox("Recalculate PE header checksum (slow!)") 94 | self.recalculate_pe_checksum.setWhatsThis("Recalculate PE header checksum (slow!)") 95 | self.recalculate_pe_checksum.setChecked(False) 96 | 97 | self.save_to_file_checkbox = QtWidgets.QCheckBox("Save to file") 98 | self.save_to_file_checkbox.setWhatsThis("Save dumped PE file to disk") 99 | self.save_to_file_checkbox.setChecked(True) 100 | 101 | self.add_to_tree_checkbox = QtWidgets.QCheckBox("Add to PE Tree") 102 | self.add_to_tree_checkbox.setWhatsThis("Add dumped PE file to tree.\n\nWarning! Addresses may be incorrect in current IDA context if image is rebased!") 103 | self.add_to_tree_checkbox.setChecked(True) 104 | 105 | button_box = QtWidgets.QDialogButtonBox() 106 | 107 | dump_button = QtWidgets.QPushButton("Dump") 108 | dump_button.setFixedHeight(25) 109 | 110 | cancel_button = QtWidgets.QPushButton("Cancel") 111 | cancel_button.setFixedHeight(25) 112 | 113 | button_box.addButton(dump_button, QtWidgets.QDialogButtonBox.AcceptRole) 114 | button_box.addButton(cancel_button, QtWidgets.QDialogButtonBox.RejectRole) 115 | 116 | button_box.accepted.connect(self.accept) 117 | button_box.rejected.connect(self.reject) 118 | 119 | # Create layout 120 | layout = QtWidgets.QVBoxLayout() 121 | 122 | widgets = [ 123 | {"PE Headers": [ 124 | (QtWidgets.QLabel("ImageBase:"), self.image_base_edit), 125 | (QtWidgets.QLabel("AddressOfEntryPoint:"), self.entry_point_edit), 126 | (QtWidgets.QLabel("SizeOfOptionalHeader:"), self.size_of_optional_header_edit), 127 | (self.recalculate_pe_checksum, None) 128 | ]}, 129 | {"Imports": [ 130 | (self.no_imports_radio, None), 131 | (self.rebuild_imports_radio, None), 132 | (self.use_existing_imports_radio, None), 133 | (QtWidgets.QLabel("IAT:"), self.iat_rva_edit), 134 | (QtWidgets.QLabel("IAT Size:"), self.iat_size_edit) 135 | ]}, 136 | {"Output": [ 137 | (self.save_to_file_checkbox, None), 138 | (self.add_to_tree_checkbox, None) 139 | ]} 140 | ] 141 | 142 | # Construct groups 143 | for widget in widgets: 144 | group_name = list(widget)[0] 145 | 146 | groupbox = QtWidgets.QGroupBox(group_name) 147 | grid = QtWidgets.QGridLayout() 148 | grid.setSpacing(10) 149 | 150 | # Add group widgets 151 | for row, (left, right) in enumerate(widget[group_name]): 152 | if right == None: 153 | # Single widget, span 2 columns 154 | grid.addWidget(left, row, 0, 1, 2) 155 | else: 156 | # Pair of widgets 157 | grid.addWidget(left, row, 0) 158 | grid.addWidget(right, row, 1) 159 | 160 | groupbox.setLayout(grid) 161 | layout.addWidget(groupbox) 162 | 163 | layout.addWidget(button_box) 164 | 165 | self.setLayout(layout) 166 | 167 | def toggle_iat_radios(self): 168 | """IAT radio button toggle""" 169 | self.iat_rva_edit.setEnabled(self.use_existing_imports_radio.isChecked()) 170 | self.iat_size_edit.setEnabled(self.use_existing_imports_radio.isChecked()) 171 | 172 | def invoke(self): 173 | """Display the dump PE dialog""" 174 | # Has the user accepted, i.e. pressed "Dump"? 175 | if self.exec_() != 0: 176 | # Read values from edit controls 177 | image_base = int(self.image_base_edit.text(), 16) 178 | entry_point = int(self.entry_point_edit.text(), 16) 179 | size_of_optional_header = int(self.size_of_optional_header_edit.text(), 16) 180 | 181 | # Determine what the user wants to do about IAT reconstruction 182 | find_patch_iat = False 183 | 184 | if self.rebuild_imports_radio.isChecked(): 185 | # Scan/rebuild IAT 186 | find_patch_iat = True 187 | 188 | iat_rva = 0 189 | iat_size = 0 190 | 191 | if self.use_existing_imports_radio.isChecked(): 192 | # Use existing IAT 193 | iat_rva = int(self.iat_rva_edit.text(), 16) 194 | iat_size = int(self.iat_size_edit.text(), 16) 195 | 196 | # Dump the PE file 197 | pe_data = pe_tree.dump_pe.DumpPEFile(self.pe, 198 | self.image_base, 199 | self.size, 200 | self.runtime, 201 | find_iat_ptrs=find_patch_iat, 202 | keep_iat_ptrs=self.no_imports_radio.isChecked(), 203 | patch_iat_ptrs=find_patch_iat, 204 | iat_rva=iat_rva, 205 | iat_size=iat_size, 206 | new_image_base=image_base, 207 | new_ep=entry_point, 208 | recalculate_pe_checksum=self.recalculate_pe_checksum.isChecked(), 209 | size_of_optional_header=size_of_optional_header).dump() 210 | 211 | # Add to tree 212 | if self.add_to_tree_checkbox.isChecked(): 213 | self.form.map_pe(image_base=image_base, data=pe_data, priority=1, filename=self.filename, opaque=self.runtime.opaque, disable_dump=True) 214 | 215 | # Save to file 216 | if self.save_to_file_checkbox.isChecked(): 217 | filename = self.runtime.ask_file("{}.dmp".format(self.filename), "Save dump", save=True) 218 | if filename: 219 | with open(filename, "wb") as ofile: 220 | ofile.write(pe_data) 221 | 222 | self.close() 223 | 224 | class ModulePicker(QtWidgets.QDialog): 225 | """Process/module picker dialog""" 226 | def __init__(self, window, title, columns, model=QtGui.QStandardItemModel, parent=None): 227 | super(ModulePicker, self).__init__(parent) 228 | 229 | self.window = window 230 | self.items = set() 231 | 232 | # Initialise the dialog 233 | self.setWindowTitle(title) 234 | 235 | self.setMinimumWidth(400) 236 | self.setMinimumHeight(600) 237 | 238 | # Create the process/module tree view and model 239 | self.treeview = QtWidgets.QTreeView() 240 | self.model = model(self.treeview) 241 | self.model.setHorizontalHeaderLabels(columns) 242 | self.model.itemChanged.connect(self.item_changed) 243 | self.treeview.setModel(self.model) 244 | self.treeview.setSizeAdjustPolicy(QtWidgets.QAbstractScrollArea.AdjustToContents) 245 | self.treeview.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers) 246 | 247 | # Create the buttons 248 | button_box = QtWidgets.QDialogButtonBox() 249 | 250 | dump_button = QtWidgets.QPushButton("Dump") 251 | dump_button.setFixedHeight(25) 252 | 253 | cancel_button = QtWidgets.QPushButton("Cancel") 254 | cancel_button.setFixedHeight(25) 255 | 256 | button_box.addButton(dump_button, QtWidgets.QDialogButtonBox.AcceptRole) 257 | button_box.addButton(cancel_button, QtWidgets.QDialogButtonBox.RejectRole) 258 | 259 | button_box.accepted.connect(self.accept) 260 | button_box.rejected.connect(self.reject) 261 | 262 | # Create the layout 263 | grid = QtWidgets.QGridLayout() 264 | grid.setSpacing(10) 265 | 266 | grid.addWidget(self.treeview, 0, 0) 267 | grid.addWidget(button_box, 1, 0) 268 | 269 | self.setLayout(grid) 270 | 271 | self.populate_processes() 272 | 273 | def populate_processes(self): 274 | pass 275 | 276 | def item_changed(self, item): 277 | """Record selected items in the tree view""" 278 | if item.checkState() == QtCore.Qt.PartiallyChecked: 279 | # Add to list of selected items 280 | self.items.add(pe_tree.qstandarditems.HashableQStandardItem(item)) 281 | 282 | # If parent item then deselect all children 283 | if item.hasChildren(): 284 | for row in range(0, item.rowCount()): 285 | child = item.child(row) 286 | child.setCheckState(QtCore.Qt.Unchecked) 287 | self.items.discard(pe_tree.qstandarditems.HashableQStandardItem(child)) 288 | 289 | elif item.checkState() == QtCore.Qt.Checked: 290 | # Add to list of selected items 291 | self.items.add(pe_tree.qstandarditems.HashableQStandardItem(item)) 292 | 293 | # If parent item then select all children 294 | if item.hasChildren(): 295 | for row in range(0, item.rowCount()): 296 | child = item.child(row) 297 | child.setCheckState(QtCore.Qt.Checked) 298 | self.items.add(pe_tree.qstandarditems.HashableQStandardItem(child)) 299 | 300 | elif item.parent() is None: 301 | # Expand the root item 302 | self.treeview.setExpanded(item.index(), True) 303 | 304 | elif item.checkState() == QtCore.Qt.Unchecked: 305 | # Remove from list of selected items 306 | self.items.discard(pe_tree.qstandarditems.HashableQStandardItem(item)) 307 | 308 | # If parent item then deselect all children 309 | if item.hasChildren(): 310 | for row in range(0, item.rowCount()): 311 | child = item.child(row) 312 | child.setCheckState(QtCore.Qt.Unchecked) 313 | self.items.discard(pe_tree.qstandarditems.HashableQStandardItem(child)) 314 | 315 | def invoke(self): 316 | # Clear selection 317 | for item in self.items.copy(): 318 | item.item.setCheckState(QtCore.Qt.Unchecked) 319 | -------------------------------------------------------------------------------- /pe_tree/dump_pe.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """Dump PE file from memory/IDB and reconstruct import address table/hint name table/directory table""" 18 | 19 | # Standard imports 20 | import sys 21 | import struct 22 | 23 | # pefile 24 | import pefile 25 | 26 | # PE Tree imports 27 | import pe_tree.form 28 | 29 | class DumpPEFile(): 30 | """Dump PE file from memory and optionally rebuild imports. 31 | 32 | There are a number of ways to perform IAT reconstruction: 33 | 34 | 1. By default the IAT will be rebuilt from the IMAGE_DIRECTORY_ENTRY_IAT virtual address, if possible. 35 | 2. If `find_iat_ptrs/patch_iat_ptrs` is set then the image disassembly will be searched for IAT xrefs which are used to construct a new IAT. 36 | 3. It is possible to override the IMAGE_DIRECTORY_ENTRY_IAT virtual address and size using `iat_rva/iat_size`. This may be useful for malware that has set the IMAGE_DIRECTORY_ENTRY_IAT entry to 0 after loading. 37 | 4. Alternatively, it is possible to specify all IAT locations in the image via `iat_ptrs`, a list of tuples containing the IAT offset, xref, module name and API name 38 | 39 | Args: 40 | pe (pefile.PE): Parsed PE file 41 | image_base (int): Base address of PE file in-memory 42 | size (int): Size of PE file in-memory 43 | runtime (pe_tree.runtime.RuntimeSignals): Runtime callbacks 44 | find_iat_ptrs (bool, optional): Scan in-memory PE for IAT references (no need for iat_ptrs or iat_rva arguments) 45 | keep_iat_ptrs (bool, optional): Keep existing IAT/IDT 46 | iat_ptrs ([(int, int, str, str)], optional): Specify IAT pointers using tuple containing IAT offset, xref, module name and API name 47 | iat_rva (int, optional): Specify address of the import address table (iat_size must be > 0) 48 | iat_size (int, optional): Specify size of the import address table (iat_rva must be != 0) 49 | new_image_base (int, optional): New OPTIONAL_HEADER.ImageBase value 50 | new_ep (int, optional): New OPTIONAL_HEADER.AddressOfEntryPoint value 51 | size_of_optional_header_edit (int, optional): New FILE_HEADER.SizeOfOptionalHeader value 52 | recalculate_pe_checksum (bool, optional): Recalculate OPTIONAL_HEADER.CheckSum (warning, this is slow!) 53 | 54 | Returns: 55 | bytes: PE data 56 | 57 | """ 58 | def __init__(self, pe, image_base, size, runtime, **kwargs): 59 | self.pe = pe 60 | self.image_base = image_base 61 | self.runtime = runtime 62 | self.size = size 63 | 64 | # Get optional arguments 65 | self.iat_ptrs = kwargs.get("iat_ptrs", []) 66 | self.iat_rva = kwargs.get("iat_rva", 0) 67 | self.iat_size = kwargs.get("iat_size", 0) 68 | self.find_iat_ptrs = kwargs.get("find_iat_ptrs", False) 69 | self.keep_iat_ptrs = kwargs.get("keep_iat_ptrs", False) 70 | self.patch_iat_ptrs = kwargs.get("patch_iat_ptrs", False) 71 | self.new_image_base = kwargs.get("new_image_base", 0) 72 | self.new_ep = kwargs.get("new_ep", 0) 73 | self.size_of_optional_header = kwargs.get("size_of_optional_header", 0) 74 | self.recalculate_pe_checksum = kwargs.get("recalculate_pe_checksum", False) 75 | 76 | # Determine architecture 77 | if pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"]: 78 | self.ptr_size = 8 79 | self.ptr_format = " 0: 187 | # Build list of modules/API names 188 | modules = [] 189 | 190 | for iat_ptr, offset, module, api in self.iat_ptrs: 191 | if len(modules) == 0: 192 | modules.append({"name": module, "apis": [], "offset": offset}) 193 | 194 | if self.find_iat_ptrs: 195 | # Building our own IAT we can combine all modules/APIs 196 | found_module = False 197 | 198 | for _module in modules: 199 | if _module["name"] == module: 200 | if api not in _module["apis"]: 201 | _module["apis"].append(api) 202 | found_module = True 203 | break 204 | 205 | if not found_module: 206 | modules.append({"name": module, "apis": [api], "offset": offset}) 207 | else: 208 | # Using an existing IAT we need to keep modules/APIs in order 209 | prev_module = modules[-1] 210 | 211 | if prev_module["name"] == module: 212 | prev_module["apis"].append(api) 213 | else: 214 | modules.append({"name": module, "apis": [api], "offset": offset}) 215 | 216 | # Construct hint/name table 217 | name_table = bytearray() 218 | for module in modules: 219 | # Add module name and align 220 | name_table += module["name"].encode() + b"\0" 221 | name_table += (len(name_table) % 2) * b"\0" 222 | name_table += b"\0\0" 223 | for api in module["apis"]: 224 | # Add hint and API name and align 225 | name_table += b"\0\0" + api.encode() + b"\0" 226 | name_table += (len(name_table) % 2) * b"\0" 227 | 228 | # Construct IDT 229 | import_rva = self.align((pe.sections[-1].VirtualAddress + pe.sections[-1].Misc_VirtualSize), pe.OPTIONAL_HEADER.SectionAlignment) 230 | idt_data = bytearray() 231 | name_table_offset = 0 232 | name_table_len = len(name_table) 233 | name_to_iat = {} 234 | 235 | iat_index = 0 236 | 237 | if self.find_iat_ptrs: 238 | # IAT will be at the end of the image in a new section 239 | iat_base = import_rva 240 | else: 241 | # IAT remains wherever it is in the image 242 | iat_base = iat_rva 243 | 244 | for module in modules: 245 | #iat_index = module["offset"] 246 | 247 | # Add descriptor - OriginalFirstThunk, TimeDateStamp, ForwarderChain, Name, FirstThunk 248 | idt_data += struct.pack(" 0: 324 | # Add new .idata section to hold IAT + IDT + hint name table 325 | pe_data = pe.trim() 326 | 327 | # Create empty section 328 | section = pefile.SectionStructure(pe.__IMAGE_SECTION_HEADER_format__) 329 | section.__unpack__(bytearray(section.sizeof())) 330 | section.set_file_offset(pe.sections[-1].get_file_offset() + pe.sections[-1].sizeof()) 331 | 332 | # Fill in section details 333 | section.Name = b".pe_tree" 334 | section.Misc_VirtualSize = len(import_data) 335 | section.VirtualAddress = pe.sections[-1].VirtualAddress + self.align(pe.sections[-1].Misc_VirtualSize, pe.OPTIONAL_HEADER.SectionAlignment) 336 | section.SizeOfRawData = self.align(len(import_data), pe.OPTIONAL_HEADER.FileAlignment) 337 | section.PointerToRawData = self.align(len(pe_data), pe.OPTIONAL_HEADER.FileAlignment) 338 | section.Characteristics = pefile.SECTION_CHARACTERISTICS["IMAGE_SCN_CNT_INITIALIZED_DATA"] | pefile.SECTION_CHARACTERISTICS["IMAGE_SCN_MEM_READ"] | pefile.SECTION_CHARACTERISTICS["IMAGE_SCN_MEM_EXECUTE"] 339 | 340 | # Update PE headers 341 | pe.FILE_HEADER.NumberOfSections += 1 342 | pe.OPTIONAL_HEADER.SizeOfImage += section.VirtualAddress + self.align(len(import_data), pe.OPTIONAL_HEADER.SectionAlignment) 343 | 344 | # Append import data 345 | pe_data += (self.align(len(pe_data), pe.OPTIONAL_HEADER.FileAlignment) - len(pe_data)) * b"\0" 346 | pe_data += import_data 347 | pe_data += (self.align(len(pe_data), pe.OPTIONAL_HEADER.FileAlignment) - len(pe_data)) * b"\0" 348 | pe.__data__ = pe_data 349 | 350 | # Add section to pefile 351 | pe.sections.append(section) 352 | pe.__structures__.append(section) 353 | 354 | # Recalculate PE checksum if enabled (warning, very slow!) 355 | if self.recalculate_pe_checksum: 356 | pe.OPTIONAL_HEADER.CheckSum = pe.generate_checksum() 357 | 358 | # Put any overlay back 359 | pe.__data__ += overlay 360 | 361 | # Return data 362 | if sys.version_info > (3,): 363 | return pe.write() 364 | 365 | return "".join(map(chr, pe.write())) 366 | 367 | def align(self, val_to_align, alignment): 368 | """Align value 369 | 370 | Arguments: 371 | val_to_align (int): Value 372 | alignment (int): Alignment multiple 373 | 374 | Returns: 375 | int: Value rounded to alignment 376 | 377 | """ 378 | return ((val_to_align + alignment - 1) // alignment) * alignment 379 | -------------------------------------------------------------------------------- /pe_tree/exceptions.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree exceptions""" 18 | 19 | class ThreadStopping(Exception): 20 | """Exception raised when the application is stopping, used to terminate threads gracefully""" 21 | -------------------------------------------------------------------------------- /pe_tree/form.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """Main form for PE Tree""" 18 | 19 | # Standard imports 20 | import os 21 | import threading 22 | import multiprocessing 23 | import webbrowser 24 | 25 | # Requests 26 | import requests 27 | 28 | # Qt imports 29 | from PyQt5 import QtCore, QtGui, QtWidgets, Qt 30 | 31 | # IDAPro imports 32 | try: 33 | import idc # pylint: disable=unused-import 34 | 35 | HAVE_IDA = True 36 | except ImportError: 37 | HAVE_IDA = False 38 | 39 | # Rekall imports 40 | try: 41 | import rekall # pylint: disable=unused-import 42 | 43 | HAVE_REKALL = True 44 | except ImportError: 45 | HAVE_REKALL = False 46 | 47 | # Ghidra imports 48 | try: 49 | import ghidra_bridge # pylint: disable=unused-import 50 | 51 | HAVE_GHIDRA = True 52 | except ImportError: 53 | HAVE_GHIDRA = False 54 | 55 | # Volatility3 imports 56 | try: 57 | import volatility3 # pylint: disable=unused-import 58 | 59 | HAVE_VOLATILITY3 = True 60 | except ImportError: 61 | HAVE_VOLATILITY3 = False 62 | 63 | # PE Tree imports 64 | import pe_tree.tree 65 | import pe_tree.map 66 | import pe_tree.contextmenu 67 | import pe_tree.qstandarditems 68 | import pe_tree.utils 69 | 70 | class PETreeForm(): 71 | """PETree GUI and helper functions 72 | 73 | Args: 74 | widget (QWidget): Parent widget 75 | application (QApplication): Application widget, or None for IDA plugin 76 | runtime (pe_tree.runtime.Runtime): Runtime callbacks 77 | 78 | """ 79 | 80 | def __init__(self, widget, application, runtime): 81 | # Get parent widget and application 82 | if not application: 83 | application = widget 84 | 85 | self.widget = widget 86 | self.application = application 87 | self.runtime = runtime 88 | 89 | # Load main font 90 | self.font = runtime.get_available_font() 91 | 92 | # Load compiler IDs 93 | self.load_compiler_ids() 94 | 95 | # Load icons 96 | resources_dir = os.path.join(runtime.get_script_dir(), "resources") 97 | 98 | self.vt_icon = self.load_icon(os.path.join(resources_dir, "virustotal.png")) 99 | self.cyberchef_icon = self.load_icon(os.path.join(resources_dir, "cyberchef.ico")) 100 | 101 | # Load window icon 102 | self.application.setWindowIcon(self.load_icon(os.path.join(resources_dir, "spear.png"))) 103 | 104 | # Create tree view and model 105 | self.treeview = QtWidgets.QTreeView() 106 | self.model = QtGui.QStandardItemModel(self.treeview) 107 | self.model.setHorizontalHeaderLabels(["", ""]) 108 | 109 | # Setup tree view 110 | self.treeview.setModel(self.model) 111 | self.treeview.doubleClicked.connect(self.row_double_clicked) 112 | self.treeview.clicked.connect(self.row_clicked) 113 | self.treeview.setContextMenuPolicy(QtCore.Qt.CustomContextMenu) 114 | self.treeview.customContextMenuRequested.connect(self.context_menu) 115 | #self.treeview.setAlternatingRowColors(True) 116 | #self.treeview.setAutoScroll(True) 117 | self.treeview.setSizeAdjustPolicy(QtWidgets.QAbstractScrollArea.AdjustToContents) 118 | self.treeview.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers) 119 | self.treeview.setItemDelegate(pe_tree.tree.TreeViewDelegate(self)) 120 | self.treeview.expanded.connect(self.expanded) 121 | self.treeview.collapsed.connect(self.collapsed) 122 | self.treeview.selectionModel().selectionChanged.connect(self.selection_changed) 123 | self.treeview.setVisible(False) 124 | self.treeview.header().hide() 125 | 126 | # Create stack of PE maps 127 | self.map_stack = QtWidgets.QStackedWidget() 128 | 129 | # Right hand pane 130 | right_pane = QtWidgets.QSplitter() 131 | 132 | right_pane.addWidget(self.treeview) 133 | 134 | if not HAVE_IDA: 135 | self.output_stack = QtWidgets.QStackedWidget() 136 | 137 | right_pane.addWidget(self.output_stack) 138 | self.output_stack.setVisible(False) 139 | 140 | right_pane.setOrientation(QtCore.Qt.Vertical) 141 | right_pane.setSizes([300, 100]) 142 | 143 | # Create the splitter 144 | self.splitter = QtWidgets.QSplitter() 145 | self.splitter.addWidget(self.map_stack) 146 | self.splitter.addWidget(right_pane) 147 | 148 | self.splitter.setSizes([100, 300]) 149 | 150 | # Create the status bar widgets 151 | self.status_widget = QtWidgets.QWidget() 152 | self.status_widget.setVisible(False) 153 | self.status_widget.setFixedHeight(25) 154 | 155 | status_layout = QtWidgets.QHBoxLayout() 156 | self.status_label = QtWidgets.QLabel(self.status_widget) 157 | 158 | self.cancel_button = QtWidgets.QPushButton("Cancel", self.status_widget) 159 | self.cancel_button.setFixedHeight(25) 160 | self.cancel_button.setFixedWidth(100) 161 | self.cancel_button.clicked.connect(self.cancel_clicked) 162 | 163 | status_layout.addWidget(self.cancel_button) 164 | status_layout.addWidget(self.status_label) 165 | 166 | self.status_widget.setLayout(status_layout) 167 | self.status_label.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Fixed) 168 | 169 | status_layout.setContentsMargins(0, 0, 0, 0) 170 | self.status_label.setContentsMargins(0, 0, 0, 0) 171 | 172 | # Create layout and add widgets 173 | self.layout = QtWidgets.QVBoxLayout() 174 | self.layout.addWidget(self.splitter) 175 | self.layout.addWidget(self.status_widget) 176 | 177 | # Populate form with layout 178 | self.widget.setLayout(self.layout) 179 | 180 | # Construct context menu actions 181 | self.context_menu_actions = pe_tree.contextmenu.CommonStandardItemContextMenu(self) 182 | 183 | # Array of all PETree() root items 184 | self.tree_roots = [] 185 | 186 | self.expanding = False 187 | 188 | # Create thread pool, events and signals 189 | self.threadpool = QtCore.QThreadPool() 190 | self.threadpool.setMaxThreadCount(1) 191 | self.stop_event = threading.Event() 192 | self.signals = pe_tree.tree.PETreeSignals() 193 | self.signals.update_ui.connect(self.update_ui) 194 | self.signals.process_file.connect(self.process_file) 195 | self.dispatcher = None 196 | 197 | if not HAVE_IDA and not HAVE_REKALL: 198 | # Use process pool for parsing PE files 199 | self.processpool = multiprocessing.Pool(1) 200 | else: 201 | # Process pool does bad things under IDA! 202 | self.processpool = None 203 | 204 | if HAVE_IDA: 205 | # For IDA we'll have to wait for threads on-close like this 206 | self.widget.destroyed.connect(self.wait_for_threads) 207 | 208 | def wait_for_threads(self): 209 | """Wait for worker threads/processes to terminate before we close""" 210 | # Signal to threads that the application is closing 211 | self.stop_event.set() 212 | 213 | # Remove any queued threads 214 | self.threadpool.clear() 215 | 216 | # Process application events whilst we wait for all threads to complete 217 | if self.dispatcher is not None: 218 | while not self.threadpool.waitForDone(10): 219 | self.dispatcher.processEvents(Qt.QEventLoop.AllEvents) 220 | 221 | # Wait for processes to complete 222 | if self.processpool is not None: 223 | self.processpool.close() 224 | self.processpool.join() 225 | 226 | def cancel_clicked(self): 227 | """Abort all map/scan threads""" 228 | # Wait for all threads/processes to complete 229 | self.wait_for_threads() 230 | 231 | # Hide the progress bar 232 | self.status_widget.setVisible(False) 233 | 234 | # Recreate the process pool (cannot use again after close) 235 | if self.processpool is not None: 236 | self.processpool = multiprocessing.Pool(1) 237 | 238 | # Clear the stop event 239 | self.stop_event.clear() 240 | 241 | def load_compiler_ids(self): 242 | """Load and parse compiler IDs from https://github.com/dishather/richprint""" 243 | try: 244 | self.compiler_ids = None 245 | 246 | filepath = os.path.join(self.runtime.get_temp_dir(), "comp_id.txt") 247 | 248 | # Download compiler IDs 249 | if not os.path.exists(filepath): 250 | r = requests.get("https://raw.githubusercontent.com/dishather/richprint/master/comp_id.txt", allow_redirects=True) 251 | with open(filepath, "wb") as compiler_ids_file: 252 | compiler_ids_file.write(r.content) 253 | 254 | with open(filepath, "r") as compiler_ids_file: 255 | # Parse compiler IDs into dict 256 | self.compiler_ids = {} 257 | 258 | for line in compiler_ids_file.readlines(): 259 | if len(line) <= 2 or line[0] == "#": 260 | continue 261 | 262 | ids = line.split(" ") 263 | 264 | # Create description entry 265 | self.compiler_ids[ids[0]] = " ".join(ids[1:]).strip() 266 | except: 267 | pass 268 | 269 | def load_icon(self, file_path): 270 | """Load icon from file 271 | 272 | Args: 273 | file_path (str): Path to icon file 274 | 275 | Returns: 276 | QIcon: Icon loaded with image from file, otherwise a blank QIcon on error 277 | 278 | """ 279 | try: 280 | if os.path.exists(file_path): 281 | with open(file_path, "rb") as ficon: 282 | return pe_tree.utils.qicon_from_ico_data(ficon.read()) 283 | except: 284 | pass 285 | 286 | return QtGui.QIcon() 287 | 288 | def set_map_view(self, item): 289 | """Set the tree view's corresponding map view to active in the map stack 290 | 291 | Args: 292 | item (pe_tree.tree.PETree): Active tree view 293 | 294 | """ 295 | if item.tree.form.map_stack.currentIndex() == item.tree.map_index: 296 | return 297 | 298 | item.tree.form.map_stack.setCurrentIndex(item.tree.map_index) 299 | 300 | if hasattr(item.tree.form, "output_stack"): 301 | item.tree.form.output_stack.setCurrentIndex(item.tree.output_view_index) 302 | 303 | # Hide the output stack if the current text edit is empty 304 | if item.tree.form.output_stack.currentWidget().toPlainText() != "": 305 | self.output_stack.setVisible(True) 306 | else: 307 | self.output_stack.setVisible(False) 308 | 309 | def expanded(self, index): 310 | """Update map view if tree node expanded 311 | 312 | Args: 313 | index (QModelIndex): Item model index 314 | 315 | """ 316 | self.set_map_view(self.model.itemFromIndex(index)) 317 | 318 | # Did we trigger the expand? If so, ignore. 319 | if not self.expanding: 320 | # Resize the column width to fit the content 321 | self.treeview.resizeColumnToContents(0) 322 | self.treeview.resizeColumnToContents(1) 323 | 324 | def collapsed(self, index): 325 | """Update map view if tree node collapsed 326 | 327 | Args: 328 | index (QModelIndex): Item model index 329 | 330 | """ 331 | self.set_map_view(self.model.itemFromIndex(index)) 332 | 333 | # Resize the column width to fit the content 334 | self.treeview.resizeColumnToContents(0) 335 | self.treeview.resizeColumnToContents(1) 336 | 337 | def row_clicked(self, index): 338 | """Display URL associated with item 339 | 340 | Args: 341 | index (QModelIndex): Item model index 342 | 343 | """ 344 | item = self.model.itemFromIndex(index.sibling(index.row(), index.column())) 345 | 346 | if item.url: 347 | webbrowser.open_new_tab(item.url) 348 | 349 | def row_double_clicked(self, index): 350 | """Navigate to item offset in IDA view 351 | 352 | Args: 353 | index (QModelIndex): Item model index 354 | 355 | """ 356 | item = self.model.itemFromIndex(index.sibling(index.row(), 1)) 357 | 358 | if item.offset == 0: 359 | return 360 | 361 | # Convert RVA to VA 362 | if item.offset < item.tree.image_base: 363 | offset = item.tree.image_base + item.offset 364 | else: 365 | offset = item.offset 366 | 367 | # Attempt to navigate to address in IDA view 368 | item.tree.form.runtime.jumpto(item, offset) 369 | 370 | def context_menu(self, point): 371 | """Tree view right-click context menu activated 372 | 373 | Args: 374 | point (QPoint): Mouse click co-ordinates 375 | 376 | """ 377 | index = self.treeview.indexAt(point) 378 | 379 | # Find item that was right-clicked 380 | item = self.model.itemFromIndex(index) 381 | 382 | if item is None: 383 | item = pe_tree.qstandarditems.NoItem(pe_tree.tree.PETree(self, "", 0, None)) 384 | 385 | # Display the item's context menu 386 | item.context_menu(self.context_menu_actions.new_menu(self, self.treeview.viewport().mapToGlobal(point), item, index)) 387 | 388 | def selection_changed(self, selection): 389 | """Change PE map in map stack 390 | 391 | Args: 392 | selection (QItemSelectionModel): PE tree view selection 393 | 394 | """ 395 | # Find selected item 396 | for index in selection.indexes(): 397 | # Update map view 398 | item = self.model.itemFromIndex(index) 399 | form = item.tree.form 400 | 401 | item.tree.form.map_stack.setCurrentIndex(item.tree.map_index) 402 | 403 | if hasattr(form, "output_stack"): 404 | item.tree.form.output_stack.setCurrentIndex(item.tree.output_view_index) 405 | 406 | # Hide the output stack if the current text edit is empty 407 | if item.tree.form.output_stack.currentWidget().toPlainText() != "": 408 | self.output_stack.setVisible(True) 409 | else: 410 | self.output_stack.setVisible(False) 411 | 412 | return 413 | 414 | def expand_items(self, index): 415 | """Expand all children in a tree view from given index (recursive) 416 | 417 | Args: 418 | index (QModelIndex): Item model index 419 | 420 | """ 421 | # Ensure the index is valid 422 | if not index.isValid(): 423 | return 424 | 425 | # Expand all children 426 | for i in range(self.model.rowCount(index)): 427 | self.expand_items(index.child(i, 0)) 428 | 429 | # Expand current index 430 | self.treeview.setExpanded(index, True) 431 | 432 | def process_file(self, filename): 433 | """Signaled via worker thread to update status bar UI. 434 | 435 | Args: 436 | filename (str): The path of the current file being processed 437 | 438 | """ 439 | self.status_widget.setVisible(True) 440 | self.status_label.setText("Processing {}".format(filename)) 441 | 442 | def update_ui(self, tree): 443 | """Add tree-view and map widgets to the main form. Signaled via worker thread. 444 | 445 | Args: 446 | tree (pe_tree.tree.PETree): New PETree to add to the form view 447 | 448 | """ 449 | # Hide the status bar if no more threads are running 450 | if self.threadpool.activeThreadCount() == 0: 451 | self.status_widget.setVisible(False) 452 | 453 | # Have we been signaled to stop? 454 | if self.stop_event.is_set(): 455 | return 456 | 457 | # Ensure we have root items 458 | if not tree or not tree.root_item or not tree.root_value_item: 459 | return 460 | 461 | # Get the "invisible root" for the tree model 462 | invisible_root_item = self.model.invisibleRootItem() 463 | has_children = invisible_root_item.hasChildren() 464 | 465 | invisible_root_item.appendRow([tree.root_item, tree.root_value_item]) 466 | 467 | # Create new output window (if needed) 468 | if hasattr(self, "output_stack"): 469 | tree.output_view = QtWidgets.QTextEdit() 470 | tree.output_view.setFont(self.font) 471 | 472 | tree.output_view_index = self.output_stack.addWidget(tree.output_view) 473 | 474 | if not has_children: 475 | # Set the output widgets as active 476 | self.output_stack.setCurrentIndex(tree.output_view_index) 477 | 478 | # Create the PE map widget and scroll area and add to stack 479 | tree.map = pe_tree.map.PEMap(tree.size) 480 | tree.map.file_size = tree.size 481 | 482 | tree.map_scroll_area = pe_tree.map.PEMapScrollArea() 483 | tree.map_scroll_area.setWidgetResizable(True) 484 | tree.map_scroll_area.setAutoFillBackground(True) 485 | tree.map_scroll_area.setWidget(tree.map) 486 | 487 | tree.map_index = self.map_stack.addWidget(tree.map_scroll_area) 488 | 489 | # Add the file regions to the PE map widget 490 | tree.map.add_regions(tree.regions) 491 | 492 | # Append this PETree object to the list of roots 493 | self.tree_roots.append(tree) 494 | 495 | # Is this the first PE file? 496 | if not has_children: 497 | # Set the map widget as active 498 | self.map_stack.setCurrentIndex(tree.map_index) 499 | 500 | self.treeview.setVisible(True) 501 | 502 | # Ensure the filename fits in the column 503 | self.treeview.resizeColumnToContents(0) 504 | 505 | def map_pe(self, filename=None, image_base=0, data=None, disable_dump=False, priority=0, opaque=None): 506 | """Starts a new thread to map PE from file/memory/data 507 | 508 | Args: 509 | filename (str, optional): Path to PE file to map 510 | image_base (int, optional): Image base of PE file to map 511 | data (bytes, optional): PE File data to map 512 | disable_dump (bool): Disable dumping of the PE file 513 | priority (int, optional): QRunnable thread priority 514 | opaque (object): Opaque object pointer passed to runtime 515 | 516 | """ 517 | self.threadpool.start(pe_tree.tree.PETree(self, filename, image_base, data, disable_dump=disable_dump, opaque=opaque), priority) 518 | 519 | def show(self): 520 | """Display the main form widget""" 521 | self.runtime.show_widget() 522 | -------------------------------------------------------------------------------- /pe_tree/ghidra.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2021 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree Ghidra application""" 18 | 19 | # Standard imports 20 | import os 21 | import sys 22 | from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter 23 | 24 | # Qt imports 25 | from PyQt5 import QtCore, QtWidgets 26 | 27 | # Ghidra imports 28 | try: 29 | import ghidra_bridge 30 | except ModuleNotFoundError: 31 | print("ghidra_bridge is required") 32 | exit() 33 | 34 | import pe_tree.window 35 | import pe_tree.runtime 36 | import pe_tree.form 37 | import pe_tree.info 38 | 39 | class GhidraRuntime(pe_tree.runtime.Runtime): 40 | """Ghidra runtime callbacks""" 41 | def __init__(self, widget, args): 42 | # Load configuration 43 | self.config_file = os.path.join(self.get_temp_dir(), "pe_tree_ghidra.ini") 44 | super(GhidraRuntime, self).__init__(widget, args) 45 | 46 | # Initialise Ghidra bridge 47 | try: 48 | self.bridge = ghidra_bridge.GhidraBridge(connect_to_host=self.args.server, connect_to_port=self.args.port, namespace=globals(), response_timeout=20) 49 | except ConnectionRefusedError: 50 | print("Please run ghidra_bridge_server_background.py under Ghidra -> Window -> Script Manager") 51 | exit() 52 | 53 | state = getState() 54 | currentProgram = state.getCurrentProgram() 55 | 56 | self.address_factory = currentProgram.getAddressFactory() 57 | self.function_manager = currentProgram.getFunctionManager() 58 | self.data_type_manager = currentProgram.getDataTypeManager() 59 | self.listing = currentProgram.getListing() 60 | self.memory = currentProgram.getMemory() 61 | self.address_space = currentProgram.getAddressFactory().getDefaultAddressSpace() 62 | 63 | def jumpto(self, item, offset): 64 | """Jump to offset in Ghidra listing view""" 65 | try: 66 | self.ret = goTo(currentProgram.addressFactory.getAddress(hex(offset))) 67 | except: 68 | self.ret = False 69 | 70 | return self.ret 71 | 72 | @QtCore.pyqtSlot(object, object) 73 | def get_bytes(self, start, size): 74 | """Read bytes from memory""" 75 | array = self.bridge.remote_import("array") 76 | buffer = self.bridge.remote_eval("array.array('b', data)", array=array, data=b"\0" * size) 77 | 78 | read = self.memory.getBytes(self.address_space.getAddress(hex(start)), buffer) 79 | 80 | if read > 0: 81 | self.ret = bytearray(self.bridge.remote_eval("[ b & 0xff for b in data]", data=buffer)) 82 | else: 83 | self.ret = b"" 84 | 85 | return self.ret 86 | 87 | @QtCore.pyqtSlot(object) 88 | def get_name(self, offset): 89 | """Get name for offset from Ghidra""" 90 | self.ret = "" 91 | return self.ret 92 | 93 | @QtCore.pyqtSlot(object) 94 | def get_segment_name(self, offset): 95 | """Get segment name for offset from Ghidra""" 96 | self.ret = "" 97 | return self.ret 98 | 99 | @QtCore.pyqtSlot(object) 100 | def is_writable(self, offset): 101 | """Determine if memory address is writable""" 102 | self.ret = False 103 | return self.ret 104 | 105 | @QtCore.pyqtSlot(object) 106 | def resolve_address(self, offset): 107 | """Get module/symbol name for address""" 108 | symbol = "" 109 | module = "" 110 | 111 | ref = currentProgram.referenceManager.getReferencesFrom(self.address_factory.getAddress(hex(offset))) 112 | 113 | if ref: 114 | ref = ref[0].toString() 115 | 116 | # Should look something like.. "->MSVCRT.DLL::__setusermatherr" 117 | if ref.startswith("->"): 118 | # Remove -> 119 | ref = ref[2:] 120 | 121 | # Split into module/API 122 | elements = ref.split("::") 123 | 124 | module = elements[0] 125 | symbol = elements[1].split(" ")[0] 126 | 127 | self.ret = (module, symbol) 128 | return self.ret 129 | 130 | @QtCore.pyqtSlot(object) 131 | def get_label(self, offset): 132 | """Get label for offest from Ghidra""" 133 | self.ret = getSymbolAt(self.address_space.getAddress(hex(offset))) 134 | return self.ret 135 | 136 | @QtCore.pyqtSlot(object, object) 137 | def make_string(self, offset, size): 138 | """Convert range to byte string in Ghidra""" 139 | self.ret = self.listing.createData(self.address_space.getAddress(hex(offset)), ghidra.program.model.data.StringDataType, size) 140 | return self.ret 141 | 142 | @QtCore.pyqtSlot(object, object, str, str, bytes) 143 | def make_segment(self, offset, size, class_name="DATA", name="pe_map", data=None): 144 | """Create a new section in Ghidra - Not implemented""" 145 | self.ret = None 146 | return self.ret 147 | 148 | @QtCore.pyqtSlot(object) 149 | def make_qword(self, offset): 150 | """Create a qword at the given offset in the listing""" 151 | self.ret = self.listing.createData(self.address_space.getAddress(hex(offset)), ghidra.program.model.data.QWordDataType) 152 | return self.ret 153 | 154 | @QtCore.pyqtSlot(object) 155 | def make_dword(self, offset): 156 | """Create a dword at the given offset in the listing""" 157 | self.ret = self.listing.createData(self.address_space.getAddress(hex(offset)), ghidra.program.model.data.DWordDataType) 158 | return self.ret 159 | 160 | @QtCore.pyqtSlot(object) 161 | def make_word(self, offset): 162 | """Create a word at the given offset in the listing""" 163 | self.ret = self.listing.createData(self.address_space.getAddress(hex(offset)), ghidra.program.model.data.WordDataType) 164 | return self.ret 165 | 166 | @QtCore.pyqtSlot(object) 167 | def make_byte(self, offset, size=1): 168 | """Create a byte at the given offset in the listing""" 169 | self.ret = self.listing.createData(self.address_space.getAddress(hex(offset)), ghidra.program.model.data.ByteDataType) 170 | return self.ret 171 | 172 | @QtCore.pyqtSlot(object, str) 173 | def make_comment(self, offset, comment): 174 | """Create a comment at the given offset in the listing""" 175 | self.ret = self.listing.setComment(self.address_space.getAddress(hex(offset)), ghidra.program.model.listing.CodeUnit.REPEATABLE_COMMENT, comment) 176 | return self.ret 177 | 178 | def find_iat_ptrs_remote(self, image_base, size, get_word, namespace=globals()): 179 | """Find all likely IAT pointers 180 | 181 | Note! For performance reasons this code should be executed remotely under the Ghidra python interpreter using bridge.remoteify() 182 | 183 | """ 184 | address_space = currentProgram.getAddressFactory().getDefaultAddressSpace() 185 | 186 | address_set = ghidra.program.model.address.AddressSet() 187 | address_set.addRange(address_space.getAddress(hex(image_base)), address_space.getAddress(hex(image_base + size))) 188 | 189 | opcode_iterator = currentProgram.getListing().getInstructions(address_set, True) 190 | 191 | # Iterate over all opcodes in the image 192 | iat_ptrs = [] 193 | 194 | while opcode_iterator.hasNext(): 195 | opcode = opcode_iterator.next() 196 | 197 | # Get the opcode mnemonic 198 | mnem = opcode.getMnemonicString().lower() 199 | 200 | # Does the opcode contain a memory address? 201 | ptr = 0 202 | 203 | if mnem in ["call", "push", "jmp"]: 204 | if opcode.getOperandType(0) & ghidra.program.model.lang.OperandType.ADDRESS: 205 | # Get memory offset for branch instructions 206 | ptr = opcode.getOpObjects(0) 207 | 208 | elif mnem in ["mov", "lea"]: 209 | if opcode.getOperandType(0) & ghidra.program.model.lang.OperandType.REGISTER and opcode.getOperandType(1) & ghidra.program.model.lang.OperandType.ADDRESS: 210 | # Get memory offset for mov/lea instructions 211 | ptr = opcode.getOpObjects(1) 212 | 213 | if ptr == 0: 214 | continue 215 | 216 | # Get the memory address value/offset 217 | try: 218 | ptr = int(ptr[0].getUnsignedValue()) 219 | except: 220 | ptr = int(ptr[0].getOffset()) 221 | 222 | # Does the instruction's memory address operand seem somewhat valid?! 223 | if ptr < 0x1000: 224 | continue 225 | 226 | # Read pointer from memory address 227 | try: 228 | # TODO: Ensure this is dynamic for x86/64! 229 | iat_offset = getInt(currentProgram.addressFactory.getAddress(hex(ptr))) 230 | except: 231 | iat_offset = 0 232 | 233 | if iat_offset == 0: 234 | continue 235 | 236 | # Ignore offset if it is in our image 237 | if iat_offset >= image_base and iat_offset < image_base + size: 238 | continue 239 | 240 | # Attempt to resolve reference (Note! do not use self.resolve_address, as this code is executed remotely ;) 241 | ref = currentProgram.referenceManager.getReferencesFrom(currentProgram.addressFactory.getAddress(hex(ptr))) 242 | 243 | if ref: 244 | ref = ref[0].toString() 245 | 246 | # Should look something like.. "->MSVCRT.DLL::__setusermatherr" 247 | if ref.startswith("->"): 248 | # Remove -> 249 | ref = ref[2:] 250 | 251 | # Split into module/API 252 | elements = ref.split("::") 253 | 254 | module = elements[0] 255 | api = elements[1].split(" ")[0] 256 | 257 | # Add IAT offset, address to patch, module name and API name to list 258 | iat_ptrs.append([iat_offset, ptr + opcode.getAddress().add(opcode.getLength() - 4).getOffset(), module, api]) 259 | 260 | return iat_ptrs 261 | 262 | @QtCore.pyqtSlot(object, object, object, object) 263 | def find_iat_ptrs(self, pe, image_base, size, get_word): 264 | """Find all likely IAT pointers""" 265 | # Perform IAT search under Ghidra python interpreter 266 | self.ret = self.bridge.remoteify(self.find_iat_ptrs_remote)(self, image_base, size, get_word) 267 | return self.ret 268 | 269 | def main(): 270 | """PE Tree Ghidra script entry-point""" 271 | # Check command line arguments 272 | parser = ArgumentParser(description="PE-Tree (Ghidra)", formatter_class=ArgumentDefaultsHelpFormatter) 273 | parser.add_argument("--server", help="Ghidra bridge server IP", default="127.0.0.1") 274 | parser.add_argument("--port", help="Ghidra bridge server port", default=4768, type=int) 275 | args = parser.parse_args() 276 | 277 | # Create PE Tree Qt application 278 | application = QtWidgets.QApplication(sys.argv) 279 | window = pe_tree.window.PETreeWindow(application, GhidraRuntime, args, open_file=False) 280 | 281 | # Try to locate the input file 282 | filename = currentProgram.getExecutablePath() 283 | 284 | if filename: 285 | # No input file found? 286 | if not os.path.isfile(filename): 287 | filename = os.path.basename(filename) 288 | 289 | if not os.path.isfile(filename): 290 | filename = runtime.ask_file(os.path.basename(filename), "Where is the input file?") 291 | 292 | if os.path.isfile(filename): 293 | # Map input file 294 | window.pe_tree_form.map_pe(image_base=currentProgram.getImageBase().getOffset(), filename=filename) 295 | 296 | sys.exit(application.exec_()) 297 | 298 | if __name__ == "__main__": 299 | main() -------------------------------------------------------------------------------- /pe_tree/hash_pe.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """Hash all portions of a PE file, this may be run as a separate process""" 18 | 19 | # Standard imports 20 | import json 21 | import hashlib 22 | 23 | import pefile 24 | 25 | def hash_data(data, entropy_H): 26 | """Calculate MD5/SHA1/SHA256 of given data 27 | 28 | Args: 29 | data (bytes): Data to calculate hashes/entropy 30 | entropy_H (pefile.SectionStructure.entropy_H): Callback function for calculating entropy 31 | 32 | Returns: 33 | dict: Dictionary of hashes/entropy 34 | 35 | """ 36 | md5 = hashlib.md5() 37 | sha1 = hashlib.sha1() 38 | sha256 = hashlib.sha256() 39 | 40 | md5.update(data) 41 | sha1.update(data) 42 | sha256.update(data) 43 | 44 | return {"md5": md5.hexdigest(), "sha1": sha1.hexdigest(), "sha256": sha256.hexdigest(), "entropy": entropy_H(data), "size": len(data)} 45 | 46 | def hash_pe_file(filename, data=None, pe=None, json_dumps=True): 47 | """Calculate PE file hashes. 48 | 49 | Either call directly or invoke via processpool:: 50 | 51 | processpool = multiprocessing.Pool(10) 52 | hashes = json.loads(processpool.apply_async(pe_tree.hash_pe.hash_pe_file, (filename,)).get()) 53 | 54 | Args: 55 | filename (str): Path to file to hash (or specify via data) 56 | data (bytes, optional): PE file data 57 | pe (pefile.PE, optional): Parsed PE file 58 | json_dumps (bool, optional): Return data as JSON 59 | 60 | Returns: 61 | dict: PE file hashes if json_dumps == False 62 | str: JSON PE file hashes if json_dumps == True 63 | 64 | """ 65 | if pe is None: 66 | pe = pefile.PE(filename) 67 | 68 | # Calculate entropy (use pefile implementation!) 69 | entropy_H = pefile.SectionStructure(pe.__IMAGE_SECTION_HEADER_format__, pe=pe).entropy_H 70 | 71 | file_hashes = {"file": {"md5": "", "sha1": "", "sha256": "", "entropy": 0.0, "size": 0}, 72 | "file_no_overlay": {"md5": "", "sha1": "", "sha256": "", "entropy": 0.0, "size": 0}, 73 | "dos_stub": {"md5": "", "sha1": "", "sha256": "", "entropy": 0.0, "size": 0}, 74 | "sections": [], 75 | "resources": [], 76 | "security_directory": {"md5": "", "sha1": "", "sha256": "", "entropy": 0.0, "size": 0}, 77 | "overlay": {"md5": "", "sha1": "", "sha256": "", "entropy": 0.0, "size": 0}} 78 | 79 | if not data: 80 | with open(filename, "rb") as f: 81 | data = f.read() 82 | 83 | # Hash entire file 84 | file_hashes["file"] = hash_data(data, entropy_H) 85 | 86 | # Hash DOS stub 87 | if pe.DOS_HEADER.e_lfanew > 64: 88 | file_hashes["dos_stub"] = hash_data(data[64:pe.DOS_HEADER.e_lfanew], entropy_H) 89 | 90 | # Hash sections 91 | for section in pe.sections: 92 | file_hashes["sections"].append({"md5": section.get_hash_md5(), "sha256": section.get_hash_sha256(), "entropy": section.get_entropy()}) 93 | 94 | # Hash resources 95 | if hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"): 96 | mapped_data = pe.get_memory_mapped_image() 97 | 98 | for resource_type in pe.DIRECTORY_ENTRY_RESOURCE.entries: 99 | if not hasattr(resource_type, "directory"): 100 | continue 101 | 102 | for resource_id in resource_type.directory.entries: 103 | if not hasattr(resource_id, "directory"): 104 | continue 105 | 106 | for resource_language in resource_id.directory.entries: 107 | if not hasattr(resource_language, "data"): 108 | continue 109 | 110 | offset = resource_language.data.struct.OffsetToData 111 | size = resource_language.data.struct.Size 112 | 113 | try: 114 | resource_data = mapped_data[offset:offset + size] 115 | except: 116 | resource_data = "" 117 | 118 | file_hashes["resources"].append(hash_data(resource_data, entropy_H)) 119 | 120 | overlay_offset = pe.get_overlay_data_start_offset() 121 | 122 | if overlay_offset: 123 | overlay_data = pe.get_overlay() 124 | 125 | security = pe.OPTIONAL_HEADER.DATA_DIRECTORY[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_SECURITY"]] 126 | 127 | if security.VirtualAddress != 0 and security.Size != 0: 128 | size = min(security.Size, len(overlay_data)) 129 | 130 | # Hash security directory 131 | file_hashes["security_directory"] = hash_data(overlay_data[:size], entropy_H) 132 | 133 | overlay_data = overlay_data[size:] 134 | overlay_offset += size 135 | 136 | # Hash overlay 137 | file_hashes["overlay"] = hash_data(overlay_data, entropy_H) 138 | file_hashes["file_no_overlay"] = hash_data(data[overlay_offset:], entropy_H) 139 | 140 | # Return JSON 141 | if json_dumps: 142 | return json.dumps(file_hashes) 143 | 144 | # Return dict 145 | return file_hashes 146 | -------------------------------------------------------------------------------- /pe_tree/info.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # 4 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # http://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | # 18 | 19 | """PE Tree information""" 20 | 21 | __author__ = "Tom Bonner" 22 | __copyright__ = "Copyright © 2021 BlackBerry Limited." 23 | __version__ = "1.0.30" 24 | __maintainer__ = "Tom Bonner" 25 | __email__ = "tbonner@blackberry.com" 26 | __license__ = "Apache License 2.0" 27 | __url__ = "https://github.com/blackberry/pe_tree" 28 | __title__ = "PE Tree" 29 | -------------------------------------------------------------------------------- /pe_tree/map.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree file region map""" 18 | 19 | # Standard imports 20 | import sys 21 | import itertools 22 | 23 | # Qt imports 24 | from PyQt5 import QtCore, QtGui, QtWidgets 25 | 26 | class FileRegion(): 27 | """Holds information about a file region for the map view 28 | 29 | Args: 30 | name (str): Name of the region 31 | start (int, optional): File offset 32 | size (int, optional): Size in bytes 33 | rva (int, optional): In-memory offset 34 | item (pe_tree.tree.PETree, optional): PE Tree 35 | 36 | """ 37 | def __init__(self, name, start=0, size=0, end=0, rva=0, item=None): 38 | if sys.version_info > (3,): 39 | self.name = name 40 | else: 41 | self.name = str(name) 42 | 43 | self.start = start 44 | self.size = size 45 | self.end = end 46 | self.rva = rva 47 | self.item = item 48 | self.rect = None 49 | self.colour = None 50 | self.hover = False 51 | 52 | if self.end == 0: 53 | self.end = self.start + self.size 54 | 55 | if self.size == 0: 56 | self.size = self.end - self.start 57 | 58 | class PEMap(QtWidgets.QGroupBox): 59 | """PE map group box widget for holding region labels 60 | 61 | Args: 62 | file_size (int): Size of the PE file 63 | parent (QWidget, optional): Parent widget 64 | 65 | """ 66 | def __init__(self, file_size, parent=None): 67 | super(PEMap, self).__init__(parent=parent) 68 | 69 | self.file_size = file_size 70 | self.colours = None 71 | self.rainbows = {} 72 | 73 | # Create group box layout for the PE map 74 | self.layout = QtWidgets.QVBoxLayout(self) 75 | 76 | # Remove space between widgets 77 | self.layout.setContentsMargins(0, 0, 0, 0) 78 | self.layout.setSpacing(0) 79 | 80 | # Set background to white (for alpha) 81 | self.setStyleSheet("background-color: white; margin-top: 0px; padding: 0px;") 82 | 83 | def add_regions(self, regions): 84 | """Add list of file regions to the PE map 85 | 86 | Args: 87 | regions ([pe_tree.map.FileRegions]: File regions 88 | 89 | """ 90 | # Rainbow colours for the PE map 91 | colours = ["#b9413c", "#ea8c2c", "#ffbd20", "#559f7a", "#4c82a4", "#764a73", "#2f0128"] 92 | 93 | wanted = len(regions) + len(colours) 94 | 95 | if wanted in self.rainbows: 96 | # Use saved colours 97 | colours = self.rainbows[wanted] 98 | else: 99 | # Make a rainbow 100 | while True: 101 | gradients = polylinear_gradient(colours, wanted) 102 | 103 | if len(gradients["hex"]) >= len(regions): 104 | break 105 | 106 | wanted += 1 107 | 108 | colours = [] 109 | 110 | for colour in gradients["hex"]: 111 | colours.append(QtGui.QColor(int(colour[1:], 16))) 112 | 113 | # Save colours 114 | self.rainbows[len(colours)] = colours 115 | 116 | self.colours = itertools.cycle(colours) 117 | 118 | # Sort regions by start offset, create labels and add to the group box layout 119 | for region in sorted(regions, key=lambda region: (region.start)): 120 | self.layout.addWidget(PEMapLabel(region, next(self.colours), self.file_size, parent=self)) 121 | 122 | class PEMapLabel(QtWidgets.QWidget): 123 | """PE map label widget 124 | 125 | Args: 126 | region (pe_tree.map.FileRegion): PE file region 127 | colour (QColor): Region background colour 128 | file_size (int): Size of the PE file 129 | parent (QWidget, optional): Parent widget 130 | 131 | """ 132 | def __init__(self, region, colour, file_size, parent=None): 133 | super(PEMapLabel, self).__init__(parent=parent) 134 | 135 | self.region = region 136 | self.colour = colour 137 | self.file_size = file_size 138 | 139 | # Initialise self 140 | self.setMouseTracking(True) 141 | self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu) 142 | self.customContextMenuRequested.connect(self.context_menu) 143 | self.setMinimumHeight(30) 144 | 145 | # Initialise font 146 | families = ["Consolas", "Monospace", "Courier"] 147 | 148 | for family in families: 149 | family = family.strip() 150 | if family in QtGui.QFontDatabase().families(): 151 | self.setFont(QtGui.QFont(family)) 152 | 153 | def paintEvent(self, event): 154 | """Draw the PE map label 155 | 156 | Args: 157 | event (QPaintEvent): Paint event 158 | 159 | """ 160 | super(PEMapLabel, self).paintEvent(event) 161 | 162 | # Get size of the label rect 163 | r = event.region().boundingRect() 164 | 165 | w = self.width() 166 | h = self.height() 167 | y = r.y() 168 | x = r.x() 169 | 170 | # Get colour and region 171 | colour = self.colour 172 | region = self.region 173 | 174 | # Same colour but with alpha 175 | colour_alpha = QtGui.QColor(colour.red(), colour.green(), colour.blue(), alpha=100 if region.hover is True else 175) 176 | 177 | # Expand the font on mouse over 178 | font = self.font() 179 | font.setWeight(QtGui.QFont.Medium) 180 | 181 | if region.hover: 182 | font.setStretch(QtGui.QFont.Expanded) 183 | else: 184 | font.setStretch(QtGui.QFont.SemiExpanded) 185 | 186 | # Draw the main rect 187 | painter = QtGui.QPainter(self) 188 | 189 | region.rect = QtCore.QRect(x, y, w, h) 190 | painter.setPen(QtCore.Qt.NoPen) 191 | painter.setBrush(QtGui.QBrush(QtGui.QColor(colour), QtCore.Qt.SolidPattern)) 192 | painter.drawRect(region.rect) 193 | 194 | # Determine width per byte of file size 195 | delta = float(w) / float(max(self.file_size, region.start + region.size)) 196 | 197 | # Draw the ratio portion over a white background 198 | painter.setBrush(QtGui.QBrush(QtGui.QColor("white"), QtCore.Qt.SolidPattern)) 199 | painter.drawRect(x + int(region.start * delta), y, max(int(region.size * delta), 2), h) 200 | 201 | painter.setBrush(QtGui.QBrush(QtGui.QColor(colour_alpha), QtCore.Qt.SolidPattern)) 202 | painter.drawRect(x + int(region.start * delta), y, max(int(region.size * delta), 2), h) 203 | 204 | # Draw drop shadow text 205 | painter.setFont(font) 206 | painter.setPen(QtGui.QPen(QtGui.QColor(63, 63, 63, 100))) 207 | painter.drawText(QtCore.QRect(x + 1, y + 1, w, h), QtCore.Qt.AlignCenter, str(region.name)) 208 | 209 | # Write the region name 210 | painter.setPen(QtGui.QPen(QtGui.QColor(244, 244, 250))) 211 | painter.drawText(QtCore.QRect(x, y, w, h), QtCore.Qt.AlignCenter, str(region.name)) 212 | 213 | painter.end() 214 | 215 | def mouseMoveEvent(self, event): 216 | """Set region hover state and redraw the PE map label 217 | 218 | Args: 219 | event (QMouseEvent): Mouse move event 220 | 221 | """ 222 | self.region.hover = True 223 | self.update() 224 | 225 | def leaveEvent(self, event): 226 | """Clear region hover state and redraw the PE map label 227 | 228 | Args: 229 | event (QMouseEvent): Mouse move event 230 | 231 | """ 232 | self.region.hover = False 233 | self.update() 234 | 235 | def mousePressEvent(self, event): 236 | """Locate item related to PE map label in tree view 237 | 238 | Args: 239 | event (QMouseEvent): Mouse press event 240 | 241 | """ 242 | if event.button() == QtCore.Qt.RightButton: 243 | return 244 | 245 | item = self.region.item 246 | form = item.tree.form 247 | index = form.model.indexFromItem(item) 248 | 249 | # Find/expand/scroll to the item in the PE tree 250 | form.treeview.setCurrentIndex(index) 251 | form.expanding = True 252 | form.expand_items(index) 253 | form.expanding = False 254 | form.treeview.resizeColumnToContents(0) 255 | form.treeview.resizeColumnToContents(1) 256 | form.treeview.scrollTo(index, QtWidgets.QAbstractItemView.PositionAtTop) 257 | 258 | def mouseDoubleClickEvent(self, event): 259 | """Locate item related to PE map label in external view 260 | 261 | Args: 262 | event (QMouseEvent): Mouse click event 263 | 264 | """ 265 | self.region.hover = False 266 | 267 | self.region.item.tree.form.runtime.jumpto(self.region.item, self.region.rva) 268 | 269 | def context_menu(self, point): 270 | """Show PE map label right-click context menu 271 | 272 | Args: 273 | point (QPoint): Right-click location 274 | 275 | """ 276 | item = self.region.item 277 | form = item.tree.form 278 | point = self.mapToGlobal(point) 279 | index = form.model.indexFromItem(item) 280 | 281 | item.context_menu(form.context_menu_actions.new_menu(form, point, item, index)) 282 | 283 | class PEMapScrollArea(QtWidgets.QScrollArea): 284 | """PE map scroll area widget""" 285 | def eventFilter(self, obj, event): 286 | if event.type() is QtCore.QEvent.MouseMove: 287 | # Force the map view to update when the mouse moves over the scrollbar 288 | self.widget().update() 289 | 290 | return super(PEMapScrollArea, self).eventFilter(obj, event) 291 | 292 | # 293 | # The following code is gratefully borrowed from: 294 | # https://github.com/bsouthga/blog/blob/master/public/posts/color-gradients-with-python.md 295 | # licensed as follows: 296 | # Copyright 2017 Ben Southgate 297 | # Permission is hereby granted, free of charge, to any person obtaining a copy 298 | # of this software and associated documentation files (the "Software"), to deal 299 | # in the Software without restriction, including without limitation the rights 300 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 301 | # copies of the Software, and to permit persons to whom the Software is 302 | # furnished to do so, subject to the following conditions: 303 | # 304 | # The above copyright notice and this permission notice shall be included in 305 | # all copies or substantial portions of the Software. 306 | # 307 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 308 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 309 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 310 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 311 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 312 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 313 | # THE SOFTWARE 314 | # 315 | 316 | def hex_to_RGB(hex): 317 | ''' "#FFFFFF" -> [255,255,255] ''' 318 | # Pass 16 to the integer function for change of base 319 | return [int(hex[i:i+2], 16) for i in range(1,6,2)] 320 | 321 | 322 | def RGB_to_hex(RGB): 323 | ''' [255,255,255] -> "#FFFFFF" ''' 324 | # Components need to be integers for hex to make sense 325 | RGB = [int(x) for x in RGB] 326 | return "#"+"".join(["0{0:x}".format(v) if v < 16 else "{0:x}".format(v) for v in RGB]) 327 | 328 | def color_dict(gradient): 329 | ''' Takes in a list of RGB sub-lists and returns dictionary of 330 | colors in RGB and hex form for use in a graphing function 331 | defined later on ''' 332 | return {"hex":[RGB_to_hex(RGB) for RGB in gradient], 333 | "r":[RGB[0] for RGB in gradient], 334 | "g":[RGB[1] for RGB in gradient], 335 | "b":[RGB[2] for RGB in gradient]} 336 | 337 | def linear_gradient(start_hex, finish_hex="#FFFFFF", n=10): 338 | ''' returns a gradient list of (n) colors between 339 | two hex colors. start_hex and finish_hex 340 | should be the full six-digit color string, 341 | inlcuding the number sign ("#FFFFFF") ''' 342 | # Starting and ending colors in RGB form 343 | s = hex_to_RGB(start_hex) 344 | f = hex_to_RGB(finish_hex) 345 | # Initilize a list of the output colors with the starting color 346 | RGB_list = [s] 347 | # Calcuate a color at each evenly spaced value of t from 1 to n 348 | for t in range(1, n): 349 | # Interpolate RGB vector for color at the current value of t 350 | curr_vector = [ 351 | int(s[j] + (float(t)/(n-1))*(f[j]-s[j])) 352 | for j in range(3) 353 | ] 354 | # Add it to our list of output colors 355 | RGB_list.append(curr_vector) 356 | 357 | return color_dict(RGB_list) 358 | 359 | def polylinear_gradient(colors, n): 360 | ''' returns a list of colors forming linear gradients between 361 | all sequential pairs of colors. "n" specifies the total 362 | number of desired output colors ''' 363 | # The number of colors per individual linear gradient 364 | n_out = int(float(n) / (len(colors) - 1)) 365 | # returns dictionary defined by color_dict() 366 | gradient_dict = linear_gradient(colors[0], colors[1], n_out) 367 | 368 | if len(colors) > 1: 369 | for col in range(1, len(colors) - 1): 370 | next = linear_gradient(colors[col], colors[col+1], n_out) 371 | for k in ("hex", "r", "g", "b"): 372 | # Exclude first point to avoid duplicates 373 | gradient_dict[k] += next[k][1:] 374 | 375 | return gradient_dict 376 | -------------------------------------------------------------------------------- /pe_tree/minidump.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2021 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree Minidump application""" 18 | 19 | # Standard imports 20 | import os 21 | import sys 22 | from argparse import ArgumentParser, FileType 23 | 24 | # pefile 25 | import pefile 26 | 27 | # Qt imports 28 | from PyQt5 import QtCore, QtWidgets, QtGui 29 | 30 | # Minidump imports 31 | try: 32 | from minidump.minidumpfile import MinidumpFile 33 | from minidump.streams.SystemInfoStream import PROCESSOR_ARCHITECTURE 34 | from minidump.streams.MemoryInfoListStream import AllocationProtect 35 | except ModuleNotFoundError: 36 | print("pip install minidump!") 37 | exit() 38 | 39 | # PE Tree imports 40 | import pe_tree.window 41 | import pe_tree.runtime 42 | import pe_tree.form 43 | import pe_tree.info 44 | 45 | class MinidumpRuntime(pe_tree.runtime.Runtime): 46 | """Minidump runtime callbacks""" 47 | def __init__(self, widget, args): 48 | # Load configuration 49 | self.config_file = os.path.join(self.get_temp_dir(), "pe_tree_minidump.ini") 50 | super(MinidumpRuntime, self).__init__(widget, args) 51 | 52 | # Exported symbols table 53 | self.exports = {} 54 | 55 | # Initialise minidump 56 | try: 57 | self.mf = MinidumpFile.parse(self.args.filename.name) 58 | except: 59 | print("Not a valid Minidump file?") 60 | exit() 61 | 62 | self.reader = self.mf.get_reader() 63 | 64 | @QtCore.pyqtSlot(object, object) 65 | def get_bytes(self, start, size): 66 | """Read bytes from memory""" 67 | self.ret = b"" 68 | 69 | offset = 0 70 | 71 | # Iterate over segments 72 | while size > 0: 73 | # Create buffered segment reader 74 | reader = self.reader.get_buffered_reader() 75 | 76 | # Move to segment 77 | reader.move(start + offset) 78 | 79 | # Calculate remaining data in the segment 80 | remaining = reader.current_segment.remaining_len(start + offset) 81 | 82 | if not remaining: 83 | break 84 | 85 | if remaining >= size: 86 | # No more data required, read and break 87 | self.ret += reader.peek(size) 88 | break 89 | 90 | # More data required, read and continue 91 | self.ret += reader.peek(remaining) 92 | 93 | offset += remaining 94 | size -= remaining 95 | 96 | self.ret = b"" if self.ret is None else self.ret 97 | return self.ret 98 | 99 | @QtCore.pyqtSlot(object) 100 | def is_writable(self, offset): 101 | """Determine if memory address is writable""" 102 | self.ret = False 103 | 104 | for info in self.mf.memory_info.infos: 105 | if offset >= info.BaseAddress and offset <= info.BaseAddress + info.RegionSize: 106 | if info.Protect is AllocationProtect.PAGE_EXECUTE_READWRITE: 107 | self.ret = True 108 | break 109 | 110 | return self.ret 111 | 112 | def _build_export_symbol_table(self): 113 | """Iterate over all load order modules in the processes address space and construct an exported function symbol table""" 114 | if len(self.exports.keys()) > 0: 115 | return 116 | 117 | # Iterate over all modules 118 | for module in self.mf.modules.modules: 119 | try: 120 | # Ensure the module has an MZ header 121 | if self.get_word(module.baseaddress) != 0x5a4d: 122 | continue 123 | 124 | # Dump and re-parse the PE 125 | pe_file = pefile.PE(data=self.read_pe(module.baseaddress)) 126 | 127 | # Check for exports data directory 128 | if hasattr(pe_file, "DIRECTORY_ENTRY_EXPORT"): 129 | # Iterate over all exports 130 | for export in pe_file.DIRECTORY_ENTRY_EXPORT.symbols: 131 | # Determine the export name and address 132 | export_name = export.name.decode("utf-8") if export.name is not None else "" 133 | export_address = export.address if export.address < pe_file.OPTIONAL_HEADER.ImageBase else export.address - pe_file.OPTIONAL_HEADER.ImageBase 134 | 135 | # Log the module name and export name for the given address 136 | self.exports[module.baseaddress + export_address] = [os.path.basename(module.name), export_name] 137 | except: 138 | continue 139 | 140 | @QtCore.pyqtSlot(object) 141 | def resolve_address(self, offset): 142 | """Get module/symbol name for address""" 143 | module = "" 144 | api = "" 145 | 146 | # Construct an exports symbol table if required 147 | self._build_export_symbol_table() 148 | 149 | # Search exports for offset 150 | if offset in self.exports: 151 | module, api = self.exports[offset] 152 | 153 | self.ret = (module, api) 154 | 155 | return self.ret 156 | 157 | class SelectProcess(pe_tree.dialogs.ModulePicker): 158 | """Process/module picker dialog""" 159 | def __init__(self, window, parent=None): 160 | super(SelectProcess, self).__init__(window, "Select module...", ["Module", "Image base"], parent=parent) 161 | 162 | def populate_processes(self): 163 | """Populate the tree view with processes""" 164 | root = self.model.invisibleRootItem() 165 | 166 | # Determine architecture specifics 167 | if self.window.runtime.reader.sysinfo.ProcessorArchitecture == PROCESSOR_ARCHITECTURE.AMD64: 168 | ptr_mask = 0xffffffffffffffff 169 | ptr_width = 16 170 | else: 171 | ptr_mask = 0xffffffff 172 | ptr_width = 8 173 | 174 | # Iterate over all modules 175 | for module in self.window.runtime.mf.modules.modules: 176 | # Determine image base (ignore if null) 177 | image_base = int(module.baseaddress) 178 | 179 | if image_base == 0: 180 | continue 181 | 182 | # Create internal filename 183 | filename = "{} - {:#08x}".format(str(module.name), image_base) 184 | 185 | # Create process tree view item 186 | process_item = QtGui.QStandardItem("{}".format(str(module.name))) 187 | process_item.setData({"filename": filename, "module": module, "image_base": image_base}, QtCore.Qt.UserRole) 188 | process_item.setCheckable(True) 189 | if root == self.model.invisibleRootItem(): 190 | process_item.setUserTristate(True) 191 | 192 | root.appendRow([process_item, QtGui.QStandardItem("0x{:0{w}x}".format(image_base & ptr_mask, w=ptr_width))]) 193 | 194 | if root == self.model.invisibleRootItem(): 195 | root = process_item 196 | 197 | self.treeview.resizeColumnToContents(0) 198 | 199 | def invoke(self): 200 | """Display the select process/modules dialog""" 201 | super(SelectProcess, self).invoke() 202 | 203 | # Has the user accepted, i.e. pressed "Dump"? 204 | if self.exec_() != 0: 205 | # Dump selected processes/modules 206 | for item in self.items: 207 | args = item.item.data(QtCore.Qt.UserRole) 208 | self.window.pe_tree_form.map_pe(filename=args["filename"], image_base=args["image_base"], opaque={"module": args["module"]}) 209 | 210 | self.close() 211 | 212 | def main(): 213 | """PE Tree Minidump script entry-point""" 214 | # Check command line arguments 215 | parser = ArgumentParser(description="PE-Tree (Minidump)") 216 | parser.add_argument("filename", help="Path to .dmp file", type=FileType("rb")) 217 | args = parser.parse_args() 218 | 219 | # Create PE Tree Qt application 220 | application = QtWidgets.QApplication(sys.argv) 221 | window = pe_tree.window.PETreeWindow(application, MinidumpRuntime, args, open_file=False) 222 | 223 | # Extend menu to include open process 224 | select_process = SelectProcess(window) 225 | 226 | open_process_action = QtWidgets.QAction("Process", window) 227 | open_process_action.setShortcut("Ctrl+Shift+P") 228 | open_process_action.setStatusTip("Open process") 229 | open_process_action.triggered.connect(select_process.invoke) 230 | window.open_menu.addAction(open_process_action) 231 | 232 | # Invoke the process/module picker dialog 233 | select_process.invoke() 234 | 235 | sys.exit(application.exec_()) 236 | 237 | if __name__ == "__main__": 238 | main() 239 | -------------------------------------------------------------------------------- /pe_tree/resources/cyberchef.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/pe_tree/resources/cyberchef.ico -------------------------------------------------------------------------------- /pe_tree/resources/spear.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/pe_tree/resources/spear.png -------------------------------------------------------------------------------- /pe_tree/resources/virustotal.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/blackberry/pe_tree/2be607fc55702293cd02cbc6cda5283452464aff/pe_tree/resources/virustotal.png -------------------------------------------------------------------------------- /pe_tree/scandir.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2021 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree directory scanning""" 18 | 19 | # Standard imports 20 | import os 21 | import zipfile 22 | 23 | # Scan dir 24 | import scandir 25 | 26 | # Qt imports 27 | from PyQt5 import QtCore 28 | 29 | class ScanDir(QtCore.QRunnable): 30 | """Scan directory thread. Responsible for traversing directories and archives. 31 | 32 | Args: 33 | pe_tree_form (pe_tree.form): PE Tree form 34 | filename (str): Path to file/folder/zip archive to scan for PE files 35 | 36 | """ 37 | def __init__(self, pe_tree_form, filename): 38 | super(ScanDir, self).__init__() 39 | 40 | self.filename = filename 41 | self.pe_tree_form = pe_tree_form 42 | 43 | def unpack_and_map(self, filename): 44 | """Attempt to extract a zip archive and map PE files""" 45 | try: 46 | # Attempt to open as zip archive 47 | with zipfile.ZipFile(filename) as zip_archive: 48 | # Iterate over all filenames 49 | for zip_filename in zip_archive.namelist(): 50 | # Iterate over passwords from the config file 51 | for password in self.pe_tree_form.runtime.get_config_option("config", "passwords", "").split(","): 52 | if password != "": 53 | # Set the password if specified 54 | zip_archive.setpassword(password.encode("utf-8")) 55 | 56 | # Attempt to extract the file 57 | try: 58 | with zip_archive.open(zip_filename, "r") as zipped_file: 59 | # Read the zipped file 60 | data = zipped_file.read() 61 | 62 | # Ensure the file has an MZ header 63 | if data[:2] == b"MZ": 64 | # Map the PE data 65 | self.pe_tree_form.map_pe(filename=os.path.join(filename, zip_filename), data=data, disable_dump=True) 66 | 67 | break 68 | except RuntimeError: 69 | # Possibly a bad password? 70 | pass 71 | 72 | # No point trying to map as a PE as well 73 | return 74 | except zipfile.BadZipFile: 75 | # Attempt to map as PE file 76 | self.pe_tree_form.map_pe(filename=filename) 77 | except OSError: 78 | pass 79 | 80 | def run(self): 81 | """Iterate over all files/folders and map PE files""" 82 | if os.path.isdir(self.filename): 83 | for root, dirs, files in scandir.walk(self.filename): 84 | for file in files: 85 | # Have we been signaled to stop? 86 | if self.pe_tree_form.stop_event.is_set(): 87 | return 88 | 89 | # Map the PE file 90 | self.pe_tree_form.map_pe(filename=os.path.join(root, file)) 91 | else: 92 | # Unpack and map PE files 93 | self.unpack_and_map(filename=self.filename) -------------------------------------------------------------------------------- /pe_tree/utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree utility functions""" 18 | 19 | # Standard imports 20 | import os 21 | import sys 22 | import string 23 | import struct 24 | from io import BytesIO 25 | import collections 26 | 27 | # Qt imports 28 | from PyQt5 import QtGui, Qt 29 | 30 | def hexdump(buffer, base=0, width=16): 31 | """Hexdump a buffer 32 | 33 | Args: 34 | buffer (bytes): Data to hexdump 35 | base (int): Base address of data 36 | width (int): Number of bytes to show per line 37 | 38 | Returns: 39 | str: Hex dump of data 40 | 41 | """ 42 | dump = "" 43 | offset = 0 44 | printable = "".join([c if ord(c) >= 0x20 and ord(c) < 0x7f else "" for c in string.printable]) 45 | 46 | if sys.version_info > (3,): 47 | _hex = lambda line: " ".join(["{:02x}".format(c) for c in line]) 48 | _ascii = lambda line: "".join([chr(c) if chr(c) in printable else "." for c in line]) 49 | else: 50 | _hex = lambda line: " ".join(["{:02x}".format(ord(c)) for c in line]) 51 | _ascii = lambda line: "".join([c if c in printable else "." for c in line]) 52 | 53 | while buffer: 54 | line = buffer[:width] 55 | 56 | dump += "0x{:08x} {:{}} {}{}".format(base + offset, _hex(line), width * 3, _ascii(line), os.linesep) 57 | 58 | offset += width 59 | buffer = buffer[width:] 60 | 61 | return dump 62 | 63 | def human_readable_filesize(size): 64 | """Convert file size in bytes to human readable format 65 | 66 | Args: 67 | size (int): Size in bytes 68 | 69 | Returns: 70 | str: Human readable file-size, i.e. 567.4 KB (580984 bytes) 71 | 72 | """ 73 | if size < 1024: 74 | return "{} bytes".format(size) 75 | remain = float(size) 76 | for unit in ["B", "KB", "MB", "GB", "TB"]: 77 | if remain < 1024.0: 78 | return "{:.1f} {} ({:d} bytes)".format(remain, unit, size) 79 | remain /= 1024.0 80 | 81 | def qicon_from_ico_data(data): 82 | """Convert ICO/PNG data to QIcon 83 | 84 | Args: 85 | data (bytes): Data starting with either an ICONDIR or BITMAPINFOHEADER 86 | 87 | Returns: 88 | QIcon: Icon loaded from data if successful, otherwise an empty icon 89 | 90 | """ 91 | # Ensure we have some data to work with 92 | if not data: 93 | return QtGui.QIcon() 94 | 95 | # Is this a PNG icon? 96 | if data[1:4].startswith(b"PNG"): 97 | # Load PNG image data 98 | return QtGui.QIcon(QtGui.QPixmap.fromImage(QtGui.QImage.fromData(data))) 99 | 100 | # Starts with BITMAPINFOHEADER? 101 | if data.startswith(b"\x28"): 102 | # Add dummy 16x16 ICONDIR header 103 | data = b"\x00\x00\x01\x00\x01\x00\x10\x10\x00\x00\x00\x00\x00\x00" + struct.pack(" 256: 145 | return QtGui.QIcon() 146 | 147 | # Contains a colour palette? 148 | if bitmap_info.colours_used != 0 or bitmap_info.bpp <= 8: 149 | # Determine the total number of colours in the palette 150 | colours = bitmap_info.colours_used if bitmap_info.colours_used > 0 else pow(2, bitmap_info.bpp) 151 | 152 | # Read RGBA palette colours 153 | palette_data = bmp.read(colours * 4) 154 | palette = [] 155 | 156 | # Parse palette colours into RGBA 157 | for i in range(colours): 158 | b, g, r, a = struct.unpack("!BBBB", palette_data[i * 4 : i * 4 + 4]) 159 | palette.append((r, g, b, 0xff)) 160 | 161 | # Determine row stride 162 | stride = ((((bitmap_info.width * bitmap_info.bpp) + 31) & ~31) >> 3) 163 | 164 | # If height is greater than width then we have alpha channel data 165 | height = bitmap_info.width 166 | 167 | # Read pixels 168 | pixel_data = bmp.read(stride * height) 169 | 170 | if sys.version_info > (3,): 171 | _ord = int 172 | else: 173 | _ord = ord 174 | 175 | # Determine pixel format 176 | if bitmap_info.bpp == 32: 177 | # 32-bit RGBA 178 | format = QtGui.QImage.Format_RGBA8888 179 | elif bitmap_info.bpp == 24: 180 | # 24-bit RGB 181 | format = QtGui.QImage.Format_RGB888 182 | elif bitmap_info.bpp == 16: 183 | # 16-bit RGB 184 | format = QtGui.QImage.Format_RGB16 185 | elif bitmap_info.bpp == 8: 186 | # 8-bit indexed, convert to 32-bit RGBA 187 | format = QtGui.QImage.Format_RGBA8888 188 | pixels = b"" 189 | for index in pixel_data: 190 | r, g, b, a = palette[_ord(index)] 191 | pixels += struct.pack("!BBBB", r, g, b, a) 192 | 193 | pixel_data = pixels 194 | elif bitmap_info.bpp == 4: 195 | # 4-bit indexed, convert to 32-bit RGBA 196 | format = QtGui.QImage.Format_RGBA8888 197 | pixels = b"" 198 | for index in pixel_data: 199 | index = _ord(index) 200 | for i in range(2): 201 | r, g, b, a = palette[(index & 0xf0) >> 4] 202 | pixels += struct.pack("!BBBB", r, g, b, a) 203 | index <<= 4 204 | 205 | pixel_data = pixels 206 | elif bitmap_info.bpp == 2: 207 | # 2-bit indexed, convert to 32-bit RGBA 208 | format = QtGui.QImage.Format_RGBA8888 209 | pixels = b"" 210 | for index in pixel_data: 211 | index = _ord(index) 212 | for i in range(4): 213 | r, g, b, a = palette[(index & 0xC0) >> 6] 214 | pixels += struct.pack("!BBBB", r, g, b, a) 215 | index <<= 2 216 | 217 | pixel_data = pixels 218 | elif bitmap_info.bpp == 1: 219 | # 1-bit indexed, convert to 32-bit RGBA 220 | format = QtGui.QImage.Format_RGBA8888 221 | pixels = b"" 222 | for index in pixel_data: 223 | index = _ord(index) 224 | for i in range(8): 225 | r, g, b, a = palette[(index & 0x80) >> 7] 226 | pixels += struct.pack("!BBBB", r, g, b, a) 227 | index <<= 1 228 | 229 | pixel_data = pixels 230 | 231 | # Create image from pixel data and mirror flip (as the bitmap is stored backwards) 232 | image = QtGui.QImage(pixel_data, bitmap_info.width, bitmap_info.width, format).mirrored(False, True) 233 | pixmap = QtGui.QPixmap.fromImage(image) 234 | 235 | # No alpha? 236 | if bitmap_info.bpp < 32: 237 | # Use pixel 0, 0 as the transparent background colour 238 | pixmap.setMask(pixmap.createMaskFromColor(image.pixelColor(0, 0), Qt.Qt.MaskInColor)) 239 | 240 | # Create icon from pixmap 241 | return QtGui.QIcon(pixmap) 242 | except: 243 | return QtGui.QIcon() 244 | -------------------------------------------------------------------------------- /pe_tree/volatility.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2021 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree Volatility3 application""" 18 | 19 | # Standard imports 20 | import os 21 | import sys 22 | from argparse import ArgumentParser, FileType 23 | 24 | # pefile 25 | import pefile 26 | 27 | # Qt imports 28 | from PyQt5 import QtCore, QtWidgets, QtGui 29 | 30 | # Volatility3 imports 31 | try: 32 | from volatility3 import framework, plugins 33 | from volatility3.framework import contexts, interfaces, symbols 34 | from volatility3.plugins.windows import pslist, modules, vadinfo 35 | except ModuleNotFoundError: 36 | print("volatility3 is required") 37 | exit() 38 | 39 | # PE Tree imports 40 | import pe_tree.window 41 | import pe_tree.dialogs 42 | import pe_tree.runtime 43 | import pe_tree.form 44 | import pe_tree.qstandarditems 45 | import pe_tree.info 46 | 47 | class VolatilityRuntime(pe_tree.runtime.Runtime): 48 | """Volatility runtime callbacks""" 49 | def __init__(self, widget, args): 50 | # Load configuration 51 | self.config_file = os.path.join(self.get_temp_dir(), "pe_tree_volatility.ini") 52 | super(VolatilityRuntime, self).__init__(widget, args) 53 | 54 | # Initialise Volatility3 55 | framework.require_interface_version(1, 0, 0) 56 | self.context = contexts.Context() 57 | 58 | # Input file in URI format 59 | single_location = "file:///{}".format(self.args.filename.name) 60 | self.context.config["automagic.LayerStacker.single_location"] = single_location 61 | self.context.config["single_location"] = single_location 62 | 63 | # Find supported plugins 64 | plugins.__path__ = framework.constants.PLUGINS_PATH 65 | framework.import_files(plugins, True) 66 | self.plugin_list = framework.list_plugins() 67 | 68 | # Initialise required plugins configuration 69 | self._automagic("windows.pslist.PsList") 70 | self._automagic("windows.vadinfo.VadInfo") 71 | self._automagic("windows.modules.Modules") 72 | 73 | @QtCore.pyqtSlot(object, object) 74 | def get_bytes(self, start, size): 75 | """Read bytes from memory""" 76 | # Create the process layer if required 77 | if "proc" in self.opaque and "proc_layer_name" not in self.opaque: 78 | if self.opaque["proc"]: 79 | self.opaque["proc_layer_name"] = self.opaque["proc"].add_process_layer() 80 | 81 | # Find the process' layer 82 | proc_layer = self.context.layers[self.opaque["proc_layer_name"]] 83 | 84 | try: 85 | # Read from the process layer 86 | self.ret = proc_layer.read(start, size, pad=True) 87 | except (framework.exceptions.PagedInvalidAddressException, framework.exceptions.InvalidAddressException): 88 | # Might be the page isn't available, fake it 89 | self.ret = b"\0" * size 90 | 91 | self.ret = b"" if self.ret is None else self.ret 92 | return self.ret 93 | 94 | @QtCore.pyqtSlot(object) 95 | def is_writable(self, offset): 96 | """Determine if memory address is writable""" 97 | self.ret = False 98 | 99 | # Require a process 100 | if "proc" in self.opaque: 101 | # Iterate over VADs 102 | for vad in self.get_vads(self.opaque["proc"]): 103 | # Does the VAD contain the offset? 104 | if offset >= vad.get_start() and offset < vad.get_end(): 105 | # Check for write protection flag 106 | if "WRITE" in vad.get_protection(vadinfo.VadInfo.protect_values(self.context, self.context.config["plugins.VadInfo.VadInfo.primary"], self.context.config["plugins.VadInfo.VadInfo.nt_symbols"]), vadinfo.winnt_protections): 107 | self.ret = True 108 | break 109 | 110 | return self.ret 111 | 112 | def _build_export_symbol_table(self): 113 | """Iterate over all load order modules in a processes address space and construct an exported function symbol table""" 114 | # Require an export symbols table 115 | if "exports" not in self.opaque: 116 | return 117 | 118 | # Require a process 119 | if "proc" not in self.opaque or self.opaque["proc"] is None: 120 | return 121 | 122 | # Skip if already built 123 | if len(self.opaque["exports"]) > 0: 124 | return 125 | 126 | if self.opaque["proc"].UniqueProcessId == 4: 127 | # Find all kernel modules 128 | export_modules = self.get_modules() 129 | 130 | # Find the kernel layer name 131 | self.opaque["proc_layer_name"] = self.context.config["plugins.Modules.Modules.primary"] 132 | else: 133 | # Find all process load order modules 134 | export_modules = self.opaque["proc"].load_order_modules() 135 | 136 | # Iterate over all modules 137 | for entry in export_modules: 138 | try: 139 | # Ensure the module has an MZ header 140 | if self.get_word(entry.DllBase) != 0x5a4d: 141 | continue 142 | 143 | # Dump and re-parse the PE 144 | pe_file = pefile.PE(data=self.read_pe(entry.DllBase)) 145 | 146 | # Check for exports data directory 147 | if hasattr(pe_file, "DIRECTORY_ENTRY_EXPORT"): 148 | # Iterate over all exports 149 | for export in pe_file.DIRECTORY_ENTRY_EXPORT.symbols: 150 | # Determine the export name and address 151 | export_name = export.name.decode("utf-8") if export.name is not None else "" 152 | export_address = export.address if export.address < pe_file.OPTIONAL_HEADER.ImageBase else export.address - pe_file.OPTIONAL_HEADER.ImageBase 153 | 154 | # Log the module name and export name for the given address 155 | self.opaque["exports"][entry.DllBase + export_address] = [entry.BaseDllName.get_string(), export_name] 156 | except: 157 | continue 158 | 159 | @QtCore.pyqtSlot(object) 160 | def resolve_address(self, offset): 161 | """Get module/symbol name for address""" 162 | module = "" 163 | api = "" 164 | 165 | # Construct an exports symbol table if required 166 | self._build_export_symbol_table() 167 | 168 | # Search exports for offset 169 | if "exports" in self.opaque: 170 | exports = self.opaque["exports"] 171 | if offset in exports: 172 | module, api = exports[offset] 173 | 174 | self.ret = (module, api) 175 | 176 | return self.ret 177 | 178 | def _automagic(self, plugin_name): 179 | """Use automagic to initialise volatility plugin configuration""" 180 | plugin = self.plugin_list[plugin_name] 181 | available_automagics = framework.automagic.available(self.context) 182 | framework.automagic.choose_automagic(available_automagics, plugin) 183 | framework.automagic.run(available_automagics, self.context, plugin, interfaces.configuration.path_join("plugins", plugin.__name__)) 184 | 185 | def get_procs(self): 186 | """Run windows.pslist.PsList plugin""" 187 | return pslist.PsList.list_processes(context=self.context, layer_name=self.context.config["plugins.PsList.PsList.primary"], symbol_table=self.context.config["plugins.PsList.PsList.nt_symbols"]) 188 | 189 | def get_modules(self): 190 | """Run windows.modules.Modules plugin""" 191 | return modules.Modules.list_modules(context=self.context, layer_name=self.context.config["plugins.Modules.Modules.primary"], symbol_table=self.context.config["plugins.Modules.Modules.nt_symbols"]) 192 | 193 | def get_vads(self, proc): 194 | """Run windows.vadinfo.VadInfo plugin""" 195 | return vadinfo.VadInfo.list_vads(proc) 196 | 197 | class SelectProcess(pe_tree.dialogs.ModulePicker): 198 | """Process/module picker dialog""" 199 | def __init__(self, window, parent=None): 200 | super(SelectProcess, self).__init__(window, "Select process...", ["Name", "PID", "PPID", "Image base", "Threads", "Handles", "Create time", "IsWow64"], parent=parent) 201 | 202 | def make_process_item(self, proc, entry, exports, expandable=False): 203 | image_base = int(entry.DllBase) 204 | # Create internal filename 205 | filename = "{} ({}) - {:#08x}".format(str(entry.BaseDllName.get_string()), proc.UniqueProcessId, image_base) 206 | 207 | # Create process tree view item 208 | process_item = QtGui.QStandardItem("{}".format(str(entry.BaseDllName.get_string()))) 209 | process_item.setData({"filename": filename, "proc": proc, "entry": entry, "exports": exports, "image_base": image_base}, QtCore.Qt.UserRole) 210 | process_item.setCheckable(True) 211 | if expandable: 212 | process_item.setUserTristate(True) 213 | 214 | return process_item 215 | 216 | def populate_processes(self): 217 | """Populate the tree view with processes""" 218 | root = self.model.invisibleRootItem() 219 | 220 | for proc in self.window.runtime.get_procs(): 221 | root = self.model.invisibleRootItem() 222 | 223 | # This will hold the export symbol table for an entire process 224 | exports = {} 225 | 226 | # System process? 227 | if proc.UniqueProcessId == 4: 228 | image_base = self.window.runtime.context.config["plugins.PsList.PsList.primary.kernel_virtual_offset"] 229 | 230 | # Create internal filename 231 | filename = "{} ({}) - {:#08x}".format("System", proc.UniqueProcessId, image_base) 232 | 233 | # Create process tree view item 234 | process_item = QtGui.QStandardItem("System") 235 | process_item.setData({"filename": filename, "proc": proc, "entry": None, "exports": exports, "image_base": image_base}, QtCore.Qt.UserRole) 236 | process_item.setUserTristate(True) 237 | process_item.setCheckable(True) 238 | 239 | # Determine kernel architecture specifics 240 | if symbols.symbol_table_is_64bit(self.window.runtime.context, self.window.runtime.context.config["plugins.PsList.PsList.nt_symbols"]): 241 | ptr_mask = 0xffffffffffffffff 242 | ptr_width = 16 243 | else: 244 | ptr_mask = 0xffffffff 245 | ptr_width = 8 246 | 247 | root.appendRow([process_item, 248 | QtGui.QStandardItem("{}".format(proc.UniqueProcessId)), 249 | QtGui.QStandardItem("{}".format(proc.InheritedFromUniqueProcessId)), 250 | QtGui.QStandardItem("0x{:0{w}x}".format(image_base & ptr_mask, w=ptr_width)), 251 | QtGui.QStandardItem(""), 252 | QtGui.QStandardItem(""), 253 | QtGui.QStandardItem(""), 254 | QtGui.QStandardItem("")]) 255 | 256 | root = process_item 257 | 258 | # Iterate over kernel modules 259 | for module in self.window.runtime.get_modules(): 260 | try: 261 | process_item = self.make_process_item(proc, module, exports) 262 | 263 | root.appendRow([process_item, 264 | QtGui.QStandardItem("{}".format(proc.UniqueProcessId)), 265 | QtGui.QStandardItem(""), 266 | QtGui.QStandardItem("0x{:0{w}x}".format(int(module.DllBase) & ptr_mask, w=ptr_width)), 267 | QtGui.QStandardItem(""), 268 | QtGui.QStandardItem(""), 269 | QtGui.QStandardItem(""), 270 | QtGui.QStandardItem("")]) 271 | 272 | except framework.exceptions.PagedInvalidAddressException: 273 | pass 274 | 275 | # Iterate over load order modules 276 | for entry in proc.load_order_modules(): 277 | try: 278 | process_item = self.make_process_item(proc, entry, exports, expandable=bool(root == self.model.invisibleRootItem())) 279 | except framework.exceptions.PagedInvalidAddressException: 280 | continue 281 | 282 | # Determine process architecture specifics 283 | if proc.get_is_wow64(): 284 | ptr_mask = 0xffffffffffffffff 285 | ptr_width = 16 286 | else: 287 | ptr_mask = 0xffffffff 288 | ptr_width = 8 289 | 290 | 291 | if root == self.model.invisibleRootItem(): 292 | # Include full details for the main process 293 | root.appendRow([process_item, 294 | QtGui.QStandardItem("{}".format(proc.UniqueProcessId)), 295 | QtGui.QStandardItem("{}".format(proc.InheritedFromUniqueProcessId)), 296 | QtGui.QStandardItem("0x{:0{w}x}".format(int(entry.DllBase) & ptr_mask, w=ptr_width)), 297 | QtGui.QStandardItem("{}".format(proc.ActiveThreads)), 298 | QtGui.QStandardItem("{}".format(proc.get_handle_count())), 299 | QtGui.QStandardItem("{}".format(proc.get_create_time())), 300 | QtGui.QStandardItem("{}".format(proc.get_is_wow64()))]) 301 | 302 | root = process_item 303 | else: 304 | # Add details for the DLL 305 | root.appendRow([process_item, 306 | QtGui.QStandardItem("{}".format(proc.UniqueProcessId)), 307 | QtGui.QStandardItem("{}".format(proc.InheritedFromUniqueProcessId)), 308 | QtGui.QStandardItem("0x{:0{w}x}".format(int(entry.DllBase) & ptr_mask, w=ptr_width)), 309 | QtGui.QStandardItem(""), 310 | QtGui.QStandardItem(""), 311 | QtGui.QStandardItem(""), 312 | QtGui.QStandardItem("")]) 313 | 314 | for i in range(0, 7): 315 | self.treeview.resizeColumnToContents(i) 316 | 317 | def invoke(self): 318 | """Display the select process/modules dialog""" 319 | super(SelectProcess, self).invoke() 320 | 321 | # Has the user accepted, i.e. pressed "Dump"? 322 | if self.exec_() != 0: 323 | # Dump selected processes/modules 324 | for item in self.items: 325 | args = item.item.data(QtCore.Qt.UserRole) 326 | self.window.pe_tree_form.map_pe(filename=args["filename"], image_base=args["image_base"], opaque={"proc": args["proc"], "entry": args["entry"], "exports": args["exports"]}) 327 | 328 | self.close() 329 | 330 | def main(): 331 | """PE Tree Volatility script entry-point""" 332 | # Check command line arguments 333 | parser = ArgumentParser(description="PE-Tree (Volatility)") 334 | parser.add_argument("filename", help="Path to memory dump", type=FileType("rb")) 335 | args = parser.parse_args() 336 | 337 | # Create PE Tree Qt application 338 | application = QtWidgets.QApplication(sys.argv) 339 | 340 | window = pe_tree.window.PETreeWindow(application, VolatilityRuntime, args, open_file=False) 341 | 342 | # Extend menu to include open process 343 | select_process = SelectProcess(window) 344 | 345 | open_process_action = QtWidgets.QAction("Process", window) 346 | open_process_action.setShortcut("Ctrl+Shift+P") 347 | open_process_action.setStatusTip("Open process") 348 | open_process_action.triggered.connect(select_process.invoke) 349 | window.open_menu.addAction(open_process_action) 350 | 351 | # Invoke the process/module picker dialog 352 | select_process.invoke() 353 | 354 | sys.exit(application.exec_()) 355 | 356 | if __name__ == '__main__': 357 | main() 358 | -------------------------------------------------------------------------------- /pe_tree/window.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2021 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree main window handling""" 18 | 19 | import os 20 | 21 | # Qt imports 22 | from PyQt5 import QtCore, QtWidgets 23 | 24 | # QDarkStyle 25 | try: 26 | import qdarkstyle 27 | 28 | HAVE_DARKSTYLE = True 29 | except: 30 | HAVE_DARKSTYLE = False 31 | 32 | # PE Tree imports 33 | import pe_tree.form 34 | import pe_tree.scandir 35 | import pe_tree.info 36 | 37 | class PETreeWindow(QtWidgets.QMainWindow): 38 | """Main window for the PE Tree application""" 39 | 40 | def __init__(self, application, runtime, args, open_file=True): 41 | super(PETreeWindow, self).__init__() 42 | 43 | self.application = application 44 | 45 | # Create container widget, runtime and PE Tree form 46 | widget = QtWidgets.QWidget() 47 | self.runtime = runtime(widget, args) 48 | self.pe_tree_form = pe_tree.form.PETreeForm(widget, application, self.runtime) 49 | self.pe_tree_form.dispatcher = application.eventDispatcher() 50 | application.aboutToQuit.connect(self.pe_tree_form.wait_for_threads) 51 | self.runtime.pe_tree_form = self.pe_tree_form 52 | 53 | if HAVE_DARKSTYLE: 54 | application.setStyleSheet(qdarkstyle.load_stylesheet(qt_api="pyqt5")) 55 | 56 | self.setWindowTitle(pe_tree.info.__title__) 57 | self.resize(1024, 768) 58 | self.setCentralWidget(widget) 59 | self.setAcceptDrops(True) 60 | 61 | # Construct application main menu 62 | self.open_menu = QtWidgets.QMenu("Open", self) 63 | 64 | open_file_action = QtWidgets.QAction("File", self) 65 | open_file_action.setShortcut("Ctrl+O") 66 | open_file_action.setStatusTip("Open PE file") 67 | open_file_action.triggered.connect(self.open_file) 68 | 69 | open_directory_action = QtWidgets.QAction("Folder", self) 70 | open_directory_action.setShortcut("Ctrl+Shift+O") 71 | open_directory_action.setStatusTip("Scan folder for PE files") 72 | open_directory_action.triggered.connect(self.open_folder) 73 | 74 | self.open_menu.addAction(open_file_action) 75 | self.open_menu.addAction(open_directory_action) 76 | 77 | exit_action = QtWidgets.QAction("Exit", self) 78 | exit_action.setShortcut("Ctrl+X") 79 | exit_action.setStatusTip("Exit") 80 | exit_action.triggered.connect(application.quit) 81 | 82 | about_action = QtWidgets.QAction("About", self) 83 | about_action.triggered.connect(self.runtime.about_box) 84 | 85 | menu = self.menuBar() 86 | file_menu = menu.addMenu("&File") 87 | file_menu.addMenu(self.open_menu) 88 | file_menu.addAction(exit_action) 89 | 90 | help_menu = menu.addMenu("&Help") 91 | help_menu.addAction(about_action) 92 | 93 | self.showMaximized() 94 | 95 | # Process command line 96 | if open_file: 97 | if args.filenames: 98 | # Map all files/folders specified on the command line 99 | for filename in args.filenames: 100 | if os.path.exists(filename): 101 | self.pe_tree_form.threadpool.start(pe_tree.scandir.ScanDir(self.pe_tree_form, filename)) 102 | else: 103 | # Ask user to select file/folder 104 | self.open_file() 105 | 106 | def dragEnterEvent(self, e): 107 | """Accept drag events containing URLs""" 108 | if e.mimeData().hasUrls: 109 | e.accept() 110 | else: 111 | e.ignore() 112 | 113 | def dragMoveEvent(self, e): 114 | """Accept drag events containing URLs""" 115 | if e.mimeData().hasUrls: 116 | e.accept() 117 | else: 118 | e.ignore() 119 | 120 | def dropEvent(self, e): 121 | """File drag/dropped onto the main form, attempt to map as PE""" 122 | if e.mimeData().hasUrls: 123 | e.accept() 124 | 125 | for url in e.mimeData().urls(): 126 | self.pe_tree_form.threadpool.start(pe_tree.scandir.ScanDir(self.pe_tree_form, QtCore.QDir.toNativeSeparators(str(url.toLocalFile()))), priority=1) 127 | else: 128 | e.ignore() 129 | 130 | def exec_open_dialog(self, dialog): 131 | """Execute open dialog and map PE files""" 132 | if dialog.exec_() == QtWidgets.QDialog.Accepted: 133 | for filename in dialog.selectedFiles(): 134 | self.pe_tree_form.threadpool.start(pe_tree.scandir.ScanDir(self.pe_tree_form, QtCore.QDir.toNativeSeparators(filename))) 135 | 136 | return True 137 | 138 | return False 139 | 140 | def open_file(self): 141 | """Prompt user to select file to map files from""" 142 | options = QtWidgets.QFileDialog.Options() 143 | options |= QtWidgets.QFileDialog.DontUseNativeDialog 144 | 145 | dialog = QtWidgets.QFileDialog() 146 | dialog.setOptions(options) 147 | dialog.setFileMode(QtWidgets.QFileDialog.AnyFile) 148 | dialog.setAcceptMode(QtWidgets.QFileDialog.AcceptOpen) 149 | 150 | return self.exec_open_dialog(dialog) 151 | 152 | def open_folder(self): 153 | """Prompt user to select folder to map files from""" 154 | options = QtWidgets.QFileDialog.Options() 155 | options |= QtWidgets.QFileDialog.DontUseNativeDialog 156 | 157 | dialog = QtWidgets.QFileDialog() 158 | dialog.setOptions(options) 159 | dialog.setFileMode(QtWidgets.QFileDialog.Directory) 160 | dialog.setAcceptMode(QtWidgets.QFileDialog.AcceptOpen) 161 | 162 | return self.exec_open_dialog(dialog) -------------------------------------------------------------------------------- /pe_tree_ida.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree IDAPython plugin""" 18 | 19 | # Standard imports 20 | import os 21 | import sys 22 | 23 | # Qt imports 24 | from PyQt5 import QtCore, QtWidgets 25 | 26 | # IDAPython imports 27 | 28 | # pylint: disable=import-error 29 | 30 | import idaapi 31 | import ida_idaapi 32 | import idc 33 | import ida_bytes 34 | import ida_name 35 | import ida_kernwin 36 | import idautils 37 | import sip 38 | 39 | # PE Tree imports 40 | if "pe_tree" in sys.modules: 41 | # Reload all PE Tree modules as IDA will keep them loaded 42 | try: 43 | from importlib import reload 44 | except ImportError: 45 | from imp import reload 46 | 47 | try: 48 | reload(pe_tree) 49 | reload(pe_tree.runtime) 50 | reload(pe_tree.form) 51 | reload(pe_tree.info) 52 | reload(pe_tree.contextmenu) 53 | reload(pe_tree.dump_pe) 54 | reload(pe_tree.exceptions) 55 | reload(pe_tree.hash_pe) 56 | reload(pe_tree.map) 57 | reload(pe_tree.qstandarditems) 58 | reload(pe_tree.tree) 59 | reload(pe_tree.utils) 60 | reload(pe_tree.dialogs) 61 | except NameError: 62 | pass 63 | 64 | import pe_tree.runtime 65 | import pe_tree.form 66 | import pe_tree.info 67 | 68 | class IDARuntime(pe_tree.runtime.Runtime): 69 | """IDA runtime callbacks""" 70 | def __init__(self, widget, args): 71 | # Load configuration 72 | self.config_file = os.path.join(self.get_temp_dir(), "pe_tree.ini") 73 | super(IDARuntime, self).__init__(widget, args) 74 | 75 | def get_temp_dir(self): 76 | """Get path to temporary directory""" 77 | self.ret = idaapi.get_user_idadir() 78 | return self.ret 79 | 80 | def ask_file(self, filename, caption, filter="", save=False): 81 | """Open/save file dialog""" 82 | if not save: 83 | # Open file dialog 84 | self.ret = ida_kernwin.ask_file(0, os.path.basename(str(filename)), str(caption)) 85 | else: 86 | # Save file dialog 87 | if filename[0] == ".": 88 | # Remove leading dot from section names 89 | filename = filename[1:] 90 | 91 | self.ret = ida_kernwin.ask_file(1, os.path.basename(str(filename)), str(caption)) 92 | 93 | return self.ret 94 | 95 | def show_widget(self): 96 | """Display the PE Tree widget""" 97 | self.ret = idaapi.display_widget(self.widget, (idaapi.PluginForm.WOPN_TAB | idaapi.PluginForm.WOPN_MENU | idaapi.PluginForm.WOPN_RESTORE | idaapi.PluginForm.WOPN_PERSIST)) 98 | 99 | return self.ret 100 | 101 | @QtCore.pyqtSlot(object, int) 102 | def jumpto(self, item, offset): 103 | """Jump to offset in IDA view""" 104 | try: 105 | self.ret = idc.jumpto(offset) 106 | except: 107 | self.ret = False 108 | 109 | return self.ret 110 | 111 | @QtCore.pyqtSlot(str) 112 | def log(self, output): 113 | """Print to output""" 114 | print(output) 115 | 116 | @QtCore.pyqtSlot(object, object) 117 | def get_bytes(self, start, size): 118 | """Read bytes from IDB""" 119 | self.ret = idc.get_bytes(start, size) 120 | return self.ret 121 | 122 | @QtCore.pyqtSlot(object) 123 | def get_byte(self, offset): 124 | """Read byte from IDB""" 125 | self.ret = ida_bytes.get_byte(offset) 126 | return self.ret 127 | 128 | @QtCore.pyqtSlot(object) 129 | def get_word(self, offset): 130 | """Read word from IDB""" 131 | self.ret = ida_bytes.get_word(offset) 132 | return self.ret 133 | 134 | @QtCore.pyqtSlot(object) 135 | def get_dword(self, offset): 136 | """Read dword from IDB""" 137 | self.ret = ida_bytes.get_dword(offset) 138 | return self.ret 139 | 140 | @QtCore.pyqtSlot(object) 141 | def get_qword(self, offset): 142 | """Read qword from IDB""" 143 | self.ret = ida_bytes.get_qword(offset) 144 | return self.ret 145 | 146 | @QtCore.pyqtSlot(object) 147 | def get_name(self, offset): 148 | """Get name for offset from IDB""" 149 | self.ret = idc.get_name(offset, ida_name.GN_VISIBLE) 150 | return self.ret 151 | 152 | @QtCore.pyqtSlot(object) 153 | def get_segment_name(self, offset): 154 | """Get segment name for offset from IDB""" 155 | self.ret = idc.get_segm_name(offset) 156 | return self.ret 157 | 158 | @QtCore.pyqtSlot(object) 159 | def is_writable(self, offset): 160 | """Determine if memory address is writable""" 161 | self.ret = bool(idaapi.getseg(offset).perm & idaapi.SEGPERM_WRITE) 162 | return self.ret 163 | 164 | @QtCore.pyqtSlot(object) 165 | def resolve_address(self, offset): 166 | """Get module/symbol name for address""" 167 | symbol = self.get_name(offset) 168 | module = self.get_segment_name(offset) 169 | 170 | if not module and "_" in symbol: 171 | # No module name for the segment, try to determine from the symbol name 172 | symbol_split = symbol.split("_") 173 | 174 | # Given a symbol, i.e. ws2_32_WSAStartup, can we find ws2_32.dll in the list of segments? 175 | for segment in idautils.Segments(): 176 | segment_name = idc.get_segm_name(segment).lower() 177 | 178 | if segment_name.startswith(symbol_split[0].lower()): 179 | new_name = "" 180 | for i in range(0, len(symbol_split)): 181 | new_name = "{}.dll".format("_".join(symbol_split[0:i])) 182 | if new_name == segment_name: 183 | break 184 | 185 | if new_name == segment_name: 186 | module = new_name 187 | break 188 | 189 | # Still nothing?! 190 | if not module and "_" in symbol: 191 | symbol_split = symbol.split("_") 192 | 193 | j = 1 194 | if symbol_split[0] == "ws2": 195 | j += 1 196 | module = "{}.dll".format("_".join(symbol_split[0:j])) 197 | else: 198 | module = "{}.dll".format(symbol_split[0]) 199 | 200 | # Strip module name from symbol name 201 | if module: 202 | module_name = module.split(".")[0].lower() 203 | 204 | if symbol[:len(module_name)].lower().startswith(module_name): 205 | symbol = symbol[len(module_name) + 1:] 206 | 207 | if not symbol: 208 | symbol = "{:x}".format(offset) 209 | 210 | self.ret = (module, symbol) 211 | return self.ret 212 | 213 | @QtCore.pyqtSlot(object) 214 | def get_label(self, offset): 215 | """Get label for offest from IDB""" 216 | self.ret = idc.GetDisasm(offset).replace("extrn ", "").split(":")[0] 217 | return self.ret 218 | 219 | @QtCore.pyqtSlot(object, object) 220 | def make_string(self, offset, size): 221 | """Convert range to byte string in IDB""" 222 | self.ret = idc.create_strlit(offset, offset + size) 223 | return self.ret 224 | 225 | @QtCore.pyqtSlot(object, object, str, str, bytes) 226 | def make_segment(self, offset, size, class_name="DATA", name="pe_map", data=None): 227 | """Create a new section in the IDB""" 228 | idc.AddSeg(offset, offset + size, 0, 1, 0, idaapi.scPub) 229 | idc.RenameSeg(offset, str(name)) 230 | idc.SetSegClass(offset, str(class_name)) 231 | #idc.SegAlign(offset, idc.saRelPara) 232 | if data: 233 | idaapi.patch_many_bytes(offset, bytes(data)) 234 | 235 | self.ret = None 236 | return self.ret 237 | 238 | @QtCore.pyqtSlot(object) 239 | def make_qword(self, offset): 240 | """Create a qword at the given offset in the IDB""" 241 | self.ret = idaapi.create_qword(offset, 8) 242 | return self.ret 243 | 244 | @QtCore.pyqtSlot(object) 245 | def make_dword(self, offset): 246 | """Create a dword at the given offset in the IDB""" 247 | self.ret = idaapi.create_dword(offset, 4) 248 | return self.ret 249 | 250 | @QtCore.pyqtSlot(object) 251 | def make_word(self, offset): 252 | """Create a word at the given offset in the IDB""" 253 | self.ret = idaapi.create_word(offset, 2) 254 | return self.ret 255 | 256 | @QtCore.pyqtSlot(object) 257 | def make_byte(self, offset, size=1): 258 | """Create a byte at the given offset in the IDB""" 259 | self.ret = idaapi.create_byte(offset, size) 260 | return self.ret 261 | 262 | @QtCore.pyqtSlot(object, str) 263 | def make_comment(self, offset, comment): 264 | """Create a comment at the given offset in the IDB""" 265 | #self.ret = idc.MakeComm(offset, comment) 266 | return self.ret 267 | 268 | @QtCore.pyqtSlot(object, str, int) 269 | def make_name(self, offset, name, flags=0): 270 | """Name the given offset in the IDB""" 271 | self.ret = idc.set_name(offset, str(name), idc.SN_NOCHECK | idc.SN_NOWARN | 0x800) 272 | return self.ret 273 | 274 | @QtCore.pyqtSlot() 275 | def get_names(self): 276 | self.ret = list(idautils.Names()) 277 | return self.ret 278 | 279 | @QtCore.pyqtSlot(object, object, object, object) 280 | def find_iat_ptrs(self, pe, image_base, size, get_word): 281 | """Find all likely IAT pointers""" 282 | iat_ptrs = [] 283 | 284 | next_offset = image_base 285 | 286 | while next_offset < image_base + size: 287 | offset = next_offset 288 | next_offset = ida_bytes.next_addr(offset) 289 | 290 | # Attempt to read the current instruction's effective memory address operand (if present) 291 | mnem = idc.print_insn_mnem(offset).lower() 292 | ptr = 0 293 | 294 | if mnem in ["call", "push", "jmp"]: 295 | if idc.get_operand_type(offset, 0) == idc.o_mem: 296 | # Get memory offset for branch instructions 297 | ptr = idc.get_operand_value(offset, 0) 298 | elif mnem in ["mov", "lea"]: 299 | if idc.get_operand_type(offset, 0) == idc.o_reg and idc.get_operand_type(offset, 1) == idc.o_mem: 300 | # Get memory offset for mov/lea instructions 301 | ptr = idc.get_operand_value(offset, 1) 302 | 303 | # Does the instruction's memory address operand seem somewhat valid?! 304 | if ptr < 0x1000: 305 | continue 306 | 307 | # Resolve pointer from memory operand 308 | iat_offset = get_word(ptr) 309 | 310 | # Ignore offset if it is in our image 311 | if image_base <= iat_offset <= image_base + size: 312 | continue 313 | 314 | # Get module and API name for offset 315 | module, api = self.resolve_address(iat_offset) 316 | 317 | # Ignore the offset if it is in a debug segment or stack etc 318 | if api and module and module.endswith(".dll"): 319 | if not iat_offset in iat_ptrs: 320 | # Add IAT offset, address to patch, module name and API name to list 321 | iat_ptrs.append((iat_offset, offset + idc.get_item_size(offset) - 4, module, api)) 322 | 323 | self.ret = iat_ptrs 324 | return self.ret 325 | 326 | @QtCore.pyqtSlot(object) 327 | def find_pe(self, cursor=False): 328 | """Search IDB for possible MZ/PE headers""" 329 | info = idaapi.get_inf_structure() 330 | 331 | mz_headers = [] 332 | 333 | # Check current cursor for MZ/PE? 334 | if cursor: 335 | # Get IDA cursor address 336 | addr = idc.here() 337 | 338 | # Get segment and end address 339 | s = idaapi.getseg(addr) 340 | e = idc.get_segm_end(addr) 341 | 342 | # Check for MZ magic 343 | if ida_bytes.get_word(addr) == 0x5a4d: 344 | # Ensure the PE header is in the segment 345 | e_lfanew = ida_bytes.get_dword(addr + 0x3c) 346 | 347 | if addr + e_lfanew + 1 < e: 348 | # Check for PE magic 349 | if ida_bytes.get_word(addr + e_lfanew) == 0x4550: 350 | # Found possible MZ/PE header 351 | mz_headers.append([addr, idc.get_segm_name(addr), info.is_64bit()]) 352 | 353 | self.ret = mz_headers 354 | return self.ret 355 | 356 | # Search all segments 357 | for seg in idautils.Segments(): 358 | s = idc.get_segm_start(seg) 359 | e = idc.get_segm_end(seg) 360 | addr = s 361 | 362 | while True: 363 | # Find first byte of MZ header 364 | addr = ida_bytes.find_byte(addr, e-addr, 0x4d, 0) 365 | 366 | if addr == ida_idaapi.BADADDR or addr >= e: 367 | break 368 | 369 | # Check for MZ magic 370 | if ida_bytes.get_word(addr) == 0x5a4d: 371 | # Ensure the PE header is in the segment 372 | e_lfanew = ida_bytes.get_dword(addr + 0x3c) 373 | 374 | if addr + e_lfanew + 1 < e: 375 | # Check for PE magic 376 | if ida_bytes.get_word(addr + e_lfanew) == 0x4550: 377 | # Found possible MZ/PE header 378 | mz_headers.append([addr, idc.get_segm_name(s), info.is_64bit()]) 379 | 380 | # Resume search from next address 381 | addr += 1 382 | 383 | self.ret = mz_headers 384 | return self.ret 385 | 386 | class pe_tree_plugin_t(idaapi.plugin_t): 387 | """Main IDAPro plugin class for PE Tree""" 388 | wanted_name = pe_tree.info.__title__ 389 | comment = pe_tree.info.__title__ 390 | version = pe_tree.info.__version__ 391 | website = pe_tree.info.__url__ 392 | help = pe_tree.info.__title__ 393 | wanted_hotkey = "" 394 | flags = 0 395 | 396 | def __init__(self): 397 | self.pe_tree_form = None 398 | 399 | def init(self): 400 | """Initialise PE Tree IDA plugin""" 401 | return ida_idaapi.PLUGIN_KEEP 402 | 403 | def term(self): 404 | """Terminate PE Tree IDA plugin""" 405 | # Wait for plugin threads to complete 406 | if self.pe_tree_form: 407 | self.pe_tree_form.wait_for_threads() 408 | 409 | def run(self, arg): 410 | """Run PE Tree IDA plugin""" 411 | # Get form and widget 412 | form = idaapi.create_empty_widget(self.comment) 413 | widget = sip.wrapinstance(int(form), QtWidgets.QWidget) 414 | 415 | runtime = IDARuntime(form, {}) 416 | 417 | self.pe_tree_form = pe_tree.form.PETreeForm(widget, None, runtime) 418 | 419 | # Try to find the IDA input file 420 | filename = idaapi.get_input_file_path() 421 | 422 | if filename: 423 | if not os.path.isfile(filename): 424 | filename = os.path.basename(filename) 425 | 426 | if not os.path.isfile(filename): 427 | filename = runtime.ask_file(os.path.basename(filename), "Where is the input file?") 428 | 429 | if os.path.isfile(filename): 430 | # Map input file 431 | self.pe_tree_form.map_pe(image_base=idaapi.get_imagebase(), filename=filename) 432 | 433 | self.pe_tree_form.treeview.setVisible(True) 434 | self.pe_tree_form.show() 435 | 436 | def PLUGIN_ENTRY(): 437 | """PE Tree IDAPython plugin entry-point""" 438 | return pe_tree_plugin_t() 439 | 440 | def main(): 441 | """PE Tree IDAPython script entry-point""" 442 | PLUGIN_ENTRY().run("") 443 | 444 | if __name__ == "__main__": 445 | main() 446 | -------------------------------------------------------------------------------- /pe_tree_rekall.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """PE Tree Rekall plugin""" 18 | 19 | # Standard imports 20 | import os 21 | import io 22 | import sys 23 | 24 | # Qt imports 25 | from PyQt5 import QtCore, QtWidgets, QtGui 26 | 27 | # Rekall imports 28 | from rekall_lib import utils 29 | 30 | # PE Tree imports 31 | import pe_tree.window 32 | import pe_tree.runtime 33 | import pe_tree.form 34 | import pe_tree.info 35 | 36 | # pylint: disable=undefined-variable 37 | 38 | class RekallRuntime(pe_tree.runtime.Runtime): 39 | """Rekall runtime callbacks""" 40 | def __init__(self, widget, args): 41 | # Load configuration 42 | self.config_file = os.path.join(self.get_temp_dir(), "pe_tree_rekall.ini") 43 | super(RekallRuntime, self).__init__(widget, args) 44 | 45 | # Initialise Rekall 46 | self.cc = session.plugins.cc() 47 | self.pedump = session.plugins.pedump() 48 | 49 | @QtCore.pyqtSlot(object, object) 50 | def read_pe(self, image_base, size=0): 51 | """Read PE image from memory""" 52 | outfd = io.BytesIO() 53 | 54 | self.pedump.WritePEFile(outfd, self.opaque["address_space"], image_base) 55 | 56 | self.ret = bytearray(outfd.getbuffer()) 57 | 58 | return self.ret 59 | 60 | @QtCore.pyqtSlot(object, object) 61 | def get_bytes(self, start, size): 62 | """Read bytes from memory""" 63 | task = self.opaque["task"] 64 | if task: 65 | self.cc.SwitchProcessContext(self.opaque["task"]) 66 | 67 | self.ret = self.opaque["address_space"].read(start, size) 68 | self.ret = b"" if self.ret is None else self.ret 69 | return self.ret 70 | 71 | @QtCore.pyqtSlot(object) 72 | def is_writable(self, offset): 73 | """Determine if memory address is writable""" 74 | self.ret = False 75 | 76 | task = self.opaque["task"] 77 | if task is not None: 78 | self.cc.SwitchProcessContext(task) 79 | 80 | # Traverse VADs 81 | for vad in task.RealVadRoot.traverse(): 82 | # Is address within the VAD region? 83 | if vad.Start <= offset <= vad.End: 84 | # Is the VAD region writable? 85 | if "WRITE" in str(vad.u.VadFlags.ProtectionEnum): 86 | self.ret = True 87 | 88 | return self.ret 89 | 90 | @QtCore.pyqtSlot(object) 91 | def resolve_address(self, offset): 92 | """Get module/symbol name for address""" 93 | module = "" 94 | api = "" 95 | 96 | task = self.opaque["task"] 97 | if task is not None: 98 | self.cc.SwitchProcessContext(task) 99 | 100 | symbols = session.address_resolver.format_address(offset) 101 | 102 | if len(symbols) > 0: 103 | symbol = symbols[0] 104 | 105 | if "!" in symbol: 106 | split_symbol = symbol.split("!") 107 | module = "{}.dll".format(split_symbol[0]) 108 | api = "".join(split_symbol[1:]) 109 | 110 | self.ret = (module, api) 111 | return self.ret 112 | 113 | class ExpandableStandardItemModel(QtGui.QStandardItemModel): 114 | """Set expandable role on process items to force the drawing of expand/collapse arrows without child items present (allows for lazy loading of modules)""" 115 | ExpandableRole = QtCore.Qt.UserRole + 500 116 | 117 | def hasChildren(self, index): 118 | """Determine if the item has children or is expandable 119 | 120 | Args: 121 | index (QModelIndex): Index of item 122 | 123 | Returns: 124 | bool: True if the item has children or is expandable 125 | 126 | """ 127 | if self.data(index, ExpandableStandardItemModel.ExpandableRole): 128 | return True 129 | 130 | return super(ExpandableStandardItemModel, self).hasChildren(index) 131 | 132 | class SelectProcess(pe_tree.dialogs.ModulePicker): 133 | """Process/module picker dialog""" 134 | def __init__(self, window, parent=None): 135 | super(SelectProcess, self).__init__(window, "Select process...", ["Name", "PID", "Image base"], model=ExpandableStandardItemModel, parent=parent) 136 | 137 | self.treeview.expanded.connect(self.populate_modules) 138 | 139 | def populate_processes(self): 140 | """Populate the tree view with processes""" 141 | root = self.model.invisibleRootItem() 142 | 143 | # Iterate over all processes 144 | for task in session.plugins.pslist(proc_regex="").filter_processes(): 145 | # Determine image base (ignore if null) 146 | image_base = int(task.Peb.ImageBaseAddress) if int(task.Peb.ImageBaseAddress) else 0 147 | 148 | if image_base == 0: 149 | continue 150 | 151 | # Create internal filename 152 | filename = "{} ({}) - {:#08x}".format(str(task.name), task.pid, image_base) 153 | 154 | # Create process tree view item 155 | process_item = QtGui.QStandardItem("{}".format(str(task.name))) 156 | process_item.setData({"filename": filename, "address_space": task.get_process_address_space(), "task": task, "image_base": image_base}, QtCore.Qt.UserRole) 157 | process_item.setUserTristate(True) 158 | process_item.setCheckable(True) 159 | process_item.setData(True, ExpandableStandardItemModel.ExpandableRole) 160 | 161 | # Determine architecture specifics 162 | if task.address_mode == "AMD64": 163 | ptr_mask = 0xffffffffffffffff 164 | ptr_width = 16 165 | else: 166 | ptr_mask = 0xffffffff 167 | ptr_width = 8 168 | 169 | root.appendRow([process_item, QtGui.QStandardItem("{}".format(task.pid)), QtGui.QStandardItem("0x{:0{w}x}".format(image_base & ptr_mask, w=ptr_width))]) 170 | 171 | self.treeview.resizeColumnToContents(0) 172 | 173 | def populate_drivers(self, parent): 174 | """Populate the tree view with drivers""" 175 | # Build a list of all address spaces 176 | address_spaces = [session.kernel_address_space] 177 | for task in session.plugins.pslist().filter_processes(): 178 | address_spaces.append(task.get_process_address_space()) 179 | 180 | # Determine architecture specifics 181 | if session.profile.metadata("arch") == "AMD64": 182 | ptr_mask = 0xffffffffffffffff 183 | ptr_width = 16 184 | else: 185 | ptr_mask = 0xffffffff 186 | ptr_width = 8 187 | 188 | for module in session.plugins.modules().lsmod(): 189 | # Ensure the driver's image base is non-null 190 | image_base = int(module.DllBase) if int(module.DllBase) else 0 191 | 192 | if image_base == 0: 193 | continue 194 | 195 | # Locate the address space 196 | address_space = None 197 | 198 | for a in address_spaces: 199 | if a.is_valid_address(image_base): 200 | address_space = a 201 | 202 | if not address_space: 203 | continue 204 | 205 | # Create internal filename 206 | base_name = os.path.basename(utils.SmartUnicode(module.BaseDllName)) 207 | filename = "{} - 0x{:0x}".format(utils.EscapeForFilesystem(base_name), int(module.DllBase)) 208 | 209 | # Create driver tree view item 210 | driver_item = QtGui.QStandardItem("{}".format(utils.EscapeForFilesystem(base_name))) 211 | driver_item.setData({"filename": filename, "task": None, "address_space": address_space, "image_base": image_base}, QtCore.Qt.UserRole) 212 | driver_item.setCheckable(True) 213 | driver_item.setCheckState(QtCore.Qt.Checked if parent.checkState() == QtCore.Qt.Checked else QtCore.Qt.Unchecked) 214 | if parent.checkState() == QtCore.Qt.Checked: 215 | self.items.add(pe_tree.qstandarditems.HashableQStandardItem(driver_item)) 216 | 217 | parent.appendRow([driver_item, QtGui.QStandardItem("N/A"), QtGui.QStandardItem("0x{:0{w}x}".format(image_base & ptr_mask, w=ptr_width))]) 218 | 219 | def populate_dlls(self, parent, task): 220 | """Populate the tree view with DLLs""" 221 | task_as = task.get_process_address_space() 222 | 223 | # Determine architecture specifics 224 | if task.address_mode == "AMD64": 225 | ptr_mask = 0xffffffffffffffff 226 | ptr_width = 16 227 | else: 228 | ptr_mask = 0xffffffff 229 | ptr_width = 8 230 | 231 | # Iterate over loaded modules (use ldrmodules plugin) 232 | for module in session.plugins.ldrmodules(pids=[task.pid]): 233 | if module.get("in_mem") and module.get("base") > 0: 234 | # Ensure the image base is non-null and not the parent process 235 | image_base = int(module.get("base")) 236 | 237 | if image_base == 0 or image_base == int(task.Peb.ImageBaseAddress): 238 | continue 239 | 240 | # Create internal filename 241 | base_name = os.path.basename(utils.SmartUnicode(module.get("in_mem_path"))) 242 | filename = "{} ({} - {}) - 0x{:0x}".format(utils.EscapeForFilesystem(base_name), task.name, task.pid, image_base) 243 | 244 | # Create DLL tree view item 245 | dll_item = QtGui.QStandardItem("{}".format(utils.EscapeForFilesystem(base_name))) 246 | dll_item.setData({"filename": filename, "task": task, "address_space": task_as, "image_base": image_base}, QtCore.Qt.UserRole) 247 | dll_item.setCheckable(True) 248 | dll_item.setCheckState(QtCore.Qt.Checked if parent.checkState() == QtCore.Qt.Checked else QtCore.Qt.Unchecked) 249 | if parent.checkState() == QtCore.Qt.Checked: 250 | self.items.add(pe_tree.qstandarditems.HashableQStandardItem(dll_item)) 251 | 252 | parent.appendRow([dll_item, QtGui.QStandardItem("{}".format(task.pid)), QtGui.QStandardItem("0x{:0{w}x}".format(image_base & ptr_mask, w=ptr_width))]) 253 | 254 | def populate_modules(self, index): 255 | """Lazy load modules to the tree view when the user expands a process""" 256 | # Find the parent item 257 | parent = self.model.itemFromIndex(index) 258 | args = parent.data(QtCore.Qt.UserRole) 259 | 260 | # Ensure the parent is a root item without any children 261 | if parent.parent() is not None or parent.rowCount() > 0: 262 | return 263 | 264 | # Get the associated task and address space 265 | task = args["task"] 266 | 267 | if task.pid == 4: 268 | # Load drivers under the system process 269 | self.populate_drivers(parent) 270 | else: 271 | # Load DLLs under the parent process 272 | self.populate_dlls(parent, task) 273 | 274 | # Expand the tree view and resize the filename column 275 | self.treeview.expand(index) 276 | self.treeview.resizeColumnToContents(0) 277 | 278 | def invoke(self): 279 | """Display the select process/modules dialog""" 280 | super(SelectProcess, self).invoke() 281 | 282 | # Has the user accepted, i.e. pressed "Dump"? 283 | if self.exec_() != 0: 284 | # Dump selected processes/modules 285 | for item in self.items: 286 | args = item.item.data(QtCore.Qt.UserRole) 287 | self.window.pe_tree_form.map_pe(filename=args["filename"], image_base=args["image_base"], opaque={"address_space": args["address_space"], "task": args["task"]}) 288 | 289 | self.close() 290 | 291 | def main(): 292 | """PE Tree Rekall script entry-point""" 293 | # Create PE Tree Qt application 294 | application = QtWidgets.QApplication(sys.argv) 295 | window = pe_tree.window.PETreeWindow(application, RekallRuntime, {}, open_file=False) 296 | 297 | # Extend menu to include open process 298 | select_process = SelectProcess(window) 299 | 300 | open_process_action = QtWidgets.QAction("Process", window) 301 | open_process_action.setShortcut("Ctrl+Shift+P") 302 | open_process_action.setStatusTip("Open process") 303 | open_process_action.triggered.connect(select_process.invoke) 304 | window.open_menu.addAction(open_process_action) 305 | 306 | # Invoke the process/module picker dialog 307 | select_process.invoke() 308 | 309 | sys.exit(application.exec_()) 310 | 311 | if __name__ == "__main__": 312 | main() 313 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pefile 2 | requests 3 | cryptography 4 | configparser 5 | asn1crypto 6 | scandir -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) 2020 BlackBerry Limited. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | """Install PE Tree standalone application or IDAPython plugin""" 18 | 19 | # Standard imports 20 | import os 21 | import sys 22 | 23 | # Setuptools imports 24 | from setuptools import setup 25 | from setuptools.command.install import install 26 | from setuptools.command.develop import develop 27 | 28 | # PE Tree imports 29 | import pe_tree.info 30 | 31 | # Common requirements for IDAPython plugin and standalone application 32 | with open("requirements.txt") as f: 33 | install_requires = f.read().splitlines() 34 | 35 | if sys.version_info >= (3, 0): 36 | install_requires.extend(["filetype", "minidump"]) 37 | 38 | INSTALL_FOR_IDA = False 39 | 40 | if "--ida" in sys.argv: 41 | sys.argv.remove("--ida") 42 | 43 | # Setup for IDAPython plugin 44 | entry_points = {} 45 | 46 | INSTALL_FOR_IDA = True 47 | else: 48 | # Setup for standalone application entry-points 49 | entry_points = {"console_scripts": ["pe-tree=pe_tree.__main__:main", 50 | "pe-tree-vol=pe_tree.volatility:main", 51 | "pe-tree-ghidra=pe_tree.ghidra:main", 52 | "pe-tree-carve=pe_tree.carve:main", 53 | "pe-tree-minidump=pe_tree.minidump:main"]} 54 | 55 | install_requires.extend(["pyqt5", "capstone"]) 56 | 57 | # Long description 58 | with open(os.path.join(os.path.abspath(os.path.dirname(__file__)), "README.md")) as f: 59 | long_description = "".join([l for l in f.readlines() if not l.startswith("!")]) 60 | 61 | def find_ida(): 62 | """Search for possible IDA Pro installations""" 63 | if INSTALL_FOR_IDA: 64 | paths = ["/opt/ida-7.{}/plugins", 65 | "~/.idapro/plugins", 66 | "%ProgramFiles%\\IDA Pro 7.{}\\plugins"] 67 | 68 | found = [] 69 | 70 | for path in paths: 71 | if "{}" in path: 72 | for i in range(0, 10): 73 | new_path = os.path.expandvars(path.format(i)) 74 | 75 | if os.path.exists(new_path): 76 | found.append(new_path) 77 | else: 78 | if os.path.exists(path): 79 | found.append(path) 80 | 81 | if len(found) > 0: 82 | print("\nTo complete the IDAPython plugin installation copy pe_tree_ida.py to:\n{}".format("\n".join(found))) 83 | 84 | def do_post_setup(): 85 | """Search for IDA installations post setup""" 86 | find_ida() 87 | 88 | class InstallCommand(install): 89 | """Extend install command""" 90 | def run(self): 91 | install.run(self) 92 | do_post_setup() 93 | 94 | class DevelopCommand(develop): 95 | """Extend develop command""" 96 | def run(self): 97 | develop.run(self) 98 | do_post_setup() 99 | 100 | setup(name="pe-tree", 101 | version=pe_tree.info.__version__, 102 | description="View Portable Executable (PE) files in a tree-view using pefile and PyQt5.", 103 | license=pe_tree.info.__license__, 104 | url=pe_tree.info.__url__, 105 | author=pe_tree.info.__author__, 106 | author_email=pe_tree.info.__email__, 107 | packages=["pe_tree"], 108 | py_modules=["pe_tree"], 109 | install_requires=install_requires, 110 | entry_points=entry_points, 111 | cmdclass={"install": InstallCommand, "develop": DevelopCommand}, 112 | python_requires=">=2.7", 113 | long_description=long_description, 114 | long_description_content_type="text/markdown", 115 | keywords=["pe_tree", 116 | "petree", 117 | "pefile", 118 | "exe", 119 | "dll", 120 | "malware", 121 | "pedump", 122 | "unpacker", 123 | "iat", 124 | "rekall", 125 | "volatility" 126 | "minidump"], 127 | classifiers=[ 128 | "Intended Audience :: Developers", 129 | "Intended Audience :: Science/Research", 130 | "Development Status :: 4 - Beta", 131 | "Programming Language :: Python", 132 | "Programming Language :: Python :: 2.7", 133 | "Programming Language :: Python :: 3.5", 134 | "Programming Language :: Python :: 3.6", 135 | "Programming Language :: Python :: 3.7", 136 | "Programming Language :: Python :: 3.8", 137 | "Environment :: X11 Applications :: Qt", 138 | "Environment :: Plugins", 139 | "Operating System :: Microsoft :: Windows", 140 | "Operating System :: MacOS :: MacOS X", 141 | "Operating System :: POSIX :: Linux", 142 | "Topic :: Security", 143 | "Topic :: Utilities", 144 | "Natural Language :: English"], 145 | include_package_data=True 146 | ) 147 | --------------------------------------------------------------------------------