#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Converts a mitmproxy dump file to a swagger schema."""
import argparse
import json
import os
import re
import sys
import traceback
import urllib
from typing import Any, Optional, Sequence, Union

import msgpack
import ruamel.yaml
from mitmproxy.exceptions import FlowReadException

from mitmproxy2swagger import console_util, swagger_util
from mitmproxy2swagger.har_capture_reader import HarCaptureReader, har_archive_heuristic
from mitmproxy2swagger.mitmproxy_capture_reader import (
    MitmproxyCaptureReader,
    mitmproxy_dump_file_huristic,
)


def path_to_regex(path):
    # replace the path template with a regex
    path = re.escape(path)
    path = path.replace(r"\{", "(?P<")
    path = path.replace(r"\}", ">[^/]+)")
    path = path.replace(r"\*", ".*")
    return "^" + path + "$"


def strip_query_string(path):
    # remove the query string from the path
    return path.split("?")[0]


def set_key_if_not_exists(dict, key, value):
    if key not in dict:
        dict[key] = value


def progress_callback(progress):
    console_util.print_progress_bar(progress)


def detect_input_format(file_path):
    har_score = har_archive_heuristic(file_path)
    mitmproxy_score = mitmproxy_dump_file_huristic(file_path)
    if "MITMPROXY2SWAGGER_DEBUG" in os.environ:
        print("har score: " + str(har_score))
        print("mitmproxy score: " + str(mitmproxy_score))
    if har_score > mitmproxy_score:
        return HarCaptureReader(file_path, progress_callback)
    return MitmproxyCaptureReader(file_path, progress_callback)

def str2re(re_str):
    return re.compile("^" + re_str + "$")

def get_abs_path(path):
    base_dir = os.getcwd()
    abs_path = os.path.join(base_dir, path)
    return abs_path


