Spaces:

anky2002
/

graph-rag

Sleeping

graph-rag / src /graph_rag_service /core /security.py

GitHub Action

Automated sync to Hugging Face

bb63c6c about 2 months ago

3.06 kB

"""
Shared security helpers.
Central home for URL/SSRF validation so both the documents router (scrape /
crawl) and the system router (ollama model discovery) enforce the same rules.
"""

import ipaddress
import socket
from urllib.parse import urlparse

from fastapi import HTTPException


def _resolve_addresses(host: str) -> set:
    """Resolve a hostname to IPv4/IPv6 addresses (used for SSRF checks)."""
    try:
        infos = socket.getaddrinfo(host, None)
        return {info[4][0] for info in infos}
    except socket.gaierror:
        return set()


def _is_private_or_reserved(ip_str: str) -> bool:
    """True if the IP literal is loopback/private/link-local/reserved/metadata."""
    try:
        ip = ipaddress.ip_address(ip_str.split("%")[0])
    except ValueError:
        return True  # unparseable -> treat as unsafe
    return (
        ip.is_private
        or ip.is_loopback
        or ip.is_link_local
        or ip.is_reserved
        or ip.is_multicast
        or ip.is_unspecified
        # AWS / cloud metadata endpoint ranges
        or ip in ipaddress.ip_network("169.254.169.0/24")
    )


def validate_public_url(url: str) -> None:
    """Reject URLs that are not public http(s) — SSRF guard.
    Blocks non-http(s) schemes (file://, gopher://, etc.), IP-literal hosts,
    alternate IP notations (hex "0x7f000001", decimal "2130706433"), and any
    hostname that resolves to a private/loopback/link-local/metadata address.
    Raises HTTPException(400) on violation.
    """
    try:
        parsed = urlparse(url)
    except ValueError:
        raise HTTPException(status_code=400, detail="Invalid URL provided.")

    if parsed.scheme not in ("http", "https"):
        raise HTTPException(
            status_code=400,
            detail="Only http:// and https:// URLs are allowed.",
        )
    host = parsed.hostname
    if not host:
        raise HTTPException(status_code=400, detail="URL must include a host.")

    # Reject non-standard IP notations before any DNS resolution: dotted-quad,
    # IPv6 (bracketed or not), bare-decimal ("2130706433" == 127.0.0.1), and
    # hex ("0x7f000001" == 127.0.0.1). Browsers (Playwright/Chromium) DO
    # interpret these alternate forms as loopback, so a plain string check
    # alone would let them through. ipaddress.ip_address is strict, so anything
    # it cannot parse is treated as unsafe by _is_private_or_reserved.
    looks_like_ip = (
        ":" in host
        or host.replace(".", "").isdigit()
        or host.lower().startswith("0x")
    )
    if looks_like_ip and _is_private_or_reserved(host):
        raise HTTPException(
            status_code=400,
            detail="URL resolves to a private/loopback address and is not allowed.",
        )

    # Resolve DNS and reject private/reserved targets
    for addr in _resolve_addresses(host):
        if _is_private_or_reserved(addr):
            raise HTTPException(
                status_code=400,
                detail="URL resolves to a private/loopback address and is not allowed.",
            )

	"""
	Shared security helpers.

	Central home for URL/SSRF validation so both the documents router (scrape /
	crawl) and the system router (ollama model discovery) enforce the same rules.
	"""

	import ipaddress
	import socket
	from urllib.parse import urlparse

	from fastapi import HTTPException


	def _resolve_addresses(host: str) -> set:
	"""Resolve a hostname to IPv4/IPv6 addresses (used for SSRF checks)."""
	try:
	infos = socket.getaddrinfo(host, None)
	return {info[4][0] for info in infos}
	except socket.gaierror:
	return set()


	def _is_private_or_reserved(ip_str: str) -> bool:
	"""True if the IP literal is loopback/private/link-local/reserved/metadata."""
	try:
	ip = ipaddress.ip_address(ip_str.split("%")[0])
	except ValueError:
	return True # unparseable -> treat as unsafe
	return (
	ip.is_private
	or ip.is_loopback
	or ip.is_link_local
	or ip.is_reserved
	or ip.is_multicast
	or ip.is_unspecified
	# AWS / cloud metadata endpoint ranges
	or ip in ipaddress.ip_network("169.254.169.0/24")
	)


	def validate_public_url(url: str) -> None:
	"""Reject URLs that are not public http(s) — SSRF guard.

	Blocks non-http(s) schemes (file://, gopher://, etc.), IP-literal hosts,
	alternate IP notations (hex "0x7f000001", decimal "2130706433"), and any
	hostname that resolves to a private/loopback/link-local/metadata address.
	Raises HTTPException(400) on violation.
	"""
	try:
	parsed = urlparse(url)
	except ValueError:
	raise HTTPException(status_code=400, detail="Invalid URL provided.")

	if parsed.scheme not in ("http", "https"):
	raise HTTPException(
	status_code=400,
	detail="Only http:// and https:// URLs are allowed.",
	)
	host = parsed.hostname
	if not host:
	raise HTTPException(status_code=400, detail="URL must include a host.")

	# Reject non-standard IP notations before any DNS resolution: dotted-quad,
	# IPv6 (bracketed or not), bare-decimal ("2130706433" == 127.0.0.1), and
	# hex ("0x7f000001" == 127.0.0.1). Browsers (Playwright/Chromium) DO
	# interpret these alternate forms as loopback, so a plain string check
	# alone would let them through. ipaddress.ip_address is strict, so anything
	# it cannot parse is treated as unsafe by _is_private_or_reserved.
	looks_like_ip = (
	":" in host
	or host.replace(".", "").isdigit()
	or host.lower().startswith("0x")
	)
	if looks_like_ip and _is_private_or_reserved(host):
	raise HTTPException(
	status_code=400,
	detail="URL resolves to a private/loopback address and is not allowed.",
	)

	# Resolve DNS and reject private/reserved targets
	for addr in _resolve_addresses(host):
	if _is_private_or_reserved(addr):
	raise HTTPException(
	status_code=400,
	detail="URL resolves to a private/loopback address and is not allowed.",
	)