# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

# This file was copied from  mitmproxy/mitmproxy/addons/serverplayback.py release tag 4.0.4
# and modified by Florin Strugariu

# Altered features:
# * returns 404 rather than dropping the whole HTTP/2 connection on the floor
# * remove the replay packages that don't have any content in their response package
from __future__ import absolute_import, print_function

import os
import json
import hashlib
from collections import defaultdict
import time
import signal

import typing

from six.moves import urllib
from mitmproxy import ctx, http
from mitmproxy import exceptions
from mitmproxy import io

# PATCHING AREA  - ALLOWS HTTP/2 WITH NO CERT SNIFFING
from mitmproxy.proxy.protocol import tls
from mitmproxy.proxy.protocol.http2 import Http2Layer, SafeH2Connection

_PROTO = {}


@property
def _alpn(self):
    proto = _PROTO.get(self.server_sni)
    if proto is None:
        return self.server_conn.get_alpn_proto_negotiated()
    if proto.startswith("HTTP/2"):
        return b"h2"
    elif proto.startswith("HTTP/1"):
        return b"h1"
    return b""


tls.TlsLayer.alpn_for_client_connection = _alpn


def _server_conn(self):
    if not self.server_conn.connected() and self.server_conn not in self.connections:
        # we can't use ctx.log in this layer
        print("Ignored CONNECT call on upstream server")
        return
    if self.server_conn.connected():
        import h2.config

        config = h2.config.H2Configuration(
            client_side=True,
            header_encoding=False,
            validate_outbound_headers=False,
            validate_inbound_headers=False,
        )
        self.connections[self.server_conn] = SafeH2Connection(
            self.server_conn, config=config
        )
    self.connections[self.server_conn].initiate_connection()
    self.server_conn.send(self.connections[self.server_conn].data_to_send())


Http2Layer._initiate_server_conn = _server_conn


def _remote_settings_changed(self, event, other_conn):
    if other_conn not in self.connections:
        # we can't use ctx.log in this layer
        print("Ignored remote settings upstream")
        return True
    new_settings = dict(
        [(key, cs.new_value) for (key, cs) in event.changed_settings.items()]
    )
    self.connections[other_conn].safe_update_settings(new_settings)
    return True


Http2Layer._handle_remote_settings_changed = _remote_settings_changed
# END OF PATCHING


