mirror of
https://github.com/ParisNeo/lollms-webui.git
synced 2024-12-29 16:48:50 +00:00
36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
|
from ascii_colors import ASCIIColors
|
||
|
|
from fastapi import HTTPException
|
||
|
|
from pathlib import Path
|
||
|
|
import re
|
||
|
|
import pytest
|
||
|
|
def sanitize_path_from_endpoint(path: str, error_text="A suspected LFI attack detected. The path sent to the server has suspicious elements in it!", exception_text="Invalid path!"):
|
||
|
|
# Fix the case of "/" at the beginning on the path
|
||
|
|
if path is None:
|
||
|
|
return path
|
||
|
|
|
||
|
|
# Regular expression to detect patterns like "...." and multiple forward slashes
|
||
|
|
suspicious_patterns = re.compile(r'(\.\.+)|(/+/)')
|
||
|
|
|
||
|
|
if suspicious_patterns.search(path) or Path(path).is_absolute():
|
||
|
|
ASCIIColors.error(error_text)
|
||
|
|
raise HTTPException(status_code=400, detail=exception_text)
|
||
|
|
|
||
|
|
path = path.lstrip('/')
|
||
|
|
return path
|
||
|
|
|
||
|
|
|
||
|
|
def test_sanitize_path_from_endpoint():
|
||
|
|
# Test a valid path
|
||
|
|
valid_path = "example/path"
|
||
|
|
assert sanitize_path_from_endpoint(valid_path) == "example/path"
|
||
|
|
|
||
|
|
# Test a path with suspicious elements
|
||
|
|
suspicious_path = "/images//D:/POC/secret.txt"
|
||
|
|
with pytest.raises(HTTPException):
|
||
|
|
sanitize_path_from_endpoint(suspicious_path)
|
||
|
|
|
||
|
|
# Add more test cases as needed
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
test_sanitize_path_from_endpoint()
|