diff --git a/relocatable.py b/relocatable.py index dd8b17d..473194e 100644 --- a/relocatable.py +++ b/relocatable.py @@ -65,15 +65,37 @@ def patch_path(output_dir: pathlib.Path, patch: str) -> None: for root, _, filenames in output_dir.walk(): for filename in filenames: path = root / filename - if path.is_symlink(): + if path.is_symlink() or path.suffix in (".pyc", ".pyo"): continue contents = path.read_bytes() if placeholder not in contents: continue - replacement = target if b"\0" in contents: - replacement = target.ljust(len(placeholder), b"\0") - path.write_bytes(contents.replace(placeholder, replacement)) + # Pad after the whole string to preserve path suffixes. + patched = re.sub( + re.escape(placeholder) + rb"[^\0]*", + lambda match: match[0].replace(placeholder, target).ljust( + len(match[0]), b"\0" + ), + contents, + ) + else: + patched = contents.replace(placeholder, target) + # shortcut: generated CPython launchers only, extend for other scripts. + if match := re.match( + b"#!" + re.escape(placeholder) + rb"(/bin/python[0-9.]*[td]*)\n", + contents, + ): + # CPython's launchers need both shell and Python quoting. + executable = target + match[1] + quoted = executable.replace(b"'", b"'\"'\"'").replace( + b"\\", b"'\"\\\\\"'" + ) + patched = ( + b"#!/bin/sh\n'exec' '" + quoted + b"' \"$0\" \"$@\"\n" + + patched.partition(b"\n")[2] + ) + path.write_bytes(patched) if contents[:4] in MACHO_MAGIC: # Rewriting a Mach-O file invalidates its (ad-hoc) code # signature, and arm64 macOS kills processes that load it. diff --git a/test_relocatable.py b/test_relocatable.py new file mode 100644 index 0000000..f13fc85 --- /dev/null +++ b/test_relocatable.py @@ -0,0 +1,95 @@ +import json +import os +import pathlib +import py_compile +import subprocess +import sys +import tempfile +import unittest + +from relocatable import patch_path, placeholder_prefix + + +class PatchPathTests(unittest.TestCase): + def test_python_shebang(self): + with tempfile.TemporaryDirectory() as directory: + root = pathlib.Path(directory).resolve() + names = ( + "with spaces", + "with\ttab", + "with ''' quotes", + 'with """ quotes', + r"with \u backslash", + "with $HOME and $(false)", + "with café", + "x" * (250 - len(os.fsencode(root)) - 1), + ) + for name in names: + with self.subTest(name=name): + prefix = root / name + bin_dir = prefix / "bin" + bin_dir.mkdir(parents=True) + interpreter = bin_dir / "python3.16" + interpreter.symlink_to(sys.executable) + script = bin_dir / "pydoc3.16" + script.write_text( + f"#!{placeholder_prefix()}/bin/python3.16\n" + "import json, sys\n" + "print(json.dumps([sys.executable, sys.argv[1:]]))\n", + encoding="utf-8", + ) + script.chmod(0o755) + patch_path(prefix, str(prefix)) + result = subprocess.run( + [script, "argument with spaces", "$literal"], + env={**os.environ, "PATH": ""}, + capture_output=True, text=True, check=True, + ) + executable, arguments = json.loads(result.stdout) + self.assertTrue(pathlib.Path(executable).samefile(interpreter)) + self.assertEqual(arguments, ["argument with spaces", "$literal"]) + + def test_prefix_replacement(self): + prefix = placeholder_prefix().encode() + target = b"/opt/python" + padding = b"\0" * (len(prefix) - len(target)) + cases = ( + (prefix + b"/lib", target + b"/lib"), + ( + b"header\0" + prefix + b"/lib\0tail", + b"header\0" + target + b"/lib" + padding + b"\0tail", + ), + ( + b"\0" + prefix + b"/lib:" + prefix + b"/lib64\0", + b"\0" + target + b"/lib:" + target + b"/lib64" + padding * 2 + b"\0", + ), + (b"\0" + prefix + b"/lib", b"\0" + target + b"/lib" + padding), + ) + with tempfile.TemporaryDirectory() as directory: + root = pathlib.Path(directory) + path = root / "file" + for original, expected in cases: + with self.subTest(original=original): + path.write_bytes(original) + patch_path(root, target.decode()) + self.assertEqual(path.read_bytes(), expected) + + def test_bytecode_is_unchanged(self): + with tempfile.TemporaryDirectory() as directory: + root = pathlib.Path(directory) + source = root / "module.py" + source.write_text("pass\n", encoding="utf-8") + for suffix in (".pyc", ".pyo"): + with self.subTest(suffix=suffix): + bytecode = root / ("module" + suffix) + py_compile.compile( + str(source), cfile=str(bytecode), + dfile=placeholder_prefix() + "/module.py", doraise=True, + ) + original = bytecode.read_bytes() + patch_path(root, "/opt/python") + self.assertEqual(bytecode.read_bytes(), original) + + +if __name__ == "__main__": + unittest.main()