From 82144e27ef6c3a9785e77edc522a1a3fa752a11a Mon Sep 17 00:00:00 2001
From: mhdzumair <mhdzumair@gmail.com>
Date: Sun, 19 Jan 2025 08:01:55 +0530
Subject: [PATCH] Refactor private IP address checks with ipaddress module

Replaced regex-based private IP detection with a utility function `is_private_ip` using the `ipaddress` module.
---
 utils/network.py           | 19 +++++++++++++------
 utils/runtime_const.py     |  4 ----
 utils/validation_helper.py |  4 ++--
 3 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/utils/network.py b/utils/network.py
index 0d7a43dc..970eb1d9 100644
--- a/utils/network.py
+++ b/utils/network.py
@@ -1,5 +1,6 @@
 import asyncio
 import logging
+from ipaddress import ip_address
 from typing import Callable, AsyncGenerator, Any, Tuple, Dict
 from urllib import parse
 from urllib.parse import urlencode, urlparse
@@ -8,11 +9,9 @@
 from fastapi.requests import Request
 
 from db.config import settings
+from db.redis_database import REDIS_ASYNC_CLIENT
 from db.schemas import UserData
 from utils import crypto
-from utils.crypto import encrypt_data
-from utils.runtime_const import PRIVATE_CIDR
-from db.redis_database import REDIS_ASYNC_CLIENT
 
 
 class CircuitBreakerOpenException(Exception):
@@ -250,7 +249,7 @@ async def get_mediaflow_proxy_public_ip(mediaflow_config) -> str | None:
         return mediaflow_config.public_ip
 
     parsed_url = urlparse(mediaflow_config.proxy_url)
-    if PRIVATE_CIDR.match(parsed_url.netloc):
+    if is_private_ip(parsed_url.netloc):
         # MediaFlow proxy URL is a private IP address
         return None
 
@@ -299,7 +298,7 @@ async def get_user_public_ip(
     # Get the user's public IP address
     user_ip = get_client_ip(request)
     # check if the user's IP address is a private IP address
-    if PRIVATE_CIDR.match(user_ip):
+    if is_private_ip(user_ip):
         # Use host public IP address.
         return None
     return user_ip
@@ -355,7 +354,7 @@ def encode_mediaflow_proxy_url(
     if encryption_api_password:
         if "api_password" not in query_params:
             query_params["api_password"] = encryption_api_password
-        encrypted_token = encrypt_data(
+        encrypted_token = crypto.encrypt_data(
             encryption_api_password, query_params, expiration, ip
         )
         encoded_params = urlencode({"token": encrypted_token})
@@ -365,3 +364,11 @@ def encode_mediaflow_proxy_url(
     # Construct the full URL
     base_url = parse.urljoin(mediaflow_proxy_url, endpoint)
     return f"{base_url}?{encoded_params}"
+
+
+def is_private_ip(ip_str: str) -> bool:
+    try:
+        ip = ip_address(ip_str)
+        return ip.is_private
+    except ValueError:
+        return False
diff --git a/utils/runtime_const.py b/utils/runtime_const.py
index f6d3c5d5..abdaedbd 100644
--- a/utils/runtime_const.py
+++ b/utils/runtime_const.py
@@ -23,10 +23,6 @@
 
 SPORTS_ARTIFACTS = get_json_data("resources/json/sports_artifacts.json")
 
-PRIVATE_CIDR = re.compile(
-    r"^(10\.|127\.|172\.(1[6-9]|2[0-9]|3[01])\.|192\.168\.)",
-)
-
 TEMPLATES = Jinja2Templates(directory="resources")
 MANIFEST_TEMPLATE = TEMPLATES.get_template("templates/manifest.json.j2")
 
diff --git a/utils/validation_helper.py b/utils/validation_helper.py
index ac3dc5e3..0f6951aa 100644
--- a/utils/validation_helper.py
+++ b/utils/validation_helper.py
@@ -8,7 +8,7 @@
 from db import schemas
 from db.config import settings
 from utils import const
-from utils.runtime_const import PRIVATE_CIDR
+from utils.network import is_private_ip
 from db.redis_database import REDIS_ASYNC_CLIENT
 
 
@@ -329,7 +329,7 @@ async def validate_mediaflow_proxy_credentials(user_data: schemas.UserData) -> d
 
     if results["message"].startswith("RequestError"):
         parsed_url = urlparse(user_data.mediaflow_config.proxy_url)
-        if PRIVATE_CIDR.match(parsed_url.netloc):
+        if is_private_ip(parsed_url.netloc):
             # MediaFlow proxy URL is a private IP address
             return {
                 "status": "success",