├── requirements.txt ├── tests ├── example_32_pie ├── example_64_pie ├── example_32_nopie ├── example_64_nopie ├── example_64_static ├── example_mingw64.exe ├── example_aarch64_static ├── makefile ├── example.c └── test_patch.py ├── examples └── test_fmt │ ├── t_fmt_32 │ ├── t_fmt_64 │ ├── t_fmt_32.patched │ ├── t_fmt_64.patched │ ├── t_fmt.c │ ├── patch_fmt_64.py │ └── patch_fmt_32.py ├── pwnpatch ├── __init__.py ├── patcher │ ├── __init__.py │ ├── base.py │ ├── elf.py │ └── pe.py ├── exceptions.py ├── main.py ├── enums.py ├── factory.py └── log_utils.py ├── README.md ├── setup.py └── .gitignore /requirements.txt: -------------------------------------------------------------------------------- 1 | keystone-engine 2 | capstone 3 | pyelftools 4 | pefile -------------------------------------------------------------------------------- /tests/example_32_pie: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/tests/example_32_pie -------------------------------------------------------------------------------- /tests/example_64_pie: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/tests/example_64_pie -------------------------------------------------------------------------------- /tests/example_32_nopie: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/tests/example_32_nopie -------------------------------------------------------------------------------- /tests/example_64_nopie: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/tests/example_64_nopie -------------------------------------------------------------------------------- /tests/example_64_static: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/tests/example_64_static -------------------------------------------------------------------------------- /examples/test_fmt/t_fmt_32: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/examples/test_fmt/t_fmt_32 -------------------------------------------------------------------------------- /examples/test_fmt/t_fmt_64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/examples/test_fmt/t_fmt_64 -------------------------------------------------------------------------------- /pwnpatch/__init__.py: -------------------------------------------------------------------------------- 1 | from pwnpatch.main import get_patcher 2 | 3 | __all__ = [ 4 | get_patcher 5 | ] -------------------------------------------------------------------------------- /tests/example_mingw64.exe: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/tests/example_mingw64.exe -------------------------------------------------------------------------------- /tests/example_aarch64_static: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/tests/example_aarch64_static -------------------------------------------------------------------------------- /examples/test_fmt/t_fmt_32.patched: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/examples/test_fmt/t_fmt_32.patched -------------------------------------------------------------------------------- /examples/test_fmt/t_fmt_64.patched: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/veritas501/pwnpatch/HEAD/examples/test_fmt/t_fmt_64.patched -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pwnpatch 2 | 3 | 用于ctf pwn的patch玩具轮子。 4 | 5 | ## Install 6 | 7 | ``` 8 | python3 -m pip install . 9 | ``` 10 | 11 | ## Quick Start 12 | 13 | 参考examples和tests。 -------------------------------------------------------------------------------- /pwnpatch/patcher/__init__.py: -------------------------------------------------------------------------------- 1 | from .base import BasePatcher 2 | from .elf import ElfPatcher 3 | from .pe import PePatcher 4 | 5 | __all__ = [BasePatcher, ElfPatcher, PePatcher] 6 | -------------------------------------------------------------------------------- /examples/test_fmt/t_fmt.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | int main(void){ 6 | char buf[0x100]; 7 | memset(buf,0,0x100); 8 | read(0,buf,0x100); 9 | printf(buf); 10 | return 0; 11 | } 12 | -------------------------------------------------------------------------------- /examples/test_fmt/patch_fmt_64.py: -------------------------------------------------------------------------------- 1 | import pwnpatch 2 | 3 | pt = pwnpatch.get_patcher('./t_fmt_64') 4 | pt.add_byte('%s\x00', 'new_fmt_str') 5 | asm_code = ''' 6 | mov rsi,rdi 7 | mov rdi, {new_fmt_str} 8 | ''' 9 | pt.hook_asm(0x400681, asm_code) 10 | pt.save() 11 | -------------------------------------------------------------------------------- /examples/test_fmt/patch_fmt_32.py: -------------------------------------------------------------------------------- 1 | import pwnpatch 2 | 3 | pt = pwnpatch.get_patcher('./t_fmt_32') 4 | pt.add_byte('%s\x00', 'new_fmt_str') 5 | asm_code = ''' 6 | pop eax 7 | push {new_fmt_str} 8 | push eax 9 | ''' 10 | pt.hook_asm(0x08048526, asm_code) 11 | pt.save() 12 | -------------------------------------------------------------------------------- /tests/makefile: -------------------------------------------------------------------------------- 1 | all: 2 | rm -f example_* 3 | gcc example.c -o example_64_pie -pie 4 | gcc example.c -o example_64_nopie -no-pie 5 | gcc example.c -o example_32_pie -pie -m32 6 | gcc example.c -o example_32_nopie -no-pie -m32 7 | gcc example.c -o example_64_static -static 8 | aarch64-linux-gnu-gcc example.c -o example_aarch64_static -static 9 | x86_64-w64-mingw32-gcc example.c -o example_mingw64 -------------------------------------------------------------------------------- /tests/example.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | int func1(){ 7 | return getpid() == 0x1337; 8 | } 9 | 10 | int func2(){ 11 | return getpid() == 0xdeadbeef; 12 | } 13 | 14 | int func3(){ 15 | char str2[]="this_is_str2"; 16 | return !strcmp("str1",str2); 17 | } 18 | 19 | int main(void){ 20 | if(func1() && func2() && func3()){ 21 | return 0; 22 | } 23 | return 1; 24 | } 25 | -------------------------------------------------------------------------------- /pwnpatch/exceptions.py: -------------------------------------------------------------------------------- 1 | class UnsupportedBinaryException(Exception): 2 | """ 3 | 不支持此格式的二进制 4 | """ 5 | 6 | pass 7 | 8 | 9 | class TODOException(Exception): 10 | """ 11 | 还不支持这个功能 12 | """ 13 | 14 | pass 15 | 16 | 17 | class PatchException(Exception): 18 | """ 19 | Patch的报错 20 | """ 21 | 22 | pass 23 | 24 | 25 | class AddException(Exception): 26 | """ 27 | Add的报错 28 | """ 29 | 30 | pass 31 | 32 | 33 | class SccException(Exception): 34 | """ 35 | 调用scc工具时发生的异常 36 | """ 37 | 38 | pass 39 | -------------------------------------------------------------------------------- /pwnpatch/main.py: -------------------------------------------------------------------------------- 1 | from typing import Union 2 | 3 | from pwnpatch.factory import PatcherFactory 4 | from pwnpatch import patcher 5 | 6 | VERSION = "v1.0.2" 7 | 8 | 9 | def get_patcher( 10 | filename: str, minimal_edit=False 11 | ) -> Union[patcher.ElfPatcher, patcher.PePatcher]: 12 | """get patcher object 13 | 14 | Args: 15 | filename (str): filename to patch 16 | minimal_edit (bool, optional): enable minimal edit mode. Defaults to False. 17 | 18 | Returns: 19 | patcher class 20 | """ 21 | return PatcherFactory.get_patcher(filename, minimal_edit) 22 | -------------------------------------------------------------------------------- /pwnpatch/enums.py: -------------------------------------------------------------------------------- 1 | from enum import Enum 2 | 3 | 4 | class Endian(Enum): 5 | NONE = 0 6 | LSB = 1 7 | MSB = 2 8 | 9 | 10 | class Arch(Enum): 11 | NONE = 0 12 | I386 = 1 13 | AMD64 = 2 14 | ARM = 3 15 | ARM64 = 4 16 | MIPS = 5 17 | 18 | 19 | class BitSize(Enum): 20 | NONE = 0 21 | X32 = 32 22 | X64 = 64 23 | 24 | 25 | class EXE(Enum): 26 | NONE = 0 27 | ELF = 1 28 | PE = 2 29 | 30 | 31 | class SegProt(Enum): 32 | NONE = 0 33 | X = 1 34 | W = 2 35 | R = 4 36 | RWX = 7 37 | RX = 5 38 | RW = 6 39 | WX = 3 40 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from setuptools import find_packages, setup 4 | 5 | # Based on https://github.com/mitmproxy/mitmproxy/blob/main/setup.py 6 | 7 | here = os.path.abspath(os.path.dirname(__file__)) 8 | 9 | # get version 10 | with open(os.path.join(here, "pwnpatch", "main.py")) as f: 11 | match = re.search(r'VERSION = "(.+?)"', f.read()) 12 | if not match: 13 | raise Exception("Can't find version string") 14 | VERSION = match.group(1) 15 | 16 | # get requirements 17 | with open(os.path.join(here, "requirements.txt")) as f: 18 | req = [] 19 | for r in f.read().strip().split("\n"): 20 | if not r.startswith('#'): 21 | req.append(r) 22 | 23 | setup( 24 | name="pwnpatch", 25 | version=VERSION, 26 | description="pwnpatch: ctf pwn patch tool", 27 | packages=find_packages( 28 | include=["pwnpatch", "pwnpatch.*"] 29 | ), 30 | include_package_data=True, 31 | python_requires=">=3.6", 32 | install_requires=req 33 | ) -------------------------------------------------------------------------------- /pwnpatch/factory.py: -------------------------------------------------------------------------------- 1 | from typing import Union 2 | 3 | from pwnpatch import enums 4 | from pwnpatch import exceptions 5 | from pwnpatch import patcher 6 | 7 | 8 | class PatcherFactory: 9 | @classmethod 10 | def is_pefile(cls, filename: str) -> bool: 11 | with open(filename, "rb") as f: 12 | magic = f.read(2) 13 | return magic == b"MZ" 14 | 15 | @classmethod 16 | def is_elf(cls, filename: str) -> bool: 17 | with open(filename, "rb") as f: 18 | magic = f.read(4) 19 | return magic == b"\x7fELF" 20 | 21 | @classmethod 22 | def get_binary_type(cls, filename: str): 23 | if cls.is_pefile(filename): 24 | return enums.EXE.PE 25 | elif cls.is_elf(filename): 26 | return enums.EXE.ELF 27 | raise exceptions.UnsupportedBinaryException("Not a PE or ELF file") 28 | 29 | @classmethod 30 | def get_patcher( 31 | cls, filename: str, minimal_edit: bool 32 | ) -> Union[patcher.ElfPatcher, patcher.PePatcher]: 33 | exe_type = cls.get_binary_type(filename) 34 | if exe_type == enums.EXE.PE: 35 | return patcher.PePatcher(filename, minimal_edit=minimal_edit) 36 | elif exe_type == enums.EXE.ELF: 37 | return patcher.ElfPatcher(filename, minimal_edit=minimal_edit) 38 | raise exceptions.UnsupportedBinaryException("unknown exe type") 39 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # idea 2 | .idea 3 | 4 | # Byte-compiled / optimized / DLL files 5 | __pycache__/ 6 | *.py[cod] 7 | *$py.class 8 | 9 | # C extensions 10 | *.so 11 | 12 | # Distribution / packaging 13 | .Python 14 | build/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | wheels/ 26 | pip-wheel-metadata/ 27 | share/python-wheels/ 28 | *.egg-info/ 29 | .installed.cfg 30 | *.egg 31 | MANIFEST 32 | 33 | # PyInstaller 34 | # Usually these files are written by a python script from a template 35 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 36 | *.manifest 37 | *.spec 38 | 39 | # Installer logs 40 | pip-log.txt 41 | pip-delete-this-directory.txt 42 | 43 | # Unit test / coverage reports 44 | htmlcov/ 45 | .tox/ 46 | .nox/ 47 | .coverage 48 | .coverage.* 49 | .cache 50 | nosetests.xml 51 | coverage.xml 52 | *.cover 53 | *.py,cover 54 | .hypothesis/ 55 | .pytest_cache/ 56 | 57 | # Translations 58 | *.mo 59 | *.pot 60 | 61 | # Django stuff: 62 | *.log 63 | local_settings.py 64 | db.sqlite3 65 | db.sqlite3-journal 66 | 67 | # Flask stuff: 68 | instance/ 69 | .webassets-cache 70 | 71 | # Scrapy stuff: 72 | .scrapy 73 | 74 | # Sphinx documentation 75 | docs/_build/ 76 | 77 | # PyBuilder 78 | target/ 79 | 80 | # Jupyter Notebook 81 | .ipynb_checkpoints 82 | 83 | # IPython 84 | profile_default/ 85 | ipython_config.py 86 | 87 | # pyenv 88 | .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 98 | __pypackages__/ 99 | 100 | # Celery stuff 101 | celerybeat-schedule 102 | celerybeat.pid 103 | 104 | # SageMath parsed files 105 | *.sage.py 106 | 107 | # Environments 108 | .env 109 | .venv 110 | env/ 111 | venv/ 112 | ENV/ 113 | env.bak/ 114 | venv.bak/ 115 | 116 | # Spyder project settings 117 | .spyderproject 118 | .spyproject 119 | 120 | # Rope project settings 121 | .ropeproject 122 | 123 | # mkdocs documentation 124 | /site 125 | 126 | # mypy 127 | .mypy_cache/ 128 | .dmypy.json 129 | dmypy.json 130 | 131 | # Pyre type checker 132 | .pyre/ -------------------------------------------------------------------------------- /pwnpatch/log_utils.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import sys 3 | 4 | __all__ = ["Log", "log", "use_style"] 5 | 6 | STYLE = { 7 | "fore": { # 前景色 8 | "black": 30, # 黑色 9 | "red": 31, # 红色 10 | "green": 32, # 绿色 11 | "yellow": 33, # 黄色 12 | "blue": 34, # 蓝色 13 | "purple": 35, # 紫红色 14 | "cyan": 36, # 青蓝色 15 | "white": 37, # 白色 16 | }, 17 | "back": { # 背景 18 | "black": 40, # 黑色 19 | "red": 41, # 红色 20 | "green": 42, # 绿色 21 | "yellow": 43, # 黄色 22 | "blue": 44, # 蓝色 23 | "purple": 45, # 紫红色 24 | "cyan": 46, # 青蓝色 25 | "white": 47, # 白色 26 | }, 27 | "mode": { # 显示模式 28 | "normal": 0, # 终端默认设置 29 | "bold": 1, # 高亮显示 30 | "underline": 4, # 使用下划线 31 | "blink": 5, # 闪烁 32 | "invert": 7, # 反白显示 33 | "hide": 8, # 不可见 34 | }, 35 | "default": { 36 | "end": 0, 37 | }, 38 | } 39 | 40 | 41 | def use_style(string, mode="", fore="", back=""): 42 | if sys.stdout.isatty(): 43 | style_list = [ 44 | STYLE["mode"].get(mode), 45 | STYLE["fore"].get(fore), 46 | STYLE["back"].get(back), 47 | ] 48 | style_list = list(filter(lambda x: x is not None, style_list)) 49 | if not style_list: 50 | style_list = [0] 51 | style = "\033[{}m".format(";".join(map(str, style_list))) 52 | end = "\033[{}m".format(STYLE["default"]["end"]) 53 | return "{}{}{}".format(style, string, end) 54 | else: 55 | return string 56 | 57 | 58 | msg_prefixes = { 59 | "success": use_style("+", "bold", "green"), 60 | "failure": use_style("-", "bold", "red"), 61 | "debug": use_style("DEBUG", "bold", "red"), 62 | "info": use_style("*", "bold", "blue"), 63 | "warning": use_style("!", "bold", "yellow"), 64 | "error": use_style("ERROR", "normal", "red"), 65 | } 66 | 67 | 68 | class Log: 69 | def __init__(self, logger=None): 70 | if not logger: 71 | self.logger = logging.getLogger(__name__) 72 | else: 73 | self.logger = logger 74 | 75 | def _log(self, level, s): 76 | self.logger.log(level, s) 77 | 78 | def set_level(self, level): 79 | self.logger.setLevel(level) 80 | 81 | def get_level(self): 82 | return logging.getLevelName(self.logger.level) 83 | 84 | def success(self, s): 85 | self._log(logging.INFO, "[{}] {}".format(msg_prefixes["success"], s)) 86 | 87 | def failure(self, s): 88 | self._log(logging.INFO, "[{}] {}".format(msg_prefixes["failure"], s)) 89 | 90 | def fail(self, s): 91 | self.failure(s) 92 | 93 | def debug(self, s): 94 | self._log(logging.DEBUG, "[{}] {}".format(msg_prefixes["debug"], s)) 95 | 96 | def info(self, s): 97 | self._log(logging.INFO, "[{}] {}".format(msg_prefixes["info"], s)) 98 | 99 | def warning(self, s): 100 | self._log(logging.WARNING, "[{}] {}".format(msg_prefixes["warning"], s)) 101 | 102 | def warn(self, s): 103 | self.warning(s) 104 | 105 | def error(self, s): 106 | self._log(logging.ERROR, "[{}] {}".format(msg_prefixes["error"], s)) 107 | 108 | def underline(self, s): 109 | return use_style(s, "underline") 110 | 111 | 112 | logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) 113 | log = Log() 114 | -------------------------------------------------------------------------------- /tests/test_patch.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import subprocess as sp 3 | 4 | from pwnpatch import get_patcher 5 | from pwnpatch.patcher import PePatcher 6 | 7 | 8 | def test_64_pie(): 9 | p = get_patcher("./sample/example_64_pie") 10 | # func1 use patch 11 | asm = ''' 12 | mov rax,1 13 | ret 14 | ''' 15 | p.patch_asm(0x1189, asm) 16 | # func2 use hook 17 | asm = ''' 18 | mov rax,0xdeadbeef 19 | ''' 20 | p.hook_asm(0x11B0, asm) 21 | # func3 use add 22 | p.add_byte("this_is_str2\x00", "str2") 23 | p.patch_asm(0x11f8, "lea rdi, [{str2}]") 24 | # test add_segment 25 | p.add_segment() 26 | p.add_segment(prot=5) 27 | p.add_segment() 28 | p.save("patch_result") 29 | io = sp.Popen("./patch_result") 30 | io.communicate() 31 | assert io.returncode == 0 32 | 33 | 34 | def test_64_nopie(): 35 | p = get_patcher("./sample/example_64_nopie") 36 | # func1 use patch 37 | asm = ''' 38 | mov rax,1 39 | ret 40 | ''' 41 | p.patch_asm(0x401176, asm) 42 | # func2 use hook 43 | asm = ''' 44 | mov rax,0xdeadbeef 45 | ''' 46 | p.hook_asm(0x40119D, asm) 47 | # func3 use add 48 | p.add_byte("this_is_str2\x00", "str2") 49 | p.patch_asm(0x4011E5, "lea rdi, [{str2}]") 50 | # test add_segment 51 | p.add_segment() 52 | p.add_segment(prot=5) 53 | p.add_segment() 54 | p.save("patch_result") 55 | io = sp.Popen("./patch_result") 56 | io.communicate() 57 | assert io.returncode == 0 58 | 59 | 60 | def test_32_pie(): 61 | p = get_patcher("./sample/example_32_pie") 62 | # func1 use patch 63 | asm = ''' 64 | mov eax,1 65 | ret 66 | ''' 67 | p.patch_asm(0x120D, asm) 68 | # func2 use hook 69 | asm = ''' 70 | mov eax,0xdeadbeef 71 | ''' 72 | p.hook_asm(0x1256, asm) 73 | # func3 use add 74 | p.add_byte("this_is_str2\x00", "str2") 75 | p.add_asm("mov edx,[esp];ret", "pc_thunk_dx") 76 | asm = ''' 77 | call {pc_thunk_dx} 78 | pc: sub edx, pc - {str2} 79 | push edx 80 | jmp 0x12AF 81 | ''' 82 | p.add_asm(asm, "patch1") 83 | p.patch_asm(0x12A8, "jmp {patch1}") 84 | # test add_segment 85 | p.add_segment() 86 | p.add_segment(prot=5) 87 | p.add_segment() 88 | p.save("patch_result") 89 | io = sp.Popen("./patch_result") 90 | io.communicate() 91 | assert io.returncode == 0 92 | 93 | 94 | def test_32_nopie(): 95 | p = get_patcher("./sample/example_32_nopie") 96 | # func1 use patch 97 | asm = ''' 98 | mov eax,1 99 | ret 100 | ''' 101 | p.patch_asm(0x080491D6, asm) 102 | # func2 use hook 103 | asm = ''' 104 | mov eax,0xdeadbeef 105 | ''' 106 | p.hook_asm(0x0804921F, asm) 107 | # func3 use add 108 | p.add_byte("this_is_str2\x00", "str2") 109 | p.add_asm("mov edx,[esp];ret", "pc_thunk_dx") 110 | asm = ''' 111 | call {pc_thunk_dx} 112 | pc: sub edx, pc - {str2} 113 | push edx 114 | jmp 0x08049278 115 | ''' 116 | p.add_asm(asm, "patch1") 117 | p.patch_asm(0x08049271, "jmp {patch1}") 118 | # test add_segment 119 | p.add_segment() 120 | p.add_segment(prot=5) 121 | p.add_segment() 122 | p.save("patch_result") 123 | io = sp.Popen("./patch_result") 124 | io.communicate() 125 | assert io.returncode == 0 126 | 127 | 128 | def test_64_static(): 129 | p = get_patcher("./sample/example_64_static") 130 | # func1 use patch 131 | asm = ''' 132 | mov rax,1 133 | ret 134 | ''' 135 | p.patch_asm(0x401CE5, asm) 136 | # func2 use hook 137 | asm = ''' 138 | mov rax,0xdeadbeef 139 | ''' 140 | p.hook_asm(0x401D0C, asm) 141 | # func3 use add 142 | p.add_byte("this_is_str2\x00", "str2") 143 | p.patch_asm(0x401D54, "lea rdi, [{str2}]") 144 | # test add_segment 145 | p.add_segment() 146 | p.add_segment(prot=5) 147 | p.add_segment() 148 | p.save("patch_result") 149 | io = sp.Popen("./patch_result") 150 | io.communicate() 151 | assert io.returncode == 0 152 | 153 | 154 | def test_aarch64_static(): 155 | p = get_patcher("./sample/example_aarch64_static") 156 | # func1 use patch 157 | asm = ''' 158 | mov x0,1 159 | ret 160 | ''' 161 | p.patch_asm(0x4006AC, asm) 162 | # TODO: aarch64 not support hook 163 | # func2 use hook patch 164 | asm = ''' 165 | CSET W0, NE 166 | ''' 167 | p.patch_asm(0x4006F0, asm) 168 | # func3 use add 169 | str2_p = p.add_byte("this_is_str2\x00", "str2") 170 | p.patch_asm(0x400740, 171 | f"adrp x0, {str2_p & ~0xfff}; add x0, x0, {str2_p & 0xfff}") 172 | # test add_segment 173 | p.add_segment() 174 | p.add_segment(prot=5) 175 | p.add_segment() 176 | p.save("patch_result") 177 | io = sp.Popen(["qemu-aarch64-static", "./patch_result"]) 178 | io.communicate() 179 | assert io.returncode == 0 180 | 181 | 182 | def test_mingw64(): 183 | p: PePatcher = get_patcher("./sample/example_mingw64.exe") 184 | # func1 use patch 185 | asm = ''' 186 | mov rax,1 187 | ret 188 | ''' 189 | p.patch_asm(0x401560, asm) 190 | # func2 use hook< patch 191 | asm = ''' 192 | mov rax,1 193 | ''' 194 | p.hook_asm(0x401593, asm) 195 | # func3 use add 196 | str2_p = p.add_byte("this_is_str2\x00", "str2") 197 | p.patch_asm(0x4015C4, "lea rcx, [{str2}]") 198 | # test add_segment 199 | p.add_segment() 200 | p.add_segment(prot=5) 201 | p.add_segment() 202 | p.save("patch_result.exe") 203 | io = sp.Popen("./patch_result.exe") 204 | io.communicate() 205 | assert io.returncode == 0 206 | 207 | 208 | if __name__ == '__main__': 209 | test_64_pie() 210 | test_64_nopie() 211 | test_32_pie() 212 | test_32_nopie() 213 | test_64_static() 214 | test_aarch64_static() 215 | test_mingw64() 216 | -------------------------------------------------------------------------------- /pwnpatch/patcher/base.py: -------------------------------------------------------------------------------- 1 | import os 2 | from abc import ABC, abstractmethod 3 | from typing import Optional, Dict, List, Union, Tuple 4 | 5 | import capstone as cs 6 | import keystone as ks 7 | 8 | from pwnpatch import enums 9 | from pwnpatch import exceptions 10 | from pwnpatch.log_utils import log 11 | 12 | 13 | class BasePatcher(ABC): 14 | def __init__(self, filename: str) -> None: 15 | # 文件信息 16 | self.rel_binary_path = filename 17 | self.abs_binary_path = os.path.abspath(self.rel_binary_path) 18 | self.abs_dir, self.binary_name = os.path.split(self.abs_binary_path) 19 | 20 | # 基础信息 basic info 21 | self.exe_type: enums.EXE = enums.EXE.NONE 22 | self.arch: enums.Arch = enums.Arch.NONE 23 | self.endian: enums.Endian = enums.Endian.NONE 24 | self.bit_size: enums.BitSize = enums.BitSize.NONE 25 | 26 | # CS and KS 27 | self.cs: Optional[cs.Cs] = None 28 | self.ks: Optional[ks.Ks] = None 29 | 30 | # raw binary data 31 | self._raw_binary_data: Optional[bytearray] = None 32 | 33 | # 符号表 34 | self.symbols: Dict = {} 35 | 36 | # internal variables 37 | self._code_cave: List[List[int]] = [] 38 | 39 | # 初始化patcher 40 | self._init_basic_info() 41 | self._init_cs_ks() 42 | 43 | def _init_cs_ks(self) -> bool: 44 | self._init_cs() 45 | self._init_ks() 46 | if self.cs is None: 47 | raise Exception("capstone init failed") 48 | if self.ks is None: 49 | raise Exception("keystone init failed") 50 | return True 51 | 52 | def _init_ks(self) -> bool: 53 | ks_arch = ks.KS_ARCH_X86 54 | ks_mode = ( 55 | ks.KS_MODE_LITTLE_ENDIAN 56 | if self.endian == enums.Endian.LSB 57 | else ks.KS_MODE_BIG_ENDIAN 58 | ) 59 | if self.arch == enums.Arch.AMD64: 60 | ks_arch = ks.KS_ARCH_X86 61 | ks_mode |= ks.KS_MODE_64 62 | elif self.arch == enums.Arch.I386: 63 | ks_arch = ks.KS_ARCH_X86 64 | ks_mode |= ks.KS_MODE_32 65 | elif self.arch == enums.Arch.MIPS: 66 | ks_arch = ks.KS_ARCH_MIPS 67 | ks_mode |= ( 68 | ks.KS_MODE_32 if self.bit_size == enums.BitSize.X32 else ks.KS_MODE_64 69 | ) 70 | elif self.arch == enums.Arch.ARM64: 71 | ks_arch = ks.KS_ARCH_ARM64 72 | elif self.arch == enums.Arch.ARM: 73 | ks_arch = ks.KS_ARCH_ARM 74 | ks_mode |= ks.KS_MODE_ARM 75 | else: 76 | raise exceptions.UnsupportedBinaryException("Unsupported binary") 77 | self.ks = ks.Ks(ks_arch, ks_mode) 78 | 79 | return True 80 | 81 | def _init_cs(self) -> bool: 82 | cs_arch = cs.CS_ARCH_X86 83 | cs_mode = ( 84 | cs.CS_MODE_LITTLE_ENDIAN 85 | if self.endian == enums.Endian.LSB 86 | else cs.CS_MODE_BIG_ENDIAN 87 | ) 88 | if self.arch == enums.Arch.AMD64: 89 | cs_arch = cs.CS_ARCH_X86 90 | cs_mode |= cs.CS_MODE_64 91 | elif self.arch == enums.Arch.I386: 92 | cs_arch = cs.CS_ARCH_X86 93 | cs_mode |= cs.CS_MODE_32 94 | elif self.arch == enums.Arch.MIPS: 95 | cs_arch = cs.CS_ARCH_MIPS 96 | cs_mode |= cs.CS_MODE_32 if self.bit_size == 32 else cs.CS_MODE_64 97 | elif self.arch == enums.Arch.ARM64: 98 | cs_arch = cs.CS_ARCH_ARM64 99 | elif self.arch == enums.Arch.ARM: 100 | cs_arch = cs.CS_ARCH_ARM 101 | cs_mode |= cs.CS_MODE_ARM 102 | else: 103 | raise exceptions.UnsupportedBinaryException("Unsupported binary") 104 | self.cs = cs.Cs(cs_arch, cs_mode) 105 | 106 | return True 107 | 108 | @abstractmethod 109 | def _init_basic_info(self) -> bool: 110 | pass 111 | 112 | @abstractmethod 113 | def _init_code_cave(self) -> bool: 114 | pass 115 | 116 | def add_range_to_code_cave(self, rva: int, size: int) -> bool: 117 | new_cave_start = rva 118 | new_cave_end = rva + size 119 | 120 | # find overlap 121 | for i, (old_cave_start, old_cave_size) in enumerate(self._code_cave): 122 | old_cave_end = old_cave_start + old_cave_size 123 | # if new code cave overlap with the old one, update the old cave 124 | if new_cave_start <= old_cave_end and new_cave_end >= old_cave_start: 125 | final_cave_start = min(new_cave_start, old_cave_start) 126 | final_cave_end = max(new_cave_end, old_cave_end) 127 | self._code_cave[i] = [ 128 | final_cave_start, 129 | final_cave_end - final_cave_start, 130 | ] 131 | return True 132 | 133 | # no overlap 134 | self._code_cave.append([rva, size]) 135 | return True 136 | 137 | @property 138 | @abstractmethod 139 | def _can_add_segment(self) -> bool: 140 | pass 141 | 142 | @property 143 | @abstractmethod 144 | def image_base(self) -> int: 145 | pass 146 | 147 | def _set_file_bytes(self, offset: int, data: Union[str, bytes, bytearray]) -> None: 148 | if offset < 0: 149 | raise Exception("set_file_bytes offset < 0, offset: {}".format(offset)) 150 | if isinstance(data, str): 151 | data = data.encode("latin-1") 152 | if offset > len(self.binary_data): 153 | self.binary_data += b"\x00" * (offset - len(self.binary_data)) 154 | self.binary_data[offset : offset + len(data)] = bytearray(data) 155 | 156 | def _get_file_bytes(self, offset: int, size: int) -> bytearray: 157 | if offset < 0: 158 | raise Exception("get_file_bytes offset < 0, offset: {}".format(offset)) 159 | if offset + size > len(self.binary_data): 160 | raise Exception( 161 | "get_file_bytes offset+size too big, val: {}, max_size: {}".format( 162 | offset + size, len(self.binary_data) 163 | ) 164 | ) 165 | return self.binary_data[offset : offset + size] 166 | 167 | def _get_byte(self, rva: int, size: int) -> bytearray: 168 | start_file_offset = self.rva_to_offset(rva) 169 | _end_file_offset = self.rva_to_offset(rva + size - 1) 170 | return self._get_file_bytes(start_file_offset, size) 171 | 172 | def _patch_byte(self, rva: int, data: Union[str, bytes, bytearray]): 173 | # 检查地址align 174 | if self.need_align and (rva % 4): 175 | # 因为patch的地址是用户传递的，因此这里只能warn一下 176 | log.warn("address {} is not 4 bytes aligned".format(hex(rva))) 177 | 178 | # 检查是否有对应的 rva 179 | start_file_offset = self.rva_to_offset(rva) 180 | _end_file_offset = self.rva_to_offset(rva + len(data) - 1) 181 | 182 | self._set_file_bytes(start_file_offset, data) 183 | 184 | def _add_byte(self, data: Union[str, bytes, bytearray]) -> int: 185 | need_length = len(data) 186 | target_rva = 0 187 | for i in range(len(self._code_cave)): 188 | add_rva, add_length = self._code_cave[i] 189 | # 判断align 190 | if self.need_align: 191 | if add_rva % 4: 192 | add_length -= 4 - add_rva % 4 193 | add_rva += 4 - add_rva % 4 194 | if add_length >= need_length: 195 | target_rva = add_rva 196 | self._code_cave[i] = [add_rva + need_length, add_length - need_length] 197 | break 198 | 199 | # 现有code cave不够大 200 | if not target_rva: 201 | # 当前格式不支持添加新段 202 | if not self._can_add_segment: 203 | raise exceptions.AddException("no enough space") 204 | 205 | seg_length = self.align(need_length) 206 | new_seg = self.add_segment(0, seg_length, 7) 207 | if new_seg: 208 | # 添加新段，并添加到code cave 209 | self._code_cave.append([new_seg, seg_length]) 210 | # 递归调用 211 | return self._add_byte(data) 212 | else: 213 | raise exceptions.AddException("Generate new memory failed") 214 | start_file_offset = self.rva_to_offset(target_rva) 215 | self._set_file_bytes(start_file_offset, data) 216 | return target_rva 217 | 218 | def _add_asm(self, asm: Union[str, bytes, bytearray]) -> int: 219 | target_rva = 0 220 | need_length = 0 221 | byte_code = bytes() 222 | for i in range(len(self._code_cave)): 223 | add_vaddr, add_length = self._code_cave[i] 224 | # 判断align 225 | if self.need_align: 226 | if add_vaddr % 4: 227 | add_length -= 4 - add_vaddr % 4 228 | add_vaddr += 4 - add_vaddr % 4 229 | # asm需要先预定位置才能判断长度 230 | byte_code, inst_cnt = self._asm(asm, add_vaddr) 231 | if not inst_cnt: 232 | return 0 233 | need_length = len(byte_code) 234 | if add_length >= need_length: 235 | target_rva = add_vaddr 236 | self._code_cave[i] = [add_vaddr + need_length, add_length - need_length] 237 | break 238 | 239 | # 现有code cave不够大 240 | if not target_rva: 241 | # 当前格式不支持添加新段 242 | if not self._can_add_segment: 243 | raise exceptions.AddException("no enough space") 244 | 245 | if not need_length: 246 | # 没有target_vaddr，因此拿image base顶一下算出大概的need_length 247 | byte_code, inst_cnt = self._asm(asm, self.image_base) 248 | if not inst_cnt: 249 | return 0 250 | need_length = len(byte_code) 251 | seg_length = self.align(need_length) 252 | new_seg = self.add_segment(0, seg_length, 7) 253 | if new_seg: 254 | # 添加新段，并添加到code cave 255 | self._code_cave.append([new_seg, seg_length]) 256 | # 递归调用 257 | return self._add_asm(asm) 258 | else: 259 | raise exceptions.AddException("Generate new memory failed") 260 | 261 | start_file_offset = self.rva_to_offset(target_rva) 262 | self._set_file_bytes(start_file_offset, byte_code) 263 | return target_rva 264 | 265 | def _label(self, rva: int, label_name: str) -> bool: 266 | """ 267 | 给指定地址起别名 268 | :param rva: 虚拟地址 269 | :param label_name: 名字，后续在asm中可以以{name}的方式引用此地址 270 | :return: 271 | """ 272 | if not label_name: 273 | return False 274 | 275 | self.symbols[label_name] = rva 276 | return True 277 | 278 | def _asm(self, asm: Union[str, bytes, bytearray], rva: int) -> Tuple[bytes, int]: 279 | """ 280 | 将汇编代码转换为机器码 281 | :param asm: 汇编代码 282 | :param rva: 虚拟地址 283 | :return: (机器码，机器码长度) 284 | """ 285 | if isinstance(asm, bytes) or isinstance(asm, bytearray): 286 | asm = asm.decode("latin-1") 287 | 288 | try: 289 | asm_formatted = asm.format(**self.symbols) 290 | except KeyError as ex: 291 | log.fail("Can't find symbol {}".format(ex)) 292 | return bytes(), 0 293 | 294 | return self.ks.asm(asm_formatted, rva, True) 295 | 296 | def dump_code_cave(self): 297 | log.info("Dumping code cave info:") 298 | for i, (cave_start, cave_size) in enumerate(self._code_cave): 299 | log.info( 300 | "\tcave-{}: 0x{:x} - 0x{:x} (0x{:x} bytes)".format( 301 | i, cave_start, cave_start + cave_size, cave_size 302 | ) 303 | ) 304 | 305 | @property 306 | def need_align(self) -> bool: 307 | """ 308 | 判断是否需要对齐 309 | :return: 310 | """ 311 | return self.arch not in [enums.Arch.AMD64, enums.Arch.I386] 312 | 313 | @property 314 | def binary_data(self) -> bytearray: 315 | if self._raw_binary_data is None: 316 | self._raw_binary_data = bytearray(open(self.abs_binary_path, "rb").read()) 317 | return self._raw_binary_data 318 | 319 | @binary_data.setter 320 | def binary_data(self, data: bytearray) -> None: 321 | self._raw_binary_data = data 322 | 323 | @property 324 | def binary_size(self) -> int: 325 | return len(self.binary_data) 326 | 327 | @abstractmethod 328 | def rva_to_offset(self, rva: int) -> int: 329 | pass 330 | 331 | @abstractmethod 332 | def add_segment(self, addr: int = 0, length: int = 0x1000, prot: int = 7) -> int: 333 | pass 334 | 335 | @staticmethod 336 | def align(rva: int, alignment=0x1000) -> int: 337 | return (rva + (alignment - 1)) - (rva + (alignment - 1)) % alignment 338 | 339 | def patch_byte( 340 | self, vaddr: int, byte: Union[str, bytes, bytearray], label: str = None 341 | ) -> bool: 342 | try: 343 | origin_bytes = self._get_byte(vaddr, len(byte)) 344 | self._patch_byte(vaddr, byte) 345 | patched_bytes = self._get_byte(vaddr, len(byte)) 346 | self._label(vaddr, label) 347 | log.success("Patch {} bytes @ {}".format(hex(len(byte)), hex(vaddr))) 348 | self.hexdump_diff(origin_bytes, patched_bytes, vaddr) 349 | except exceptions.PatchException as ex: 350 | log.fail(ex) 351 | return False 352 | 353 | return True 354 | 355 | def patch_asm( 356 | self, 357 | vaddr: int, 358 | asm: str, 359 | label: str = None, 360 | ljust: Tuple[int, Union[bytes, bytearray]] = (0, b"\x00"), 361 | ) -> bool: 362 | fill_size = ljust[0] 363 | fill_byte = ljust[1] 364 | byte_code, inst_cnt = self._asm(asm, vaddr) 365 | if not inst_cnt: 366 | return False 367 | 368 | try: 369 | if fill_size and len(byte_code) > fill_size: 370 | log.fail( 371 | "asm result size {} > fill_size {}".format( 372 | len(byte_code), fill_size 373 | ) 374 | ) 375 | return False 376 | fetch_size = ( 377 | len(byte_code) if not fill_size else max(fill_size, len(byte_code)) 378 | ) 379 | origin_bytes = self._get_byte(vaddr, fetch_size) 380 | self._patch_byte(vaddr, byte_code) 381 | if fill_size: 382 | self._patch_byte( 383 | vaddr + len(byte_code), 384 | bytes((fill_byte[0],)) * (fill_size - len(byte_code)), 385 | ) 386 | patched_bytes = self._get_byte(vaddr, fetch_size) 387 | self._label(vaddr, label) 388 | log.success( 389 | "Patch {} bytes asm @ {}".format(hex(len(byte_code)), hex(vaddr)) 390 | ) 391 | self.hexdump_diff(origin_bytes, patched_bytes, vaddr) 392 | except exceptions.PatchException as ex: 393 | log.fail(ex) 394 | return False 395 | 396 | return True 397 | 398 | def add_byte(self, byte: Union[str, bytes, bytearray], label: str = None) -> int: 399 | try: 400 | target_vaddr = self._add_byte(byte) 401 | except exceptions.AddException as ex: 402 | log.fail(ex) 403 | return 0 404 | if not target_vaddr: 405 | log.fail( 406 | "Add {} bytes @ {} failed".format(hex(len(byte)), hex(target_vaddr)) 407 | ) 408 | return 0 409 | self._label(target_vaddr, label) 410 | log.success("Add {} bytes @ {}".format(hex(len(byte)), hex(target_vaddr))) 411 | return target_vaddr 412 | 413 | def add_asm(self, asm: str, label: str = None) -> int: 414 | try: 415 | target_vaddr = self._add_asm(asm) 416 | except exceptions.AddException as ex: 417 | log.fail(ex) 418 | return 0 419 | if not target_vaddr: 420 | log.fail("Add asm @ {} failed".format(hex(target_vaddr))) 421 | return 0 422 | self._label(target_vaddr, label) 423 | log.success("Add asm @ {}".format(hex(target_vaddr))) 424 | return target_vaddr 425 | 426 | def label(self, vaddr: int, label_name: str) -> bool: 427 | if self._label(vaddr, label_name): 428 | log.success("name {} with {}".format(hex(vaddr), label_name)) 429 | return True 430 | log.fail("name {} with {} failed".format(hex(vaddr), label_name)) 431 | return False 432 | 433 | def save(self, filename: str = None) -> bool: 434 | """ 435 | 保存patch后的binary 436 | :param filename: 保存文件名（默认值为 <源文件名>.patched<.后缀>） 437 | :return: 438 | """ 439 | if not filename: 440 | bin_name = self.binary_name 441 | bin_postfix = os.path.splitext(bin_name)[-1].lower() 442 | bin_true_name = "".join(os.path.splitext(bin_name)[:-1]) 443 | filename = "{}.patched{}".format(bin_true_name, bin_postfix) 444 | 445 | with open(filename, "wb") as f: 446 | f.write(self.binary_data) 447 | 448 | os.chmod(filename, 0o755) 449 | log.success( 450 | "Binary saved @ {}".format(log.underline(os.path.abspath(filename))) 451 | ) 452 | 453 | return True 454 | 455 | def hexdump(self, data: Union[str, bytes, bytearray], length: int = 16): 456 | if isinstance(data, str): 457 | data = data.encode("latin-1") 458 | elif not isinstance(data, (bytes, bytearray)): 459 | raise TypeError("data must be str, bytes, or bytearray") 460 | 461 | for i in range(0, len(data), length): 462 | chunk = data[i : i + length] 463 | hex_parts = [] 464 | for j in range(0, len(chunk), 4): 465 | group = chunk[j : j + 4] 466 | hex_parts.append(" ".join(f"{byte:02x}" for byte in group)) 467 | hex_line = " ".join(hex_parts) 468 | 469 | print(f"{i:08x} {hex_line}") 470 | 471 | def hexdump_diff( 472 | self, 473 | data1: Union[str, bytes, bytearray], 474 | data2: Union[str, bytes, bytearray], 475 | base_address: int = 0, 476 | length: int = 16, 477 | ): 478 | """ 479 | 以 hexdump 格式对比两份数据，并展示差异。 480 | 481 | :param data1: 第一份数据 (str, bytes, or bytearray) 482 | :param data2: 第二份数据 (str, bytes, or bytearray) 483 | :param length: 每行显示的字节数，默认 16 484 | """ 485 | 486 | # 转换字符串为字节 487 | def to_bytes(data): 488 | if isinstance(data, str): 489 | return data.encode("latin-1") 490 | if not isinstance(data, (bytes, bytearray)): 491 | raise TypeError("data must be str, bytes, or bytearray") 492 | return data 493 | 494 | # 转换数据为字节序列 495 | data1 = to_bytes(data1) 496 | data2 = to_bytes(data2) 497 | 498 | # 确保两份数据长度一致（长数据用短数据填充） 499 | if len(data1) > len(data2): 500 | data2 += data1[len(data2) :] 501 | elif len(data2) > len(data1): 502 | data1 += data2[len(data1) :] 503 | 504 | max_len = max(len(data1), len(data2)) 505 | 506 | for i in range(0, max_len, length): 507 | chunk1 = data1[i : i + length] 508 | chunk2 = data2[i : i + length] 509 | 510 | # 格式化第一份数据：每 4 字节添加空格 511 | def format_chunk(chunk): 512 | groups = [ 513 | " ".join( 514 | f"{byte:02x}" if byte is not None else ".." 515 | for byte in chunk[j : j + 4] 516 | ) 517 | for j in range(0, len(chunk), 4) 518 | ] 519 | return " ".join(groups) 520 | 521 | hex_line1 = format_chunk(chunk1) 522 | print(f"{i+base_address:08x} {hex_line1}") 523 | 524 | # 第二行：标记差异 525 | diff_line = [] 526 | for b1, b2 in zip(chunk1, chunk2): 527 | if b1 == b2: 528 | diff_line.append(None) 529 | else: 530 | diff_line.append(b2) 531 | diff_line_str = format_chunk(diff_line) 532 | print(f">>>>>>>> {diff_line_str}") 533 | -------------------------------------------------------------------------------- /pwnpatch/patcher/elf.py: -------------------------------------------------------------------------------- 1 | import io 2 | from typing import Optional, Union 3 | 4 | from elftools.elf import constants 5 | from elftools.elf.elffile import ELFFile 6 | from elftools.elf.sections import Section 7 | from elftools.elf.segments import Segment 8 | 9 | from pwnpatch import enums 10 | from pwnpatch import exceptions 11 | from pwnpatch.log_utils import log 12 | from pwnpatch.patcher.base import BasePatcher 13 | 14 | 15 | class ElfPatcher(BasePatcher): 16 | def __init__(self, filename: str, minimal_edit=False) -> None: 17 | self._elffile_stream = open(filename, "rb") 18 | self._elffile: ELFFile = ELFFile(self._elffile_stream) 19 | 20 | # internal variables 21 | self._phdr_reallocated: bool = False 22 | 23 | super().__init__(filename) 24 | if not minimal_edit: 25 | self._init_code_cave() 26 | 27 | def _init_basic_info(self) -> bool: 28 | # 初始化基础信息 29 | self.exe_type = enums.EXE.ELF 30 | e_machine = self._elffile.header.e_machine 31 | self.arch = { 32 | "EM_386": enums.Arch.I386, 33 | "EM_X86_64": enums.Arch.AMD64, 34 | "EM_ARM": enums.Arch.ARM, 35 | "EM_AARCH64": enums.Arch.ARM64, 36 | "EM_MIPS": enums.Arch.MIPS, 37 | }.get(e_machine, enums.Arch.NONE) 38 | if self.arch == enums.Arch.NONE: 39 | raise exceptions.UnsupportedBinaryException( 40 | "Unsupported arch: {}".format(e_machine) 41 | ) 42 | is_little_endian = self._elffile.little_endian 43 | self.endian = enums.Endian.LSB if is_little_endian else enums.Endian.MSB 44 | bit_size = self._elffile.header.e_ident["EI_CLASS"] 45 | self.bit_size = { 46 | "ELFCLASS32": enums.BitSize.X32, 47 | "ELFCLASS64": enums.BitSize.X64, 48 | }.get(bit_size, enums.BitSize.NONE) 49 | if self.bit_size == enums.BitSize.NONE: 50 | raise exceptions.UnsupportedBinaryException( 51 | "Unsupported bit size: {}".format(bit_size) 52 | ) 53 | 54 | return True 55 | 56 | def _init_code_cave(self) -> bool: 57 | eh_frame_section = self._elffile.get_section_by_name(".eh_frame") 58 | if not eh_frame_section: 59 | return False 60 | 61 | # 如果 .eh_frame的权限为只读，需要patch一下segment的权限 +x 62 | eh_frame_segment = self._find_segment_hold_section(eh_frame_section) 63 | if not eh_frame_segment: 64 | log.warning( 65 | "Can't find which segment hold .eh_frame section, skip add .eh_frame section to code cave" 66 | ) 67 | return True 68 | if not eh_frame_segment.header.p_flags & constants.P_FLAGS.PF_X: 69 | log.success("Modify .eh_frame section prot +x") 70 | segment_idx = self._get_phdr_idx(eh_frame_segment) 71 | eh_frame_segment.header.p_flags |= constants.P_FLAGS.PF_X 72 | self._replace_phdr_by_idx(eh_frame_segment, segment_idx) 73 | self._reload() 74 | 75 | self.add_range_to_code_cave( 76 | eh_frame_section["sh_addr"], eh_frame_section["sh_size"] 77 | ) 78 | log.info("Add .eh_frame to code cave") 79 | return True 80 | 81 | @property 82 | def _phoff(self) -> int: 83 | return self._elffile.header.e_phoff 84 | 85 | @property 86 | def _phnum(self) -> int: 87 | return self._elffile.header.e_phnum 88 | 89 | @property 90 | def _phentsize(self) -> int: 91 | return self._elffile.header.e_phentsize 92 | 93 | @property 94 | def image_base(self) -> int: 95 | min_addr = (1 << 64) - 1 96 | for segment in self._elffile.iter_segments(): 97 | if segment.header.p_type == "PT_LOAD": 98 | min_addr = min(min_addr, segment.header.p_vaddr) 99 | return min_addr 100 | 101 | def _find_segment_hold_section(self, section: Section) -> Optional[Segment]: 102 | for segment in self._elffile.iter_segments(): 103 | if segment.section_in_segment(section): 104 | return segment 105 | 106 | return None 107 | 108 | def _get_phdr_idx(self, segment: Segment) -> int: 109 | target_phdr_data = self._elffile.structs.Elf_Phdr.build(segment.header) 110 | 111 | for i in range(self._elffile.num_segments()): 112 | seg = self._elffile.get_segment(i) 113 | tmp_phdr_data = self._elffile.structs.Elf_Phdr.build(seg.header) 114 | if target_phdr_data == tmp_phdr_data: 115 | return i 116 | 117 | raise Exception("can't find target phdr") 118 | 119 | def _replace_phdr_by_idx(self, segment: Segment, idx: int) -> bool: 120 | if idx > self._phnum: 121 | log.warning("segment index out of range: {} > {}".format(idx, self._phnum)) 122 | return False 123 | 124 | phdr_file_offset = self._phoff + self._phentsize * idx 125 | phdr_data = self._elffile.structs.Elf_Phdr.build(segment.header) 126 | self._set_file_bytes(phdr_file_offset, phdr_data) 127 | return True 128 | 129 | @property 130 | def _can_add_segment(self) -> bool: 131 | return True 132 | 133 | def _add_segment(self, addr: int = 0, length: int = 0x1000, prot: int = 7) -> int: 134 | if not self._phdr_reallocated: 135 | self._alloc_new_segment_for_phdr() 136 | 137 | if not addr: 138 | # 自动计算地址 139 | new_seg = self._auto_next_segment(length, prot) 140 | if not new_seg: 141 | return 0 142 | else: 143 | new_seg = self._elffile.structs.Elf_Phdr.parse( 144 | bytes(self._elffile.structs.Elf_Phdr.sizeof()) 145 | ) 146 | new_seg.p_type = "PT_LOAD" 147 | new_seg.p_flags = prot 148 | new_seg.p_offset = self.align(self.binary_size) 149 | new_seg.p_vaddr = addr 150 | new_seg.p_paddr = addr 151 | new_seg.p_filesz = length 152 | new_seg.p_memsz = length 153 | new_seg.p_align = 0x1000 154 | 155 | # segment内容用全0填充 156 | self._set_file_bytes(new_seg.p_offset, bytes(new_seg.p_filesz)) 157 | 158 | pht_data = bytearray() 159 | for i in range(self._elffile.num_segments()): 160 | seg = self._elffile.get_segment(i) 161 | phdr = seg.header 162 | if phdr.p_type == "PT_PHDR": 163 | # 修改PHDR的size 164 | phdr.p_filesz += self._phentsize 165 | phdr.p_memsz += self._phentsize 166 | pht_data += self._elffile.structs.Elf_Phdr.build(phdr) 167 | pht_data += self._elffile.structs.Elf_Phdr.build(new_seg) 168 | 169 | # 写入新 phdr 数据 170 | self._set_file_bytes(self._phoff, pht_data) 171 | # 更新并写入 ehdr 172 | self._elffile.header.e_phnum += 1 173 | new_ehdr = self._elffile.structs.Elf_Ehdr.build(self._elffile.header) 174 | self._set_file_bytes(0, new_ehdr) 175 | self._reload() 176 | 177 | log.success( 178 | "Add segment @ {}({}) {}{}{}".format( 179 | hex(new_seg.p_vaddr), 180 | hex(new_seg.p_memsz), 181 | "r" if new_seg.p_flags & constants.P_FLAGS.PF_R else "-", 182 | "w" if new_seg.p_flags & constants.P_FLAGS.PF_W else "-", 183 | "x" if new_seg.p_flags & constants.P_FLAGS.PF_X else "-", 184 | ) 185 | ) 186 | return new_seg.p_vaddr 187 | 188 | def _auto_next_segment(self, length: int, prot: int): 189 | aligned_file_end = self.align(self.binary_size) 190 | next_segment_addr = aligned_file_end 191 | for i in range(self._elffile.num_segments()): 192 | seg = self._elffile.get_segment(i) 193 | if seg.header.p_type == "PT_LOAD": 194 | seg_start = seg.header.p_vaddr - self.image_base 195 | seg_end = self.align(seg_start + seg.header.p_memsz) 196 | if seg_start > aligned_file_end: 197 | continue 198 | next_segment_addr = max(next_segment_addr, seg_end) 199 | 200 | res_seg = self._elffile.structs.Elf_Phdr.parse( 201 | bytes(self._elffile.structs.Elf_Phdr.sizeof()) 202 | ) 203 | res_seg.p_type = "PT_LOAD" 204 | res_seg.p_flags = prot 205 | res_seg.p_offset = next_segment_addr 206 | res_seg.p_vaddr = next_segment_addr + self.image_base 207 | res_seg.p_paddr = next_segment_addr + self.image_base 208 | res_seg.p_filesz = length 209 | res_seg.p_memsz = length 210 | res_seg.p_align = 0x1000 211 | return res_seg 212 | 213 | def _alloc_new_segment_for_phdr(self) -> int: 214 | new_seg_for_phdr = self._auto_next_segment(0x1000, 5) 215 | self._set_file_bytes(new_seg_for_phdr.p_paddr, bytes(new_seg_for_phdr.p_filesz)) 216 | pht_data = bytearray() 217 | for i in range(self._elffile.num_segments()): 218 | seg = self._elffile.get_segment(i) 219 | phdr = seg.header 220 | if phdr.p_type == "PT_PHDR": 221 | # 修改PHDR file_offset 和 size 222 | phdr.p_offset = new_seg_for_phdr.p_offset 223 | phdr.p_paddr = new_seg_for_phdr.p_paddr 224 | phdr.p_vaddr = new_seg_for_phdr.p_vaddr 225 | phdr.p_filesz += self._phentsize 226 | phdr.p_memsz += self._phentsize 227 | pht_data += self._elffile.structs.Elf_Phdr.build(phdr) 228 | pht_data += self._elffile.structs.Elf_Phdr.build(new_seg_for_phdr) 229 | # 写入新 phdr 数据 230 | self._set_file_bytes(new_seg_for_phdr.p_offset, pht_data) 231 | # 更新并写入 ehdr 232 | self._elffile.header.e_phoff = new_seg_for_phdr.p_offset 233 | self._elffile.header.e_phnum += 1 234 | new_ehdr = self._elffile.structs.Elf_Ehdr.build(self._elffile.header) 235 | self._set_file_bytes(0, new_ehdr) 236 | self._reload() 237 | 238 | self._phdr_reallocated = True 239 | return new_seg_for_phdr.p_paddr 240 | 241 | def _segment_from_virtual_address(self, virtual_addr: int) -> Segment: 242 | for seg in self._elffile.iter_segments(): 243 | if seg.header.p_type == "PT_LOAD": 244 | if ( 245 | seg.header.p_vaddr 246 | <= virtual_addr 247 | < seg.header.p_vaddr + seg.header.p_memsz 248 | ): 249 | return seg 250 | 251 | raise Exception( 252 | "Can't found segment from virtual address: 0x{:x}".format(virtual_addr) 253 | ) 254 | 255 | def _reload(self) -> bool: 256 | """ 257 | 当patch对header等重要数据做出修改后，需要更新一下elffile 258 | :return: 259 | """ 260 | self._elffile = ELFFile(io.BytesIO(self.binary_data)) 261 | return True 262 | 263 | def _hook_byte_intel(self, vaddr: int, byte: Union[str, bytes, bytearray]) -> bool: 264 | if isinstance(byte, str): 265 | byte = byte.encode("latin-1") 266 | if isinstance(byte, bytes): 267 | byte = bytearray(byte) 268 | byte += b"\xc3" # append ret 269 | self.add_byte(byte, "user_hook_{}".format(hex(vaddr))) 270 | return self._hook_intel(vaddr) 271 | 272 | def _hook_asm_intel(self, vaddr: int, asm: str) -> bool: 273 | asm += ";ret;" # append ret 274 | addr = self._add_asm(asm) 275 | if not addr: 276 | return False 277 | self._label(addr, "user_hook_{}".format(hex(vaddr))) 278 | return self._hook_intel(vaddr) 279 | 280 | def _hook_intel(self, vaddr: int) -> bool: 281 | x86_jmp_inst_length = 5 282 | disasm_data = bytearray(self._get_byte(vaddr, 0x20)) 283 | cs_disasm = self.cs.disasm(disasm_data, vaddr) 284 | origin1_jmp_addr = vaddr 285 | tmp_len = 0 286 | for i, disasm in enumerate(cs_disasm): 287 | if i == 0 or tmp_len < x86_jmp_inst_length: 288 | tmp_len += len(disasm.bytes) 289 | origin1_jmp_length = tmp_len 290 | origin2_jmp_addr = vaddr + origin1_jmp_length 291 | backup_length = origin1_jmp_length + x86_jmp_inst_length 292 | backup1 = self._add_byte(bytearray(backup_length)) 293 | if not backup1: 294 | return False 295 | self._label(backup1, "backup1_{}".format(hex(vaddr))) 296 | backup2 = self._add_byte(bytearray(backup_length)) 297 | if not backup2: 298 | return False 299 | self._label(backup2, "backup2_{}".format(hex(vaddr))) 300 | if self.arch == enums.Arch.I386: 301 | detour1_asm = """ 302 | pushad 303 | call get_pc 304 | pc: mov esi,edi 305 | sub edi, pc - {} 306 | sub esi, pc - {{backup2_{}}} 307 | mov ecx, {} /* patch_byte_len */ 308 | cld 309 | rep movsb 310 | popad 311 | call {{user_hook_{}}} /* hook_func */ 312 | jmp {} /* origin1 */ 313 | get_pc: 314 | mov edi, [esp] 315 | ret 316 | """.format(vaddr, hex(vaddr), backup_length, hex(vaddr), origin1_jmp_addr) 317 | 318 | detour2_asm = """ 319 | pushf 320 | pushad 321 | call get_pc 322 | pc: mov esi,edi 323 | sub edi, pc - {} 324 | sub esi, pc - {{backup1_{}}} 325 | mov ecx, {} /* patch_byte_len */ 326 | cld 327 | rep movsb 328 | popad 329 | popf 330 | jmp {} /* origin2 */ 331 | get_pc: 332 | mov edi, [esp] 333 | ret 334 | """.format(vaddr, hex(vaddr), backup_length, origin2_jmp_addr) 335 | elif self.arch == enums.Arch.AMD64: 336 | detour1_asm = """ 337 | push rdi /* backup reg */ 338 | push rsi 339 | push rcx 340 | lea rdi, [{}] 341 | lea rsi, [{{backup2_{}}}] 342 | mov rcx, {} /* patch_byte_len */ 343 | cld 344 | rep movsb 345 | pop rcx /* restore reg */ 346 | pop rsi 347 | pop rdi 348 | call {{user_hook_{}}} /* hook_func */ 349 | jmp {} /* origin1 */ 350 | """.format(vaddr, hex(vaddr), backup_length, hex(vaddr), origin1_jmp_addr) 351 | 352 | detour2_asm = """ 353 | pushf 354 | push rdi /* backup reg */ 355 | push rsi 356 | push rcx 357 | lea rdi, [{}] 358 | lea rsi, [{{backup1_{}}}] 359 | mov rcx, {} /* patch_byte_len */ 360 | cld 361 | rep movsb 362 | pop rcx /* restore reg */ 363 | pop rsi 364 | pop rdi 365 | popf 366 | jmp {} /* origin2 */ 367 | """.format(vaddr, hex(vaddr), backup_length, origin2_jmp_addr) 368 | else: 369 | # should not be here 370 | return False 371 | 372 | detour1 = self._add_asm(detour1_asm) 373 | if not detour1: 374 | return False 375 | self._label(detour1, "detour1_{}".format(hex(vaddr))) 376 | detour2 = self._add_asm(detour2_asm) 377 | if not detour2: 378 | return False 379 | self._label(detour2, "detour2_{}".format(hex(vaddr))) 380 | origin1_jmp, inst_cnt = self._asm( 381 | "jmp {{detour1_{}}}".format(hex(vaddr)), origin1_jmp_addr 382 | ) 383 | origin2_jmp, inst_cnt = self._asm( 384 | "jmp {{detour2_{}}}".format(hex(vaddr)), origin2_jmp_addr 385 | ) 386 | backup1_bytes = bytearray(self._get_byte(vaddr, backup_length)) 387 | backup1_bytes[: len(origin1_jmp)] = origin1_jmp 388 | backup2_bytes = bytearray(self._get_byte(vaddr, backup_length)) 389 | backup2_bytes[origin1_jmp_length:] = origin2_jmp 390 | self._patch_byte(backup1, backup1_bytes) 391 | self._patch_byte(backup2, backup2_bytes) 392 | self._patch_byte(vaddr, backup1_bytes) 393 | segment = self._segment_from_virtual_address(vaddr) 394 | if not segment.header.p_flags & constants.P_FLAGS.PF_W: 395 | log.success("Modify hook segment prot +w") 396 | segment_idx = self._get_phdr_idx(segment) 397 | segment.header.p_flags |= constants.P_FLAGS.PF_W 398 | self._replace_phdr_by_idx(segment, segment_idx) 399 | self._reload() 400 | return True 401 | 402 | def rva_to_offset(self, rva: int) -> int: 403 | for segment in self._elffile.iter_segments(): 404 | # 检查虚拟地址是否在该段的范围内 405 | if segment["p_vaddr"] <= rva < segment["p_vaddr"] + segment["p_memsz"]: 406 | # 计算偏移量并返回 407 | return rva - segment["p_vaddr"] + segment["p_offset"] 408 | 409 | # 如果没有找到对应的段，可能是节中的地址 410 | for section in self._elffile.iter_sections(): 411 | # 检查虚拟地址是否在该节的范围内 412 | if section["sh_addr"] <= rva < section["sh_addr"] + section["sh_size"]: 413 | # 计算偏移量并返回 414 | return rva - section["sh_addr"] + section["sh_offset"] 415 | 416 | raise exceptions.PatchException("no such virtual address: 0x{:x}".format(rva)) 417 | 418 | def add_segment(self, addr: int = 0, length: int = 0x1000, prot: int = 7) -> int: 419 | """ 420 | 新增一个段 421 | :param addr: 虚拟地址，默认为0表示自动计算 422 | :param length: 长度，默认0x1000 423 | :param prot: 保护权限，默认rwx 424 | :return: 成功则返回申请的地址，失败返回0 425 | """ 426 | if addr & 0xFFF: 427 | log.fail("address 0x{:x} is not page aligned".format(addr)) 428 | return 0 429 | return self._add_segment(addr, length, prot) 430 | 431 | def set_execstack(self, enabled: bool) -> bool: 432 | """ 433 | 开启或关闭execstack（只支持ELF） 434 | :param enabled: 是否开启execstack 435 | :return: 436 | """ 437 | 438 | for i in range(self._elffile.num_segments()): 439 | seg = self._elffile.get_segment(i) 440 | if seg.header.p_type == "PT_GNU_STACK": 441 | if enabled: 442 | seg.header.p_flags = ( 443 | constants.P_FLAGS.PF_R 444 | | constants.P_FLAGS.PF_W 445 | | constants.P_FLAGS.PF_X 446 | ) 447 | else: 448 | seg.header.p_flags = constants.P_FLAGS.PF_R | constants.P_FLAGS.PF_W 449 | 450 | # 更新ph_data 451 | if not self._replace_phdr_by_idx(seg, i): 452 | log.fail("Modify stack phdr failed") 453 | return False 454 | 455 | log.success( 456 | "Set execstack -> {}".format("enabled" if enabled else "disabled") 457 | ) 458 | return True 459 | 460 | log.fail("Can't find GNU_STACK segment") 461 | return False 462 | 463 | def set_norelro(self) -> bool: 464 | """ 465 | 取消ELF的relro 466 | :return: 467 | """ 468 | 469 | for i in range(self._elffile.num_segments()): 470 | seg = self._elffile.get_segment(i) 471 | if seg.header.p_type == "PT_GNU_RELRO": 472 | seg.header.p_type = "PT_NULL" 473 | if not self._replace_phdr_by_idx(seg, i): 474 | log.fail("Modify phdr failed") 475 | return False 476 | log.success("Set norelro -> enabled") 477 | return True 478 | 479 | log.fail("Can't find GNU_RELRO segment") 480 | return False 481 | 482 | def hook_byte( 483 | self, vaddr: int, byte: Union[str, bytes, bytearray], label: str = None 484 | ) -> bool: 485 | """ 486 | 在指定地址插入bytes作为hook代码 487 | :param vaddr: 虚拟地址 488 | :param byte: bytes数组 489 | :param label: 此hook的名字，后续在asm中可以以{name}的方式引用这个hook的地址 490 | :return: 491 | """ 492 | if self.arch in [enums.Arch.I386, enums.Arch.AMD64]: 493 | if self._hook_byte_intel(vaddr, byte): 494 | self._label(vaddr, label) 495 | log.success("Hook @ {}".format(hex(vaddr))) 496 | return True 497 | else: 498 | log.fail("Hook @ {} failed".format(hex(vaddr))) 499 | return False 500 | else: 501 | raise exceptions.TODOException("Only support hook i386 & amd64 for now") 502 | 503 | def hook_asm(self, vaddr: int, asm: str, label: str = None) -> bool: 504 | """ 505 | 在指定地址插入asm作为hook代码 506 | :param vaddr: 虚拟地址 507 | :param asm: 汇编代码 508 | :param label: 此hook的名字，后续在asm中可以以{name}的方式引用这个hook的地址 509 | :return: 510 | """ 511 | if self.arch in [enums.Arch.I386, enums.Arch.AMD64]: 512 | if self._hook_asm_intel(vaddr, asm): 513 | self._label(vaddr, label) 514 | log.success("Hook @ {}".format(hex(vaddr))) 515 | return True 516 | else: 517 | log.fail("Hook @ {} failed".format(hex(vaddr))) 518 | return False 519 | else: 520 | raise exceptions.TODOException("Only support hook i386 & amd64 for now") 521 | -------------------------------------------------------------------------------- /pwnpatch/patcher/pe.py: -------------------------------------------------------------------------------- 1 | import math 2 | from collections import namedtuple 3 | from typing import Tuple, Union 4 | 5 | import pefile 6 | from pefile import PE 7 | 8 | from pwnpatch import enums 9 | from pwnpatch import exceptions 10 | from pwnpatch.log_utils import log 11 | from pwnpatch.patcher.base import BasePatcher 12 | 13 | RealignedSection = namedtuple( 14 | "RealignedSection", 15 | [ 16 | "name", 17 | "src_offset", 18 | "src_size", 19 | "src_buf", 20 | "dst_offset", 21 | "dst_size", 22 | "dst_padding", 23 | ], 24 | ) 25 | 26 | 27 | class PePatcher(BasePatcher): 28 | def __init__(self, filename: str, need_detail=True, minimal_edit=False): 29 | self.fast_load = not need_detail 30 | 31 | with open(filename, "rb") as f: 32 | self._pefile: PE = pefile.PE(data=f.read(), fast_load=self.fast_load) 33 | 34 | super().__init__(filename) 35 | if not minimal_edit: 36 | self._init_code_cave() 37 | 38 | def _init_basic_info(self): 39 | # 初始化基础信息 40 | self.exe_type = enums.EXE.PE 41 | machine = self._pefile.NT_HEADERS.FILE_HEADER.Machine 42 | self.arch = { 43 | pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_I386"]: enums.Arch.I386, 44 | pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"]: enums.Arch.AMD64, 45 | pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_ARM"]: enums.Arch.ARM, 46 | pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_ARM64"]: enums.Arch.ARM64, 47 | }.get(machine, enums.Arch.NONE) 48 | if self.arch == enums.Arch.NONE: 49 | raise exceptions.UnsupportedBinaryException( 50 | "Unsupported arch: 0x{:x} ({})".format( 51 | machine, pefile.MACHINE_TYPE[machine] 52 | ) 53 | ) 54 | self.endian = enums.Endian.LSB 55 | self.bit_size = { 56 | pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_I386"]: enums.BitSize.X32, 57 | pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"]: enums.BitSize.X64, 58 | pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_ARM"]: enums.BitSize.X32, 59 | pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_ARM64"]: enums.BitSize.X64, 60 | }.get(machine, enums.BitSize.NONE) 61 | if self.bit_size == enums.BitSize.NONE: 62 | raise exceptions.UnsupportedBinaryException( 63 | "Unsupported arch: 0x{:x} ({})".format( 64 | machine, pefile.MACHINE_TYPE[machine] 65 | ) 66 | ) 67 | 68 | return True 69 | 70 | @property 71 | def has_overlay(self): 72 | """ 73 | 判断PE文件末尾是否有overlay数据 74 | :return: 75 | """ 76 | return self._pefile.get_overlay_data_start_offset() is not None 77 | 78 | @property 79 | def binary_data(self) -> bytearray: 80 | if not isinstance(self._pefile.__data__, bytearray): 81 | self._pefile.__data__ = bytearray(self._pefile.__data__) 82 | return self._pefile.__data__ 83 | 84 | @binary_data.setter 85 | def binary_data(self, data: bytearray) -> None: 86 | self._pefile.__data__ = data 87 | 88 | def _init_code_cave(self) -> bool: 89 | res = False 90 | for i, section in enumerate(self._pefile.sections): 91 | # 判断类似.text段的flags 92 | characteristics = section.Characteristics 93 | target_flags = [ 94 | "IMAGE_SCN_CNT_CODE", 95 | "IMAGE_SCN_CNT_INITIALIZED_DATA", 96 | "IMAGE_SCN_MEM_EXECUTE", 97 | "IMAGE_SCN_MEM_READ", 98 | ] 99 | target_characteristics_mask = 0 100 | for flag_name in target_flags: 101 | target_characteristics_mask |= pefile.SECTION_CHARACTERISTICS[flag_name] 102 | 103 | # flags 不满足条件，跳过下一个 104 | if ( 105 | characteristics & target_characteristics_mask 106 | != target_characteristics_mask 107 | ): 108 | continue 109 | 110 | file_size = section.SizeOfRawData 111 | virtual_size = section.Misc_VirtualSize 112 | if file_size > virtual_size: 113 | # 更新 virtual size, 否则后续 rva_to_offset 地址可能会转换出错 114 | section.Misc_VirtualSize = file_size 115 | section.Misc = section.Misc_VirtualSize 116 | section.Misc_PhysicalAddress = section.Misc_VirtualSize 117 | 118 | self._reload(self.fast_load) 119 | 120 | # 添加到 code cave 121 | self.add_range_to_code_cave( 122 | self.image_base + section.VirtualAddress + virtual_size, 123 | file_size - virtual_size, 124 | ) 125 | 126 | res |= True 127 | log.info( 128 | "add range 0x{:x} - 0x{:x} (0x{:x} bytes) to code cave".format( 129 | self.image_base + section.VirtualAddress + virtual_size, 130 | self.image_base + section.VirtualAddress + file_size, 131 | file_size - virtual_size, 132 | ) 133 | ) 134 | 135 | return res 136 | 137 | @property 138 | def _can_add_segment(self) -> bool: 139 | return True 140 | 141 | def _reload(self, fast_load=True): 142 | self._pefile = pefile.PE(data=self._pefile.write(), fast_load=fast_load) 143 | 144 | def _has_room_for_new_section_header(self) -> bool: 145 | max_header_addr = math.inf 146 | for i, section in enumerate(self._pefile.sections): 147 | if section.PointerToRawData == 0: 148 | continue 149 | max_header_addr = min(max_header_addr, section.PointerToRawData) 150 | 151 | section_header_size = self._pefile.sections[0].sizeof() 152 | now_header_end = ( 153 | self._pefile.sections[-1].get_file_offset() + section_header_size 154 | ) 155 | 156 | return (max_header_addr - now_header_end) // section_header_size >= 1 157 | 158 | def _fix_image_size(self): 159 | size_of_image = 0 160 | section_alignment = self._pefile.NT_HEADERS.OPTIONAL_HEADER.SectionAlignment 161 | file_alignment = self._pefile.NT_HEADERS.OPTIONAL_HEADER.FileAlignment 162 | for section in self._pefile.sections: 163 | # 计算区段的舍入后大小 164 | virtual_size = self.align(section.Misc_VirtualSize, section_alignment) 165 | 166 | # 更新映像的大小 167 | size_of_image = max(size_of_image, section.VirtualAddress + virtual_size) 168 | 169 | # 确保整个映像的大小是分页大小的倍数 170 | size_of_image = self.align(size_of_image, file_alignment) 171 | self._pefile.OPTIONAL_HEADER.SizeOfImage = size_of_image 172 | 173 | def _alloc_room_for_new_section_header(self) -> bool: 174 | if self._has_room_for_new_section_header(): 175 | return False 176 | 177 | max_header_addr = math.inf 178 | for i, section in enumerate(self._pefile.sections): 179 | if section.PointerToRawData == 0: 180 | continue 181 | max_header_addr = min(max_header_addr, section.PointerToRawData) 182 | 183 | section_header_size = self._pefile.sections[0].sizeof() 184 | alignment = self._pefile.NT_HEADERS.OPTIONAL_HEADER.FileAlignment 185 | new_max_header_addr = self.align( 186 | max_header_addr + section_header_size, alignment 187 | ) 188 | 189 | # mod from https://gist.github.com/williballenthin/d43cbc98fa127211c9099f46d2e73d2c 190 | # the offset at which the current section should begin 191 | dst_offset = new_max_header_addr 192 | # list of RealignedSection instances 193 | dst_secs = [] 194 | for section in sorted(self._pefile.sections, key=lambda s: s.PointerToRawData): 195 | dst_size = self.align(section.SizeOfRawData, alignment) 196 | padding = bytes(dst_size - section.SizeOfRawData) 197 | 198 | # collect pointers to the section data 199 | sec = RealignedSection( 200 | section.Name, 201 | section.PointerToRawData, 202 | section.SizeOfRawData, 203 | self.binary_data[ 204 | section.PointerToRawData : section.PointerToRawData 205 | + section.SizeOfRawData 206 | ], 207 | dst_offset, 208 | dst_size, 209 | padding, 210 | ) 211 | 212 | log.debug( 213 | "\tresizing {}\toffset: 0x{:x}\traw size: 0x{:x} \t--> offset: 0x{:x}\traw size: 0x{:x}".format( 214 | section.Name.decode().strip("\x00"), 215 | section.PointerToRawData, 216 | section.SizeOfRawData, 217 | dst_offset, 218 | dst_size, 219 | ) 220 | ) 221 | 222 | dst_secs.append(sec) 223 | 224 | # fixup the section pointers 225 | section.PointerToRawData = dst_offset 226 | section.SizeOfRawData = dst_size 227 | dst_offset += dst_size 228 | 229 | mod_buf = self._pefile.write() 230 | ret = [mod_buf[:max_header_addr], bytes(new_max_header_addr - max_header_addr)] 231 | for sec in dst_secs: 232 | ret.append(sec.src_buf) 233 | ret.append(sec.dst_padding) 234 | 235 | self._pefile = pefile.PE(data=b"".join(ret), fast_load=self.fast_load) 236 | self._fix_image_size() 237 | self._pefile.NT_HEADERS.OPTIONAL_HEADER.SizeOfHeaders = new_max_header_addr 238 | self._reload(self.fast_load) 239 | return True 240 | 241 | def _add_segment(self, addr: int = 0, length: int = 0x1000, prot: int = 7) -> int: 242 | self._alloc_room_for_new_section_header() 243 | if not self._has_room_for_new_section_header(): 244 | raise Exception("no room for new section header") 245 | 246 | def auto_last_segment() -> Tuple[int, int]: 247 | virtual_addr = 0 248 | file_offset = 0 249 | for s in self._pefile.sections: 250 | virtual_addr = self.align( 251 | max(virtual_addr, s.VirtualAddress + s.Misc_VirtualSize), 252 | self._pefile.NT_HEADERS.OPTIONAL_HEADER.SectionAlignment, 253 | ) 254 | file_offset = self.align( 255 | max(file_offset, s.PointerToRawData + s.SizeOfRawData), 256 | self._pefile.NT_HEADERS.OPTIONAL_HEADER.FileAlignment, 257 | ) 258 | 259 | return virtual_addr, file_offset 260 | 261 | def get_new_section_offset() -> int: 262 | s = max(self._pefile.sections, key=lambda ss: ss.get_file_offset()) 263 | return s.get_file_offset() + s.sizeof() 264 | 265 | last_rva, last_file_offset = auto_last_segment() 266 | 267 | target_flags = [ 268 | "IMAGE_SCN_CNT_CODE", 269 | "IMAGE_SCN_CNT_INITIALIZED_DATA", 270 | "IMAGE_SCN_MEM_EXECUTE", 271 | "IMAGE_SCN_MEM_READ", 272 | ] 273 | if prot & enums.SegProt.R.value: 274 | target_flags.append("IMAGE_SCN_MEM_READ") 275 | if prot & enums.SegProt.W.value: 276 | target_flags.append("IMAGE_SCN_MEM_WRITE") 277 | if prot & enums.SegProt.X.value: 278 | target_flags.append("IMAGE_SCN_MEM_EXECUTE") 279 | target_characteristics = 0 280 | for flag_name in target_flags: 281 | target_characteristics |= pefile.SECTION_CHARACTERISTICS[flag_name] 282 | 283 | new_section = pefile.SectionStructure( 284 | self._pefile.__IMAGE_SECTION_HEADER_format__ 285 | ) 286 | new_section.__unpack__(bytes(new_section.sizeof())) 287 | 288 | # 设置 section 属性 289 | new_section.set_file_offset(get_new_section_offset()) 290 | new_section.Name = b".patch" 291 | new_section.Misc_VirtualSize = length 292 | new_section.Misc = new_section.Misc_VirtualSize 293 | new_section.Misc_PhysicalAddress = new_section.Misc_VirtualSize 294 | new_section.VirtualAddress = last_rva 295 | new_section.SizeOfRawData = length 296 | new_section.PointerToRawData = last_file_offset 297 | new_section.PointerToRelocations = 0 298 | new_section.PointerToLinenumbers = 0 299 | new_section.NumberOfRelocations = 0 300 | new_section.NumberOfLinenumbers = 0 301 | new_section.Characteristics = target_characteristics 302 | 303 | overlay_offset = self._pefile.get_overlay_data_start_offset() 304 | if overlay_offset is not None: 305 | overlay_data = self.binary_data[overlay_offset:] 306 | else: 307 | overlay_data = None 308 | 309 | self._set_file_bytes( 310 | new_section.PointerToRawData, bytes(new_section.SizeOfRawData) 311 | ) 312 | 313 | # 添加新的 section 到 sections 数组 314 | self._pefile.sections.append(new_section) 315 | self._pefile.__structures__.append(new_section) 316 | 317 | # 更新 PE 头信息 318 | self._pefile.FILE_HEADER.NumberOfSections += 1 319 | self._fix_image_size() 320 | self._reload(self.fast_load) 321 | 322 | # 追加之前保存的 overlay 数据 323 | overlay_offset = self._pefile.get_overlay_data_start_offset() 324 | if overlay_data is not None: 325 | self._set_file_bytes(overlay_offset, overlay_data) 326 | 327 | return last_rva + self.image_base 328 | 329 | def _hook_byte_intel(self, vaddr: int, byte: Union[str, bytes, bytearray]) -> bool: 330 | if isinstance(byte, str): 331 | byte = byte.encode("latin-1") 332 | if isinstance(byte, bytes): 333 | byte = bytearray(byte) 334 | byte += b"\xc3" # append ret 335 | self.add_byte(byte, "user_hook_{}".format(hex(vaddr))) 336 | return self._hook_intel(vaddr) 337 | 338 | def _hook_asm_intel(self, vaddr: int, asm: str) -> bool: 339 | asm += ";ret;" # append ret 340 | addr = self._add_asm(asm) 341 | if not addr: 342 | return False 343 | self._label(addr, "user_hook_{}".format(hex(vaddr))) 344 | return self._hook_intel(vaddr) 345 | 346 | def _hook_intel(self, vaddr: int) -> bool: 347 | x86_jmp_inst_length = 5 348 | disasm_data = bytearray(self._get_byte(vaddr, 0x20)) 349 | cs_disasm = self.cs.disasm(disasm_data, vaddr) 350 | origin1_jmp_addr = vaddr 351 | tmp_len = 0 352 | for i, disasm in enumerate(cs_disasm): 353 | if i == 0 or tmp_len < x86_jmp_inst_length: 354 | tmp_len += len(disasm.bytes) 355 | origin1_jmp_length = tmp_len 356 | origin2_jmp_addr = vaddr + origin1_jmp_length 357 | backup_length = origin1_jmp_length + x86_jmp_inst_length 358 | backup1 = self._add_byte(bytearray(backup_length)) 359 | if not backup1: 360 | return False 361 | self._label(backup1, "backup1_{}".format(hex(vaddr))) 362 | backup2 = self._add_byte(bytearray(backup_length)) 363 | if not backup2: 364 | return False 365 | self._label(backup2, "backup2_{}".format(hex(vaddr))) 366 | if self.arch == enums.Arch.I386: 367 | detour1_asm = """ 368 | pushad 369 | call get_pc 370 | pc: mov esi,edi 371 | sub edi, pc - {} 372 | sub esi, pc - {{backup2_{}}} 373 | mov ecx, {} /* patch_byte_len */ 374 | cld 375 | rep movsb 376 | popad 377 | call {{user_hook_{}}} /* hook_func */ 378 | jmp {} /* origin1 */ 379 | get_pc: 380 | mov edi, [esp] 381 | ret 382 | """.format(vaddr, hex(vaddr), backup_length, hex(vaddr), origin1_jmp_addr) 383 | 384 | detour2_asm = """ 385 | pushf 386 | pushad 387 | call get_pc 388 | pc: mov esi,edi 389 | sub edi, pc - {} 390 | sub esi, pc - {{backup1_{}}} 391 | mov ecx, {} /* patch_byte_len */ 392 | cld 393 | rep movsb 394 | popad 395 | popf 396 | jmp {} /* origin2 */ 397 | get_pc: 398 | mov edi, [esp] 399 | ret 400 | """.format(vaddr, hex(vaddr), backup_length, origin2_jmp_addr) 401 | elif self.arch == enums.Arch.AMD64: 402 | detour1_asm = """ 403 | push rdi /* backup reg */ 404 | push rsi 405 | push rcx 406 | lea rdi, [{}] 407 | lea rsi, [{{backup2_{}}}] 408 | mov rcx, {} /* patch_byte_len */ 409 | cld 410 | rep movsb 411 | pop rcx /* restore reg */ 412 | pop rsi 413 | pop rdi 414 | call {{user_hook_{}}} /* hook_func */ 415 | jmp {} /* origin1 */ 416 | """.format(vaddr, hex(vaddr), backup_length, hex(vaddr), origin1_jmp_addr) 417 | 418 | detour2_asm = """ 419 | pushf 420 | push rdi /* backup reg */ 421 | push rsi 422 | push rcx 423 | lea rdi, [{}] 424 | lea rsi, [{{backup1_{}}}] 425 | mov rcx, {} /* patch_byte_len */ 426 | cld 427 | rep movsb 428 | pop rcx /* restore reg */ 429 | pop rsi 430 | pop rdi 431 | popf 432 | jmp {} /* origin2 */ 433 | """.format(vaddr, hex(vaddr), backup_length, origin2_jmp_addr) 434 | else: 435 | # should not be here 436 | return False 437 | 438 | detour1 = self._add_asm(detour1_asm) 439 | if not detour1: 440 | return False 441 | self._label(detour1, "detour1_{}".format(hex(vaddr))) 442 | detour2 = self._add_asm(detour2_asm) 443 | if not detour2: 444 | return False 445 | self._label(detour2, "detour2_{}".format(hex(vaddr))) 446 | origin1_jmp, inst_cnt = self._asm( 447 | "jmp {{detour1_{}}}".format(hex(vaddr)), origin1_jmp_addr 448 | ) 449 | origin2_jmp, inst_cnt = self._asm( 450 | "jmp {{detour2_{}}}".format(hex(vaddr)), origin2_jmp_addr 451 | ) 452 | backup1_bytes = bytearray(self._get_byte(vaddr, backup_length)) 453 | backup1_bytes[: len(origin1_jmp)] = origin1_jmp 454 | backup2_bytes = bytearray(self._get_byte(vaddr, backup_length)) 455 | backup2_bytes[origin1_jmp_length:] = origin2_jmp 456 | self._patch_byte(backup1, backup1_bytes) 457 | self._patch_byte(backup2, backup2_bytes) 458 | self._patch_byte(vaddr, backup1_bytes) 459 | 460 | section = self._pefile.get_section_by_rva(vaddr - self.image_base) 461 | if ( 462 | not section.Characteristics 463 | & pefile.SECTION_CHARACTERISTICS["IMAGE_SCN_MEM_WRITE"] 464 | ): 465 | log.success("Modify hook section prot +w") 466 | section.Characteristics |= pefile.SECTION_CHARACTERISTICS[ 467 | "IMAGE_SCN_MEM_WRITE" 468 | ] 469 | self._reload() 470 | return True 471 | 472 | @property 473 | def image_base(self) -> int: 474 | return self._pefile.NT_HEADERS.OPTIONAL_HEADER.ImageBase 475 | 476 | def rva_to_offset(self, rva: int) -> int: 477 | return self._pefile.get_offset_from_rva(rva - self.image_base) 478 | 479 | def add_segment(self, addr: int = 0, length: int = 0x1000, prot: int = 7) -> int: 480 | """ 481 | 新增一个段 482 | :param addr: 虚拟地址，默认为0表示自动计算 483 | :param length: 长度，默认0x1000 484 | :param prot: 保护权限，默认rwx 485 | :return: 成功则返回申请的地址，失败返回0 486 | """ 487 | if addr & 0xFFF: 488 | log.fail("address 0x{:x} is not page aligned".format(addr)) 489 | return 0 490 | new_segment = self._add_segment(addr, length, prot) 491 | if new_segment: 492 | log.success( 493 | "Add segment @ 0x{:x}, size: 0x{:x}".format(new_segment, length) 494 | ) 495 | return new_segment 496 | 497 | def save(self, filename: str = None) -> bool: 498 | # fix checksum 499 | self._reload() 500 | self._pefile.OPTIONAL_HEADER.CheckSum = self._pefile.generate_checksum() 501 | 502 | self._reload(self.fast_load) 503 | return super().save(filename) 504 | 505 | def hook_byte( 506 | self, vaddr: int, byte: Union[str, bytes, bytearray], label: str = None 507 | ) -> bool: 508 | """ 509 | 在指定地址插入bytes作为hook代码 510 | :param vaddr: 虚拟地址 511 | :param byte: bytes数组 512 | :param label: 此hook的名字，后续在asm中可以以{name}的方式引用这个hook的地址 513 | :return: 514 | """ 515 | if self.arch in [enums.Arch.I386, enums.Arch.AMD64]: 516 | if self._hook_byte_intel(vaddr, byte): 517 | self._label(vaddr, label) 518 | log.success("Hook @ {}".format(hex(vaddr))) 519 | return True 520 | else: 521 | log.fail("Hook @ {} failed".format(hex(vaddr))) 522 | return False 523 | else: 524 | raise exceptions.TODOException("Only support hook i386 & amd64 for now") 525 | 526 | def hook_asm(self, vaddr: int, asm: str, label: str = None) -> bool: 527 | """ 528 | 在指定地址插入asm作为hook代码 529 | :param vaddr: 虚拟地址 530 | :param asm: 汇编代码 531 | :param label: 此hook的名字，后续在asm中可以以{name}的方式引用这个hook的地址 532 | :return: 533 | """ 534 | if self.arch in [enums.Arch.I386, enums.Arch.AMD64]: 535 | if self._hook_asm_intel(vaddr, asm): 536 | self._label(vaddr, label) 537 | log.success("Hook @ {}".format(hex(vaddr))) 538 | return True 539 | else: 540 | log.fail("Hook @ {} failed".format(hex(vaddr))) 541 | return False 542 | else: 543 | raise exceptions.TODOException("Only support hook i386 & amd64 for now") 544 | --------------------------------------------------------------------------------