extractSystemCalls.py (2938B)
1 # This Source Code Form is subject to the terms of the Mozilla Public 2 # License, v. 2.0. If a copy of the MPL was not distributed with this 3 # file, You can obtain one at http://mozilla.org/MPL/2.0/. 4 5 import argparse 6 import os 7 import re 8 from pathlib import Path 9 10 11 def extract_dll_methods(file: str, list_dll, list_rejects, windows_methods): 12 accepted = {} 13 rejected = {} 14 with open(file, "rb") as f: 15 data = f.read() 16 pattern = re.compile(rb"[a-zA-Z]\w{2,}") 17 for name in pattern.findall(data): 18 str_name = name.decode("ascii") 19 if str_name in windows_methods: 20 accepted[str_name] = windows_methods[str_name] 21 else: 22 rejected[str_name] = 1 23 if list_rejects: 24 for name in sorted(rejected.keys()): 25 print(f"{name}") 26 elif list_dll: 27 for name in sorted(accepted.keys()): 28 print(f"{name} ({accepted[name]})") 29 else: 30 for name in sorted(accepted.keys()): 31 print(f"{name}") 32 33 34 def extract_crate_methods(crate: str): 35 methods = {} 36 pattern = re.compile('link!\\("([\\w\\.-]+)".+?\\s+fn\\s+(\\w+)\\(') 37 for dirpath, dirnames, filenames in os.walk(crate): 38 for filename in filenames: 39 with open(os.path.join(dirpath, filename)) as f: 40 data = f.read() 41 for method in pattern.findall(data): 42 methods[method[1]] = method[0] 43 return methods 44 45 46 def main(): 47 examples = """examples: 48 python dom/media/tools/extractSystemCalls.py /path/to/widevinecdm.dll""" 49 50 parser = argparse.ArgumentParser( 51 description="Extract Windows system call usage from Windows DLL", 52 epilog=examples, 53 formatter_class=argparse.RawDescriptionHelpFormatter, 54 ) 55 parser.add_argument( 56 "plugin", 57 help="path to plugin", 58 ) 59 parser.add_argument( 60 "--list-rejects", action="store_true", help="output unknown strings instead" 61 ) 62 parser.add_argument( 63 "--list-dll", action="store_true", help="include the system dll name" 64 ) 65 parser.add_argument( 66 "--windows-sys-crate", help="alternate path to the rust windows-sys crate" 67 ) 68 args = parser.parse_args() 69 70 # Pull out a list of win32 APIs from one of our Rust crates that we can 71 # cross reference against to filter out the unrelated strings. 72 if args.windows_sys_crate is not None: 73 crate_path = args.windows_sys_crate 74 else: 75 crate_path = ( 76 Path(__file__).parents[3].joinpath("third_party", "rust", "windows-sys") 77 ) 78 windows_methods = extract_crate_methods(crate_path) 79 80 # RtlGenRandom maps to SystemFunction036 81 # https://learn.microsoft.com/en-us/windows/win32/api/ntsecapi/nf-ntsecapi-rtlgenrandom 82 windows_methods["SystemFunction036"] = "advapi32.dll" 83 84 extract_dll_methods(args.plugin, args.list_dll, args.list_rejects, windows_methods) 85 86 87 if __name__ == "__main__": 88 main()