diff --git a/lollms/security.py b/lollms/security.py index 39cd5ef..7bd4352 100644 --- a/lollms/security.py +++ b/lollms/security.py @@ -1,10 +1,13 @@ from fastapi import HTTPException from ascii_colors import ASCIIColors +from urllib.parse import urlparse +import socket from pathlib import Path from typing import List import os import re + def sanitize_path(path:str, allow_absolute_path:bool=False, error_text="Absolute database path detected", exception_text="Detected an attempt of path traversal. Are you kidding me?"): if path is None: return path @@ -58,6 +61,22 @@ def validate_path(path, allowed_paths:List[str|Path]): # If the path is not within any of the allowed paths, return False return False +def is_allowed_url(url): + # Check if url is legit + parsed_url = urlparse(url) + # Check if scheme is not http or https, return False + if parsed_url.scheme not in ['http', 'https']: + return False + + hostname = parsed_url.hostname + + try: + ip_address = socket.gethostbyname(hostname) + except socket.gaierror: + return False + + return not ip_address.startswith('127.') or ip_address.startswith('192.168.') or ip_address.startswith('10.') or ip_address.startswith('172.') + if __name__=="__main__": sanitize_path_from_endpoint("main")