71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Patch pinned TransformerLab source to tolerate symlinked home directories."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
PATCH_MARKER = "with symlinked TransformerLab home directories."
|
|
TARGET_BLOCK = re.compile(
|
|
r"(?P<indent>[ \t]+)# The following prevents path traversal attacks:.*?"
|
|
r"(?P=indent)# now get the file contents",
|
|
re.DOTALL,
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--transformerlab-dir", required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
def repair_plugins_router(path: Path) -> bool:
|
|
source = path.read_text(encoding="utf-8")
|
|
if PATCH_MARKER in source:
|
|
return False
|
|
|
|
replacement = (
|
|
" # The following prevents path traversal attacks while remaining compatible\n"
|
|
" # with symlinked TransformerLab home directories.\n"
|
|
" plugin_dir = Path(await lab_dirs.plugin_dir_by_name((pluginId)))\n"
|
|
" resolved_plugin_dir = plugin_dir.resolve()\n"
|
|
' final_path = (plugin_dir / f"{filename}{file_ext}").resolve()\n'
|
|
"\n"
|
|
" try:\n"
|
|
" final_path.relative_to(resolved_plugin_dir)\n"
|
|
" except ValueError:\n"
|
|
' return {"message": f"File {filename}{file_ext} is outside plugin directory"}\n'
|
|
"\n"
|
|
" # now get the file contents"
|
|
)
|
|
updated, count = TARGET_BLOCK.subn(replacement, source, count=1)
|
|
if count != 1:
|
|
raise RuntimeError(f"Could not find path traversal block in {path}")
|
|
|
|
path.write_text(updated, encoding="utf-8")
|
|
return True
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
root = Path(args.transformerlab_dir).expanduser().resolve()
|
|
plugins_router = root / "src" / "transformerlab" / "routers" / "experiment" / "plugins.py"
|
|
if not plugins_router.exists():
|
|
print(f"missing TransformerLab plugins router: {plugins_router}", file=sys.stderr)
|
|
return 1
|
|
|
|
changed = repair_plugins_router(plugins_router)
|
|
if changed:
|
|
print(f"patched {plugins_router}")
|
|
else:
|
|
print(f"already patched {plugins_router}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|