import re from pathlib import Path import pytest from ascii_colors import ASCIIColors from fastapi import HTTPException 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()