#!/usr/bin/python3 -u
# Scan a local web server: the "q" parameter reflected verbatim must be
# reported as XSS, and the same parameter HTML-escaped must come back clean.
import html
import json
import os
import subprocess
import tempfile
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlsplit


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def do_GET(self):
        url = urlsplit(self.path)
        value = parse_qs(url.query).get("q", [""])[0]
        if url.path == "/vulnerable":
            reflected = value
        elif url.path == "/safe":
            reflected = html.escape(value)
        else:
            self.send_error(404)
            return
        body = f"<html><body><p>Results for {reflected}</p></body></html>".encode()
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *args):
        pass


def scan(url):
    print(f"Scanning {url}")
    process = subprocess.run(["dalfox", "scan", url, "--format", "json", "--silence"],
                             stdout=subprocess.PIPE, text=True, timeout=300)
    report = json.loads(process.stdout)
    print(f"  exit code: {process.returncode}, requests: {report['meta']['total_requests']}, "
          f"findings: {len(report['findings'])}")
    for finding in report["findings"]:
        print(f"  [{finding['type']}] {finding['param']}: {finding['payload']}")
    return process.returncode, report


server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
threading.Thread(target=server.serve_forever, daemon=True).start()
url = f"http://127.0.0.1:{server.server_port}"
print(f"Test server listening on {url}")

with tempfile.TemporaryDirectory() as directory:
    os.environ["XDG_CONFIG_HOME"] = directory

    code, report = scan(f"{url}/vulnerable?q=test")
    assert code == 1, report
    assert any(finding["type"] == "V" and finding["param"] == "q"
               for finding in report["findings"]), report
    print("OK: reflected XSS reported in parameter q")

    code, report = scan(f"{url}/safe?q=test")
    assert code == 0, report
    assert not report["findings"], report
    assert report["meta"]["target_summary"][0]["status"] == "clean", report
    print("OK: escaped parameter reported as clean")
