From 131d76cb36d601b46c711faa3df02cb0c534e3cf Mon Sep 17 00:00:00 2001 From: retr0reg Date: Sun, 10 Mar 2024 20:28:16 +0800 Subject: [PATCH] Added is_allowed_url --- lollms/security.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) 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")