from typing import Set, Union, Dict, Optional
import json
from dataclasses import dataclass
from mitmproxy import ctx
from mitmproxy.addonmanager import Loader
from mitmproxy.http import HTTPFlow


"""
This extension implements support for domain fronting.

Usage:

   mitmproxy -s examples/contrib/domain_fronting.py --set domainfrontingfile=./domain_fronting.json

In the following basic example, www.example.com will be used for DNS requests and SNI values
but the secret.example.com value will be used for the HTTP host header:

    {
        "mappings": [
            {
                "patterns": ["secret.example.com"],
                "server": "www.example.com"
            }
        ]
    }

The following example demonstrates the usage of a wildcard (at the beginning of the domain name only):

    {
        "mappings": [
            {
                "patterns": ["*.foo.example.com"],
                "server": "www.example.com"
            }
        ]
    }

In the following example, we override the HTTP host header:

    {
        "mappings": [
            {
                "patterns": ["foo.example"],
                "server": "www.example.com",
                "host": "foo.proxy.example.com"
            }
        ]
    }

"""


@dataclass
class Mapping:
    server: Union[str, None]
    host: Union[str, None]


class HttpsDomainFronting:

    # configurations for regular ("foo.example.com") mappings:
    star_mappings: Dict[str, Mapping]

    # Configurations for star ("*.example.com") mappings:
    strict_mappings: Dict[str, Mapping]

    def __init__(self) -> None:
        self.strict_mappings = {}
        self.star_mappings = {}

    def _resolve_addresses(self, host: str) -> Optional[Mapping]:
        mapping = self.strict_mappings.get(host)
        if mapping is not None:
            return mapping

        index = 0
        while True:
            index = host.find(".", index)
            if index == -1:
                break
            super_domain = host[(index + 1):]
            mapping = self.star_mappings.get(super_domain)
            if mapping is not None:
                return mapping
            index += 1

        return None

    def load(self, loader: Loader) -> None:
        loader.add_option(
            name="domainfrontingfile",
            typespec=str,
            default="./fronting.json",
            help="Domain fronting configuration file",
        )

    def _load_configuration_file(self, filename: str) -> None:
        config = json.load(open(filename, "rt"))
        strict_mappings: Dict[str, Mapping] = {}
        star_mappings: Dict[str, Mapping] = {}
        for mapping in config["mappings"]:
            item = Mapping(server=mapping.get("server"), host=mapping.get("host"))
            for pattern in mapping["patterns"]:
                if pattern.startswith("*."):
                    star_mappings[pattern[2:]] = item
                else:
                    strict_mappings[pattern] = item
        self.strict_mappings = strict_mappings
        self.star_mappings = star_mappings

    def configure(self, updated: Set[str]) -> None:
        if "domainfrontingfile" in updated:
            domain_fronting_file = ctx.options.domainfrontingfile
            self._load_configuration_file(domain_fronting_file)

    def request(self, flow: HTTPFlow) -> None:
        if not flow.request.scheme == "https":
            return
        # We use the host header to dispatch the request:
        target = flow.request.host_header
        if target is None:
            return
        mapping = self._resolve_addresses(target)
        if mapping is not None:
            flow.request.host = mapping.server or target
            flow.request.headers["host"] = mapping.host or target


addons = [HttpsDomainFronting()]