fix python dll finding mechanism for venvs

This commit is contained in:
cosine0 2022-10-27 09:59:28 +09:00
parent 153e92426f
commit af77d37912
1 changed files with 19 additions and 8 deletions

View File

@ -6,8 +6,8 @@ import sys
import tempfile import tempfile
from pathlib import Path from pathlib import Path
MAIN_DIR = os.path.dirname(os.path.abspath(inspect.getframeinfo(inspect.currentframe()).filename)) MAIN_DIR = Path(inspect.getframeinfo(inspect.currentframe()).filename).parent
INTERPRETER_DIR = os.path.dirname(os.path.abspath(sys.executable)) INTERPRETER_DIR = Path(sys.executable).parent
injected_script = ''' injected_script = '''
import time import time
@ -34,22 +34,33 @@ def main():
envs = os.environ.copy() envs = os.environ.copy()
if 'PATH' in os.environ: if 'PATH' in os.environ:
envs['PATH'] = os.path.pathsep.join([r'C:\windows\syswow64', INTERPRETER_DIR, os.environ['PATH']]) envs['PATH'] = os.path.pathsep.join([
r'C:\windows\syswow64',
INTERPRETER_DIR.absolute().as_posix(),
os.environ['PATH']])
else: else:
envs['PATH'] = INTERPRETER_DIR envs['PATH'] = INTERPRETER_DIR.absolute().as_posix()
if os.path.basename(INTERPRETER_DIR).lower() == 'scripts': if os.path.basename(INTERPRETER_DIR).lower() == 'scripts':
# in venv # in venv
library_path = Path(INTERPRETER_DIR) / 'lib' / 'site-packages' library_path = INTERPRETER_DIR / 'lib' / 'site-packages'
venv_setup = f'import sys\nsys.path.append(' \ venv_setup = f'import sys\nsys.path.append(' \
f'{library_path.absolute().as_posix()!r})' f'{library_path.absolute().as_posix()!r})'
with open(INTERPRETER_DIR.parent / 'pyvenv.cfg') as f:
for line in f:
if line.startswith('base-prefix'):
envs['PATH'] = os.path.pathsep.join([
line.split('=')[1].strip(),
envs['PATH']])
break
else: else:
venv_setup = '' venv_setup = ''
with tempfile.NamedTemporaryFile(suffix='.py', delete=False) as f: with tempfile.NamedTemporaryFile(suffix='.py', delete=False) as f:
formatted_script = injected_script.format(working_dir=os.path.abspath(os.curdir), formatted_script = injected_script.format(
script_path=script_path, working_dir=os.path.abspath(os.curdir),
venv_setup=venv_setup).encode('utf_8') script_path=script_path,
venv_setup=venv_setup).encode('utf_8')
f.write(formatted_script) f.write(formatted_script)
injector_path = Path(MAIN_DIR) / 'mayhem' / 'tools' / 'python_injector.py' injector_path = Path(MAIN_DIR) / 'mayhem' / 'tools' / 'python_injector.py'