Skip to content

Commit

Permalink
Refactor private IP address checks with ipaddress module
Browse files Browse the repository at this point in the history
Replaced regex-based private IP detection with a utility function `is_private_ip` using the `ipaddress` module.
  • Loading branch information
mhdzumair committed Jan 19, 2025
1 parent 4e6b814 commit 82144e2
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 12 deletions.
19 changes: 13 additions & 6 deletions utils/network.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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})
Expand All @@ -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
4 changes: 0 additions & 4 deletions utils/runtime_const.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
4 changes: 2 additions & 2 deletions utils/validation_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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",
Expand Down

0 comments on commit 82144e2

Please sign in to comment.