# This file was copied from mitmproxy/addons/serverplayback.py release tag 2.0.2 and modified by
# Benjamin Smedberg

# Altered features:
# * --kill returns 404 rather than dropping the whole HTTP/2 connection on the floor
# * best-match response handling is used to improve success rates

import hashlib
import urllib
import sys
from collections import defaultdict
from typing import Any  # noqa
from typing import List  # noqa

from mitmproxy import ctx
from mitmproxy import exceptions
from mitmproxy import io
from mitmproxy import http


class ServerPlayback:
    def __init__(self, replayfiles):
        self.options = None
        self.replayfiles = replayfiles
        self.flowmap = {}

    def load(self, flows):
        for i in flows:
            if i.response:
                l = self.flowmap.setdefault(self._hash(i.request), [])
                l.append(i)

    def clear(self):
        self.flowmap = {}

    def _parse(self, r):
        """
            Return (path, queries, formdata, content) for a request.
        """
        _, _, path, _, query, _ = urllib.parse.urlparse(r.url)
        queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True)
        queries = defaultdict(list)
        for k, v in queriesArray:
            queries[k].append(v)

        content = None
        formdata = None
        if r.raw_content != b'':
            if r.multipart_form:
                formdata = r.multipart_form
            elif r.urlencoded_form:
                formdata = r.urlencoded_form
            else:
                content = r.content
        return (path, queries, formdata, content)

    def _hash(self, r):
        """
            Calculates a loose hash of the flow request.
        """
        path, queries, _, _ = self._parse(r)

        key = [str(r.port), str(r.scheme), str(r.method), str(path)]  # type: List[Any]
        if not self.options.server_replay_ignore_host:
            key.append(r.host)

        if len(queries):
            key.append("?")

        return hashlib.sha256(
            repr(key).encode("utf8", "surrogateescape")
        ).digest()

    def _match(self, request_a, request_b):
        """
            Calculate a match score between two requests.
            Match algorithm:
              * identical query keys: 3 points
              * matching query param present: 1 point
              * matching query param value: 3 points
              * identical form keys: 3 points
              * matching form param present: 1 point
              * matching form param value: 3 points
              * matching body (no multipart or encoded form): 4 points
        """
        match = 0

        path_a, queries_a, form_a, content_a = self._parse(request_a)
        path_b, queries_b, form_b, content_b = self._parse(request_b)

        keys_a = set(queries_a.keys())
        keys_b = set(queries_b.keys())
        if keys_a == keys_b:
            match += 3

        for key in keys_a:
            values_a = set(queries_a[key])
            values_b = set(queries_b[key])
            if len(values_a) == len(values_b):
                match += 1
            if values_a == values_b:
                match += 3

        if form_a and form_b:
            keys_a = set(form_a.keys())
            keys_b = set(form_b.keys())
            if keys_a == keys_b:
                match += 3

            for key in keys_a:
                values_a = set(form_a.get_all(key))
                values_b = set(form_b.get_all(key))
                if len(values_a) == len(values_b):
                    match += 1
                if values_a == values_b:
                    match += 3

        elif content_a and (content_a == content_b):
            match += 4

        return match

    def next_flow(self, request):
        """
            Returns the next flow object, or None if no matching flow was
            found.
        """
        hsh = self._hash(request)
        flows = self.flowmap.get(hsh, None)
        if flows is None:
            return None

        # if it's an exact match, great!
        if len(flows) == 1:
            candidate = flows[0]
            if (candidate.request.url == request.url and
               candidate.request.raw_content == request.raw_content):
                ctx.log.info("For request {} found exact replay match".format(request.url))
                return candidate

        # find the best match between the request and the available flow candidates
        match = -1
        flow = None
        ctx.log.debug("Candiate flows for request: {}".format(request.url))
        for candidate_flow in flows:
            candidate_match = self._match(candidate_flow.request, request)
            ctx.log.debug("  score={} url={}".format(candidate_match, candidate_flow.request.url))
            if candidate_match > match:
                match = candidate_match
                flow = candidate_flow
        ctx.log.info("For request {} best match {} with score=={}".format(request.url,
                     flow.request.url, match))
        return candidate_flow

    def configure(self, options, updated):
        self.options = options
        self.clear()
        try:
            flows = io.read_flows_from_paths(self.replayfiles)
        except exceptions.FlowReadException as e:
            raise exceptions.OptionsError(str(e))
        self.load(flows)

    def request(self, f):
        if self.flowmap:
            rflow = self.next_flow(f.request)
            if rflow:
                response = rflow.response.copy()
                response.is_replay = True
                if self.options.refresh_server_playback:
                    response.refresh()
                f.response = response
            elif self.options.replay_kill_extra:
                ctx.log.warn(
                    "server_playback: killed non-replay request {}".format(
                        f.request.url
                    )
                )
                f.response = http.HTTPResponse.make(404, b'', {'content-type': 'text/plain'})


def start():
    files = sys.argv[1:]
    print("Replaying from files: {}".format(files))
    return ServerPlayback(files)
