From 28ee567a9a120967215ff19b96ab7515ce469620 Mon Sep 17 00:00:00 2001
From: Saifeddine ALOUI <aloui.seifeddine@gmail.com>
Date: Mon, 22 Jul 2024 00:57:22 +0200
Subject: [PATCH] Removed windows absolute paths

---
 lollms/security.py | 50 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 43 insertions(+), 7 deletions(-)

diff --git a/lollms/security.py b/lollms/security.py
index 9374fdb..8fe762d 100644
--- a/lollms/security.py
+++ b/lollms/security.py
@@ -183,10 +183,11 @@ def sanitize_path(path: str, allow_absolute_path: bool = False, error_text="Abso
 
     # Normalize path to use forward slashes
     path = path.replace('\\', '/')
-
-    if not allow_absolute_path and path.strip().startswith("/"):
+    path = path.strip()
+    if not allow_absolute_path and (path.startswith("/") or (len(path) == 2 and path[1] == ':')):
         raise HTTPException(status_code=400, detail=exception_text)
 
+
     # Regular expression to detect patterns like "....", multiple forward slashes, and command injection attempts like $(whoami)
     suspicious_patterns = re.compile(r'(\.\.+)|(/+/)|(\$\(.*\))')
 
@@ -288,8 +289,43 @@ def is_allowed_url(url):
 
 
 if __name__=="__main__":
-    sanitize_path_from_endpoint("main")
-    sanitize_path_from_endpoint("cat/main")
-    print("Main passed")
-    sanitize_path_from_endpoint(".../user")
-    print("hi")
+    test_cases = [
+        # Unix-style paths
+        ("valid/path/to/file.txt", False),
+        ("../../etc/passwd", False),
+        ("/absolute/path/file.txt", False),
+        ("relative/path/file.txt", False),
+        ("valid/path/with/..", False),
+        ("valid/path/with/./file.txt", False),
+        ("another/valid/path/file.txt", True),
+        ("/absolute/path/allowed.txt", True),
+        ("$(whoami)", False),
+        ("path/with/unauthorized&chars", False),
+        (None, False),
+
+        # Windows-style paths
+        (r"valid\path\to\file.txt", False),
+        (r"..\..\etc\passwd", False),
+        (r"C:\absolute\path\file.txt", False),
+        (r"relative\path\file.txt", False),
+        (r"valid\path\with\..", False),
+        (r"valid\path\with\.\file.txt", False),
+        (r"another\valid\path\file.txt", True),
+        (r"C:\absolute\path\allowed.txt", True),
+        (r"$(whoami)", False),
+        (r"path\with\unauthorized&chars", False),
+
+        # New test cases with C: drive
+        (r"C:\valid\path\to\file.txt", False),
+        (r"C:\another\valid\path\file.txt", True),
+        (r"C:\..\etc\passwd", False),
+        (r"C:\valid\path\with\..", False),
+        (r"C:", False),
+    ]
+
+    for path, allow_absolute in test_cases:
+        try:
+            sanitized = sanitize_path(path, allow_absolute)
+            print(f"Original: {path}, Sanitized: {sanitized}")
+        except HTTPException as e:
+            print(f"Original: {path}, Exception: {e.detail}")