[RTY-260032]: add modular approach and whitelist-blacklist authorization

Author: Ravindrayadav04

app/routes.py

@@ -34,8 +34,19 @@
     remove_cache_key,
     rev_cache,
 )
-from app.utils.config import DOMAIN, MAX_RECENT_URLS, CACHE_PURGE_TOKEN, QR_DIR
-from app.utils.helper import generate_code, is_valid_url, format_date
+from app.utils.config import (
+    DOMAIN,
+    MAX_RECENT_URLS,
+    CACHE_PURGE_TOKEN,
+    QR_DIR,
+)
+from app.utils.helper import (
+    generate_code,
+    sanitize_url,
+    is_valid_url,
+    authorize_url,
+    format_date,
+)
 from app.utils.qr import generate_qr_with_logo
 
 # templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))
@@ -98,11 +109,18 @@ async def create_short_url(
     qr_type: str = Form("short"),
 ):
     session = request.session
+    original_url = sanitize_url(original_url)  # sanitize the URL input
 
-    if not original_url or not is_valid_url(original_url):
+    if not original_url or not is_valid_url(original_url):  # validate the URL
         session["error"] = "Please enter a valid URL."
         return RedirectResponse("/", status_code=status.HTTP_303_SEE_OTHER)
 
+    if not authorize_url(
+        original_url
+    ):  # authorize the URL based on whitelist/blacklist
+        session["error"] = "This domain is not allowed."
+        return RedirectResponse("/", status_code=status.HTTP_303_SEE_OTHER)
+
     short_code: Optional[str] = get_short_from_cache(original_url)
 
     if not short_code and db.is_connected():
@@ -329,10 +347,14 @@ class ShortenRequest(BaseModel):
 
 @api_v1.post("/shorten")
 def shorten_api(payload: ShortenRequest):
-    original_url = payload.url
+    original_url = sanitize_url(payload.url)
+
     if not is_valid_url(original_url):
         return JSONResponse(status_code=400, content={"error": "INVALID_URL"})
 
+    if not authorize_url(original_url):
+        return JSONResponse(status_code=400, content={"error": "DOMAIN_NOT_ALLOWED"})
+
     short_code = get_short_from_cache(original_url)
     if not short_code:
         short_code = generate_code()

app/utils/config.py

@@ -82,8 +82,16 @@ def _get_int(key: str, default: int) -> int:
 MAX_URL_LENGTH = 2048
 
 # for making the qr constant
-# Base project paths
 BASE_DIR = Path(__file__).resolve().parent.parent  # app/
 PROJECT_ROOT = BASE_DIR.parent  # project root
-# QR directory constant
-QR_DIR = PROJECT_ROOT / "assets" / "images" / "qr"
+QR_DIR = PROJECT_ROOT / "assets" / "images" / "qr"  # QR directory constant
+
+
+# for the check of the url which is blacklist or whitelist
+whitelist_urls: set[str] = set()
+blacklist_urls: set[str] = {
+    "malware.com",
+    "phishing.site",
+    "badsite.test",
+    "spam.test",
+}

app/utils/helper.py

@@ -3,13 +3,18 @@
 from datetime import datetime, timezone
 from zoneinfo import ZoneInfo
 from typing import Union
-from app.utils.config import SHORT_CODE_LENGTH
+from app.utils.config import SHORT_CODE_LENGTH, whitelist_urls, blacklist_urls
 from urllib.parse import urlparse
 import ipaddress
 
 
+# for sanitization the url
+def sanitize_url(url: str) -> str:
+    return url.strip()
+
+
+# for validating the url
 def is_valid_url(url: str) -> bool:
-    url = url.strip()  # sanitize here
 
     try:
         parsed = urlparse(url)
@@ -43,6 +48,29 @@ def is_valid_url(url: str) -> bool:
         return False
 
 
+# for authorizing the url based on whitelist and blacklist
+def authorize_url(url: str) -> bool:
+    """Check whitelist / blacklist rules."""
+    hostname = urlparse(url).hostname
+
+    if hostname is None:
+        return False
+
+    hostname = hostname.lower()
+
+    # block blacklist domains
+    if any(hostname.endswith(domain) for domain in blacklist_urls):
+        return False
+
+    # allow only whitelist domains (if defined)
+    if whitelist_urls and not any(
+        hostname.endswith(domain) for domain in whitelist_urls
+    ):
+        return False
+
+    return True
+
+
 def generate_code(length: int = SHORT_CODE_LENGTH) -> str:
     chars = string.ascii_letters + string.digits
     return "".join(random.choice(chars) for _ in range(length))