class AlternateServerPlayback:
    def __init__(self):
        ctx.master.addons.remove(ctx.master.addons.get("serverplayback"))
        self.flowmap = {}
        self.configured = False
        self.netlocs = defaultdict(lambda: defaultdict(int))
        self.calls = []
        self._done = False
        self._replayed = 0
        self._not_replayed = 0
        self._recordings_used = 0
        self.mitm_version = ctx.mitmproxy.version.VERSION

        ctx.log.info("MitmProxy version: %s" % self.mitm_version)

    def load(self, loader):
        loader.add_option(
            "server_replay_files",
            typing.Sequence[str],
            [],
            "Replay server responses from a saved file.",
        )
        loader.add_option(
            "upload_dir",
            str,
            "",
            "Upload directory",
        )

    def load_flows(self, flows):
        """
        Replay server responses from flows.
        """
        for i in flows:
            if i.type == "websocket":
                # Mitmproxy can't replay WebSocket packages.
                ctx.log.info(
                    "Recorded response is a WebSocketFlow. Removing from recording list as"
                    "  WebSockets are disabled"
                )
            elif i.response and self.mitm_version in ("4.0.2", "4.0.4", "5.1.1"):
                # see: https://github.com/mitmproxy/mitmproxy/issues/3856
                f = self.flowmap.setdefault(
                    self._hash(i), {"flow": None, "reply_count": 0}
                )
                # overwrite with new flow if already hashed
                f["flow"] = i
            else:
                ctx.log.info(
                    "Recorded request %s has no response. Removing from recording list"
                    % i.request.url
                )
        ctx.master.addons.trigger("update", [])

    def load_files(self, paths):
        try:
            if "," in paths[0]:
                paths = paths[0].split(",")
            for path in paths:
                ctx.log.info("Loading flows from %s" % path)
                if not os.path.exists(path):
                    raise Exception("File does not exist!")
                try:
                    flows = io.read_flows_from_paths([path])
                except exceptions.FlowReadException as e:
                    raise exceptions.CommandError(str(e))
                self.load_flows(flows)
                proto = os.path.join(os.path.dirname(path), "metadata.json")
                if os.path.exists(proto):
                    ctx.log.info("Loading proto info from %s" % proto)
                    with open(proto) as f:
                        recording_info = json.loads(f.read())
                    if recording_info.get("http_protocol", False):
                        ctx.log.info(
                            "Replaying file {} recorded on {}".format(
                                path, recording_info["recording_date"]
                            )
                        )
                        _PROTO.update(recording_info["http_protocol"])
                    else:
                        ctx.log.warn(
                            "Replaying file {} has no http_protocol info.".format(proto)
                        )
        except Exception as e:
            ctx.log.error("Could not load recording file! Stopping playback process!")
            ctx.log.error(str(e))
            ctx.master.shutdown()

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

        # unquote url
        # See Bug 1509835
        _, _, path, _, query, _ = urllib.parse.urlparse(urllib.parse.unquote(r.url))
        queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True)

        key = [str(r.port), str(r.scheme), str(r.method), str(path)]
        key.append(str(r.raw_content))
        key.append(r.host)

        for p in queriesArray:
            key.append(p[0])
            key.append(p[1])

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

    def next_flow(self, request):
        """
        Returns the next flow object, or None if no matching flow was
        found.
        """
        hsh = self._hash(request)
        if hsh in self.flowmap:
            if self.flowmap[hsh]["reply_count"] == 0:
                self._recordings_used += 1
            self.flowmap[hsh]["reply_count"] += 1
            # return the most recently added flow with this hash
            return self.flowmap[hsh]["flow"]

    def configure(self, updated):
        if not self.configured and ctx.options.server_replay_files:
            self.configured = True
            self.load_files(ctx.options.server_replay_files)

    def done(self):
        if self._done or not ctx.options.upload_dir:
            return

        replay_confidence = float(self._replayed) / (
            self._replayed + self._not_replayed
        )
        recording_proportion_used = (
            0
            if self._recordings_used == 0
            else float(self._recordings_used) / len(self.flowmap)
        )
        stats = {
            "totals": dict(self.netlocs),
            "calls": self.calls,
            "replayed": self._replayed,
            "not-replayed": self._not_replayed,
            "replay-confidence": int(replay_confidence * 100),
            "recording-proportion-used": int(recording_proportion_used * 100),
        }
        file_name = (
            "mitm_netlocs_%s.json"
            % os.path.splitext(os.path.basename(ctx.options.server_replay_files[0]))[0]
        )
        path = os.path.normpath(os.path.join(ctx.options.upload_dir, file_name))
        try:
            with open(path, "w") as f:
                f.write(json.dumps(stats, indent=2, sort_keys=True))
        finally:
            self._done = True

    def request(self, f):
        if self.flowmap:
            try:
                rflow = self.next_flow(f)
                if rflow:
                    response = rflow.response.copy()
                    response.is_replay = True
                    # Refresh server replay responses by adjusting date, expires and
                    # last-modified headers, as well as adjusting cookie expiration.
                    response.refresh()

                    f.response = response
                    self._replayed += 1
                else:
                    # returns 404 rather than dropping the whole HTTP/2 connection
                    ctx.log.warn(
                        "server_playback: killed non-replay request {}".format(
                            f.request.url
                        )
                    )
                    f.response = http.HTTPResponse.make(
                        404, b"", {"content-type": "text/plain"}
                    )
                    self._not_replayed += 1

                # collecting stats only if we can dump them (see .done())
                if ctx.options.upload_dir:
                    parsed_url = urllib.parse.urlparse(
                        urllib.parse.unquote(f.request.url)
                    )
                    self.netlocs[parsed_url.netloc][f.response.status_code] += 1
                    self.calls.append(
                        {
                            "time": str(time.time()),
                            "url": f.request.url,
                            "response_status": f.response.status_code,
                        }
                    )
            except Exception as e:
                ctx.log.error("Could not generate response! Stopping playback process!")
                ctx.log.info(e)
                ctx.master.shutdown()

        else:
            ctx.log.error("Playback library is empty! Stopping playback process!")
            ctx.master.shutdown()
            return


playback = AlternateServerPlayback()

if hasattr(signal, "SIGBREAK"):
    # allows the addon to dump the stats even if mitmproxy
    # does not call done() like on windows termination
    # for this, the parent process sends CTRL_BREAK_EVENT which
    # is received as an SIGBREAK event
    def _shutdown(sig, frame):
        ctx.master.shutdown()

    signal.signal(signal.SIGBREAK, _shutdown)

addons = [playback]
