"""
Tests for XML Entity Expansion (XXE) protection in nltk.downloader.
Verifies that defusedxml blocks Billion Laughs (XML bomb) payloads
when parsing the package index from a remote URL.
"""
import io
import unittest
from unittest.mock import patch
from defusedxml import EntitiesForbidden
from nltk.downloader import Downloader
# Billion Laughs payload: ~1KB input expands to ~3GB in memory
BILLION_LAUGHS_XML = """\
]>
&lol9;
"""
class TestDownloaderXXE(unittest.TestCase):
"""Verify that the downloader rejects XML bomb payloads."""
def test_billion_laughs_blocked(self):
"""defusedxml must reject entity expansion in index XML."""
downloader = Downloader()
xml_bytes = BILLION_LAUGHS_XML.encode("utf-8")
with patch("nltk.downloader.urlopen", return_value=io.BytesIO(xml_bytes)):
with self.assertRaises(EntitiesForbidden):
downloader._update_index(url="http://example.com/malicious.xml")
def test_valid_xml_parses(self):
"""Normal XML index should parse without error."""
valid_xml = """\