diff --git a/nethook_helpers.py b/nethook_helpers.py new file mode 100644 index 00000000..a97c5a85 --- /dev/null +++ b/nethook_helpers.py @@ -0,0 +1,113 @@ +import sys +import pathlib +import shutil +import os +import typing as t +import psutil +import time +from collections import Counter + + +STEAM_PATH = R'C:\Program Files (x86)\Steam\steam.exe' +GALAXY_PATH = R'C:\Program Files (x86)\GOG Galaxy\GalaxyClient.exe' +NETHOOK_PATH = None # R'C:COMPLETE_ME\\NetHook2\NetHook2.dll' + + +def find_steam_ps(initial: t.Optional[psutil.Process]=None) -> t.Optional[psutil.Process]: + if initial and initial.is_running(): + return initial + for p in psutil.process_iter(['exe'], ad_value=''): + if STEAM_PATH == p.info['exe']: + return p + return None + + +def restart_steam(): + p = psutil.Popen([STEAM_PATH, "-shutdown", "-silent"]) + p.wait() + steam_ps = find_steam_ps() + if steam_ps: + steam_ps.wait(timeout=10) + return psutil.Popen(STEAM_PATH) + + +def restart_galaxy(): + p = psutil.Popen([GALAXY_PATH, '/command=shutdown']) + p.wait() + psutil.Popen(GALAXY_PATH) + + +def inject_nethook(steam_ps=None): + """Requires running as admin""" + timeout = time.time() + 10 + while True: + steam_ps = find_steam_ps(steam_ps) + assert steam_ps, 'NetHook requires Steam running' + if time.time() > timeout: + print('timeout on loading steamclint.dll by Steam') + break + if [dll for dll in steam_ps.memory_maps() if 'steamclient.dll' in dll.path]: + print('steamclient.dll loaded. Injecting nethook') + break + cmd = f'rundll32 "{NETHOOK_PATH}",Inject' + psutil.Popen(cmd) + + +def copy_results_to_common_dir(dest_dir, must_include=''): + os.makedirs(dest_dir, exist_ok=True) + steam_nethook_dir = pathlib.PurePath(STEAM_PATH).parent / 'nethook' + for d in os.listdir(steam_nethook_dir): + if steam_nethook_dir.name == d: # avoid recurssion + continue + for f in os.listdir(steam_nethook_dir / d): + if must_include not in f: + continue + src = steam_nethook_dir / d / f + dest = pathlib.PurePath(dest_dir) / (src.name + '_' + d + src.suffix) # originalname_timestamp.bin + shutil.copy(src, dest) + + +def print_statistics(first_signals=10): + steam_nethook_dir = pathlib.PurePath(STEAM_PATH).parent / 'nethook' + signals = Counter() + for root, dirs, files in os.walk(steam_nethook_dir): + signal_names = [] + for i, f in enumerate(sorted(files)): + if i > first_signals: + break + signal_names.append(f.split('_')[-1]) + signals.update(signal_names) + print(signals) + + +def main(): + if NETHOOK_PATH is None: + print('set up nethook path in script first') + + try: + no = int(sys.argv[1]) + except (IndexError, ValueError): + print('Usage: