lollms-webui/tests/pentests/test_sanitize/test.py

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!"):
    if path.strip().startswith("/"):
        raise HTTPException(status_code=400, detail=exception_text)
    # 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 = "/D:/POC/secret.txt"
    
    #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()