def main(override_args: Optional[Sequence[str]] = None):
    parser = argparse.ArgumentParser(
        description="Converts a mitmproxy dump file or HAR to a swagger schema."
    )
    parser.add_argument(
        "-i",
        "--input",
        help="The input mitmproxy dump file or HAR dump file (from DevTools)",
        required=True,
    )
    parser.add_argument(
        "-o",
        "--output",
        help="The output swagger schema file (yaml). If it exists, new endpoints will be added",
        required=True,
    )
    parser.add_argument("-p", "--api-prefix", help="The api prefix", required=True)
    parser.add_argument(
        "-e",
        "--examples",
        action="store_true",
        help="Include examples in the schema. This might expose sensitive information.",
    )
    parser.add_argument(
        "-hd",
        "--headers",
        action="store_true",
        help="Include headers in the schema. This might expose sensitive information.",
    )
    parser.add_argument(
        "-f",
        "--format",
        choices=["flow", "har"],
        help="Override the input file format auto-detection.",
    )
    parser.add_argument(
        "-r",
        "--param-regex",
        default="[0-9]+",
        help="Regex to match parameters in the API paths. Path segments that match this regex will be turned into parameter placeholders.",
    )
    parser.add_argument(
        "-s",
        "--suppress-params",
        action="store_true",
        help="Do not include API paths that have the original parameter values, only the ones with placeholders.",
    )
    parser.add_argument(
        "-eh",
        "--excluded-headers",
        action="extend",
        nargs="+",
        type=str,
        default=[],
        help="Exclude the headers in schema e.g. -eh 'sec-ch-ua-platform' 'cache-control'",
    )   
    parser.add_argument(
        "-pn",
        "--populate-new",
        action="store_true",
        default=False,
        help="Populate doc for new endpoints in actual requests",
    )
    parser.add_argument(
        "-em",
        "--excluded-methods",
        action="extend",
        nargs="+",
        type=str,
        default=[],
        help="Exclude the request with specific method in schema e.g. -em 'OPTIONS' 'PUT'",
    )  
    parser.add_argument(
        "-w",
        "--overwrite",
        action="store_true",
        help="Overwrite the whole swagger file",
    )
    parser.add_argument(
        "-at",
        "--add-tags",
        action="store_true",
        help="Add swagger tag to the endpoint automatically base on the url",
    )
    parser.add_argument(
        "-tj",
        "--tags-json",
        help='The JSON file mapping URLs to tags, the content can be e.g. {"auth":["/api/auth/.+", "/api/account/.+"]}',
    )
    
    args = parser.parse_args(override_args)

    try:
        args.param_regex = str2re(args.param_regex)
    except re.error as e:
        print(
            f"{console_util.ANSI_RED}Invalid path parameter regex: {e}{console_util.ANSI_RESET}"
        )
        sys.exit(1)

    yaml = ruamel.yaml.YAML()

    capture_reader: Union[MitmproxyCaptureReader, HarCaptureReader]
    if args.format == "flow" or args.format == "mitmproxy":
        capture_reader = MitmproxyCaptureReader(args.input, progress_callback)
    elif args.format == "har":
        capture_reader = HarCaptureReader(args.input, progress_callback)
    else:
        capture_reader = detect_input_format(args.input)

    swagger = None

    # try loading the existing swagger file
    if not args.overwrite:
        try:
            with open(get_abs_path(args.output), "r") as f:
                swagger = yaml.load(f)
        except FileNotFoundError:
            print("No existing swagger file found. Creating new one.")
    if swagger is None:
        swagger = ruamel.yaml.comments.CommentedMap(
            {
                "openapi": "3.0.0",
                "info": {
                    "title": args.input + " Mitmproxy2Swagger",
                    "version": "1.0.0",
                },
            }
        )

    if args.add_tags or args.tags_json:
        swagger["tags"] = []

    tags = None
    if args.tags_json:
        with open(get_abs_path(args.tags_json), "r") as f:
            tags = json.load(f)


    # strip the trailing slash from the api prefix
    args.api_prefix = args.api_prefix.rstrip("/")
    args.excluded_headers = [h.lower() for h in args.excluded_headers]
    args.excluded_methods = [h.lower() for h in args.excluded_methods]

    if "servers" not in swagger or swagger["servers"] is None:
        swagger["servers"] = []

    # add the server if it doesn't exist
    if not any(server["url"] == args.api_prefix for server in swagger["servers"]):
        swagger["servers"].append(
            {"url": args.api_prefix, "description": "The default server"}
        )

    if "paths" not in swagger or swagger["paths"] is None:
        swagger["paths"] = {}

    if "x-path-templates" not in swagger or swagger["x-path-templates"] is None:
        swagger["x-path-templates"] = []

    path_templates = []
    for path in swagger["paths"]:
        path_templates.append(path)

    # also add paths from the the x-path-templates array
    if "x-path-templates" in swagger and swagger["x-path-templates"] is not None:
        for path in swagger["x-path-templates"]:
            path_templates.append(path)

    # new endpoints will be added here so that they can be added as comments in the swagger file
    new_path_templates = []
    path_template_regexes = [re.compile(path_to_regex(path)) for path in path_templates]

    def is_param(param_value):
        return args.param_regex.match(param_value) is not None

    def get_suggested_path(path):
        suggested_path = None
        segments = path.split("/")
        has_param = any(is_param(segment) for segment in segments)
        if has_param:
            # replace digit segments with {id}, {id1}, {id2} etc
            new_segments = []
            param_id = 0
            for segment in segments:
                if is_param(segment):
                    param_name = "id" + str(param_id)
                    if param_id == 0:
                        param_name = "id"
                    new_segments.append("{" + param_name + "}")
                    param_id += 1
                else:
                    new_segments.append(segment)
            suggested_path = "/".join(new_segments)
        return suggested_path

    common_prefix = None
    ref_path = None

    try:
        for req in capture_reader.captured_requests():
            # strip the api prefix from the url
            url = req.get_matching_url(args.api_prefix)

            if url is None:
                continue
            method = req.get_method().lower()
            path = strip_query_string(url).removeprefix(args.api_prefix)
            status = req.get_response_status_code()

            # check if the path matches any of the path templates, and save the index
            path_template_index = None
            for i, path_template_regex in enumerate(path_template_regexes):
                if path_template_regex.match(path):
                    path_template_index = i
                    break
            if path_template_index is None:
                if method not in args.excluded_methods:
                    if path not in new_path_templates:
                        new_path_templates.append(path)
                    if args.populate_new:
                        path_templates.append(get_suggested_path(path) or path)
                        path_template_regexes.append(re.compile(path_to_regex(path)))
                        path_template_index = -1
                    else:
                        continue
            
            if method in args.excluded_methods:
                continue

            path_template_to_set = path_templates[path_template_index]
            set_key_if_not_exists(swagger["paths"], path_template_to_set, {})

            if not ref_path:
                ref_path = path_template_to_set
                common_prefix = path_template_to_set
            else:
                path_prefix = os.path.commonpath([ref_path, path_template_to_set])
                if len(path_prefix)<len(common_prefix):
                    common_prefix = path_prefix

            set_key_if_not_exists(
                swagger["paths"][path_template_to_set],
                method,
                {
                    "tags": [],
                    "summary": swagger_util.path_template_to_endpoint_name(
                        method, path_template_to_set
                    ),
                    "responses": {},
                },
            )

            if args.tags_json:
                for tag, endpoints in tags.items():
                    if any(str2re(ep).match(path_template_to_set) for ep in endpoints):
                        if tag not in swagger["paths"][path_template_to_set][method]["tags"]:
                            swagger["paths"][path_template_to_set][method]["tags"].append(tag)
                            swagger['tags'].append(tag)
                            break

            params = swagger_util.url_to_params(url, path_template_to_set)
            if args.headers:
                headers_request = swagger_util.request_to_headers(
                    req.get_request_headers() if not args.excluded_headers else \
                        {h:v for h,v in req.get_request_headers().items() if h.lower() not in args.excluded_headers}
                )
                params.extend(headers_request)
            if params is not None and len(params) > 0:
                set_key_if_not_exists(
                    swagger["paths"][path_template_to_set][method], "parameters", params
                )

            if method not in ["get", "head"]:
                body = req.get_request_body()
                if body is not None:
                    body_val = None
                    content_type = None
                    # try to parse the body as json
                    try:
                        body_val = json.loads(req.get_request_body())
                        content_type = "application/json"
                    except UnicodeDecodeError:
                        pass
                    except json.decoder.JSONDecodeError:
                        pass

                    # try to parse the body as msgpack, if it's not json
                    if body_val is None:
                        try:
                            body_val = msgpack.loads(req.get_request_body())
                            content_type = "application/msgpack"
                        except Exception:
                            pass

                    if content_type is None:
                        # try to parse the body as form data
                        try:
                            body_val_bytes: Any = dict(
                                urllib.parse.parse_qsl(
                                    body, encoding="utf-8", keep_blank_values=True
                                )
                            )
                            body_val = {}
                            did_find_anything = False
                            for key, value in body_val_bytes.items():
                                did_find_anything = True
                                body_val[key.decode("utf-8")] = value.decode("utf-8")
                            if did_find_anything:
                                content_type = "application/x-www-form-urlencoded"
                            else:
                                body_val = None
                        except UnicodeDecodeError:
                            pass

                    if body_val is not None:
                        content_to_set = {
                            "content": {
                                content_type: {
                                    "schema": swagger_util.value_to_schema(body_val)
                                }
                            }
                        }
                        if args.examples:
                            content_to_set["content"][content_type]["example"] = (
                                swagger_util.limit_example_size(body_val)
                            )
                        set_key_if_not_exists(
                            swagger["paths"][path_template_to_set][method],
                            "requestBody",
                            content_to_set,
                        )

            response_body = req.get_response_body()
            if response_body is not None:
                # try parsing the response as json
                try:
                    response_parsed = json.loads(response_body)
                    response_content_type = "application/json"
                except UnicodeDecodeError:
                    response_parsed = None
                except json.decoder.JSONDecodeError:
                    response_parsed = None

                if response_parsed is None:
                    # try parsing the response as msgpack, if it's not json
                    try:
                        response_parsed = msgpack.loads(response_body)
                        response_content_type = "application/msgpack"
                    except Exception:
                        response_parsed = None

                if response_parsed is not None:
                    resp_data_to_set = {
                        "description": req.get_response_reason(),
                        "content": {
                            response_content_type: {
                                "schema": swagger_util.value_to_schema(response_parsed)
                            }
                        },
                    }
                    if args.examples:
                        resp_data_to_set["content"][response_content_type][
                            "example"
                        ] = swagger_util.limit_example_size(response_parsed)
                    if args.headers:
                        resp_data_to_set["headers"] = swagger_util.response_to_headers(
                            req.get_response_headers()
                        )

                    set_key_if_not_exists(
                        swagger["paths"][path_template_to_set][method]["responses"],
                        str(status),
                        resp_data_to_set,
                    )

            if (
                "responses" in swagger["paths"][path_template_to_set][method]
                and len(swagger["paths"][path_template_to_set][method]["responses"])
                == 0
            ):
                # add a default response if there were no responses detected,
                # this is for compliance with the OpenAPI spec
                content_type = (
                    req.get_response_headers().get("content-type") or "text/plain"
                )

                swagger["paths"][path_template_to_set][method]["responses"]["200"] = {
                    "description": "OK",
                    "content": {},
                }

    except FlowReadException as e:
        print(f"Flow file corrupted: {e}")
        traceback.print_exception(*sys.exc_info())
        print(
            f"{console_util.ANSI_RED}Failed to parse the input file as '{capture_reader.name()}'. "
        )
        if not args.format:
            print(
                f"It might happen that the input format as incorrectly detected. Please try using '--format flow' or '--format har' to specify the input format.{console_util.ANSI_RESET}"
            )
        sys.exit(1)
    except ValueError as e:
        print(f"ValueError: {e}")
        # print stack trace
        traceback.print_exception(*sys.exc_info())
        print(
            f"{console_util.ANSI_RED}Failed to parse the input file as '{capture_reader.name()}'. "
        )
        if not args.format:
            print(
                f"It might happen that the input format as incorrectly detected. Please try using '--format flow' or '--format har' to specify the input format.{console_util.ANSI_RESET}"
            )
        sys.exit(1)

    if args.add_tags and ref_path != common_prefix:
        tags = set()
        tag_index = len(common_prefix.split("/")) if common_prefix!="/" else 1
        for path, p_value in swagger["paths"].items():
            for method, m_value in p_value.items():
                m_value["tags"] = [path.split("/")[tag_index]]
                tags.add(m_value["tags"][0])
        swagger["tags"].extend([{"name": tag} for tag in tags])


    new_path_templates.sort()

    # add suggested path templates
    # basically inspects urls and replaces segments containing only numbers with a parameter
    new_path_templates_with_suggestions = []
    origin_paths = []
    path_prefix = "ignore:" if not args.populate_new else "new:"
    for path in new_path_templates:
        # check if path contains number-only segments
        suggested_path = get_suggested_path(path)
        # prepend the suggested path to the new_path_templates list
        if suggested_path is not None:
            if suggested_path not in new_path_templates_with_suggestions:
                new_path_templates_with_suggestions.append(path_prefix + suggested_path)
            if path not in origin_paths:
                origin_paths.append(path)
        else:
            new_path_templates_with_suggestions.append(path_prefix+ path)

    if not args.suppress_params:
        new_path_templates_with_suggestions.extend([path_prefix+op for op in origin_paths])
    else:
        swagger["x-path-templates"][:] = [
            p for p in swagger["x-path-templates"] 
            if p.split("ignore:")[-1].split("new:")[-1].strip() not in origin_paths
        ]

    # remove the ending comments not to add them twice

    # append the contents of new_path_templates_with_suggestions to swagger['x-path-templates']
    swagger["x-path-templates"].extend(new_path_templates_with_suggestions)

    # remove elements already generated
    swagger["x-path-templates"] = [
        path for path in swagger["x-path-templates"] if path.split("ignore:")[-1].strip() not in swagger["paths"]
    ]

    # remove duplicates while preserving order
    def f7(seq):
        seen = set()
        seen_add = seen.add
        return [x for x in seq if not (x in seen or seen_add(x))]

    swagger["x-path-templates"] = f7(swagger["x-path-templates"])

    swagger["x-path-templates"] = ruamel.yaml.comments.CommentedSeq(
        swagger["x-path-templates"]
    )
    swagger["x-path-templates"].yaml_set_start_comment(
        "Remove the ignore: prefix to generate an endpoint with its URL\n"+
        "Lines that are closer to the top take precedence, the matching is greedy\n"+
        "An endpoint prefixed with `new:` is an endpoint with Doc being populated though it is not in x-path-templates"
    )
    # save the swagger file
    with open(args.output, "w") as f:
        yaml.dump(swagger, f)
    print("Done!")


if __name__ == "__main__":
    main()
