#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# This file is part of the Wapiti project (https://wapiti-scanner.github.io)
# Copyright (C) 2023-2025 Nicolas SURRIBAS
# Copyright (C) 2024 Cyberwatch
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
import asyncio
import re
import sys
from traceback import print_tb
from typing import Tuple, List, AsyncIterator, Dict, Optional, Deque
from logging import getLogger, WARNING, CRITICAL
from http.cookiejar import CookieJar
from urllib.parse import urlparse
import inspect
import json

from mitmproxy import addons
from mitmproxy.master import Master
from mitmproxy.options import Options
from mitmproxy.http import Request as MitmRequest
import httpx
# pylint: disable=protected-access
from playwright._impl._errors import TargetClosedError
from playwright.async_api import async_playwright, Error as PlaywrightError

from wapitiCore.net import Request
from wapitiCore.net.cookies import mitm_jar_to_cookiejar
from wapitiCore.net.response import Response
from wapitiCore.net.crawler import AsyncCrawler
from wapitiCore.net.classes import CrawlerConfiguration
from wapitiCore.net.async_stickycookie import AsyncStickyCookie
from wapitiCore.net.explorer import Explorer, EXCLUDED_MEDIA_EXTENSIONS
from wapitiCore.net.scope import Scope, wildcard_translate
from wapitiCore.main.log import log_verbose, log_blue, logging
from wapitiCore.parsers.html_parser import Html

# Mime types that a browser will commonly display instead of downloading
KNOWN_INTERPRETED_TYPES = (
    "text/plain", "text/html", "application/javascript", "text/javascript", "text/css", "application/json",
    "application/x-javascript", "image/webp", "application/manifest+json", "application/json+protobuf",
    "image/avif", "image/png", "image/gif", "image/x-icon", "font/woff2", "image/jpeg", "image/svg+xml",
    "application/ld+json", "font/ttf", "font/woff", "application/xhtml+xml",
)


def is_interpreted_type(mime_type: str) -> bool:
    for known_mime_type in KNOWN_INTERPRETED_TYPES:
        if mime_type.startswith(known_mime_type):
            return True
    return False


def set_playwright_log_level(level: int = WARNING):
    # Create logger
    logger = getLogger('playwright')
    logger.setLevel(level)


set_playwright_log_level(CRITICAL)


def decode_key_value_dict(bytes_dict: Dict[bytes, bytes], multi: bool = True) -> List[Tuple[str, str]]:
    result = []
    for key, value in bytes_dict.items(multi=multi):
        if isinstance(key, bytes):
            key = key.decode("utf-8", errors="ignore")
        if isinstance(value, bytes):
            value = value.decode("utf-8", errors="ignore")
        result.append((key, value))
    return result


def mitm_to_wapiti_request(mitm_request: MitmRequest) -> Optional[Request]:
    post_params = []
    enctype = ""
    if mitm_request.urlencoded_form:
        post_params.extend(decode_key_value_dict(mitm_request.urlencoded_form))

    if mitm_request.multipart_form:
        enctype = "multipart/form-data"
        # This is wrong but mitmproxy doesn't give files entries (along with filename, mime type)
        # so we will consider values as simple fields but still send using multipart encoding
        post_params.extend(decode_key_value_dict(mitm_request.multipart_form))

    if not enctype and mitm_request.method in ["POST", "PUT", "PATCH"]:
        enctype = mitm_request.headers.get("Content-Type", "")

    if enctype == "application/json":
        try:
            json.loads(mitm_request.text)
        except json.JSONDecodeError:
            # Ignore request if JSON data seems invalid
            return None

        post_params = mitm_request.text

    request_headers = httpx.Headers(decode_key_value_dict(mitm_request.headers))

    return Request(
        path=mitm_request.url,
        method=mitm_request.method,
        referer=mitm_request.headers.get("Referer", ""),
        post_params=post_params or None,
        enctype=enctype,
        headers=request_headers
    )


class MitmFlowToWapitiRequests:
    def __init__(self, data_queue: asyncio.Queue, headers: httpx.Headers, scope: Scope, drop_cookies: bool = False):
        self._queue = data_queue
        self._headers = headers
        self._drop_cookies = drop_cookies
        self._scope = scope

    async def request(self, flow):
        for key, value in self._headers.items():
            # This will use our user-agent too
            flow.request.headers[key] = value

    async def response(self, flow):
        if flow.request.method.upper() == "CONNECT":
            return

        if 400 <= flow.response.status_code < 500:
            # Those are certainly broken links, and we don't want to deal with that
            return

        if self._drop_cookies:
            if "set-cookie" in flow.response.headers:
                del flow.response.headers["set-cookie"]

        content_type = flow.response.headers.get("Content-Type", "text/plain").split(";")[0]
        flow.response.stream = False

        # We only need this for the automated part, a human without the headless crawler will be able to deal
        # with download popups. Let's keep in mind we should tweak that later.
        is_forced_download = flow.response.headers.get("content-disposition", "").startswith("attachment")
        if not is_interpreted_type(content_type) or is_forced_download:
            flow.response.status_code = 200
            flow.response.content = b"Lasciate ogne speranza, voi ch'intrate."
            flow.response.headers["content-type"] = "text/plain"
            flow.response.headers.pop("content-disposition", "")
            return

        redirection = flow.response.headers.get("location", "")
        if redirection.startswith(("http://", "https://")) and not self._scope.check(redirection):
            flow.response.status_code = 200
            flow.response.content = b"Lasciate ogne speranza, voi ch'intrate."
            flow.response.headers["content-type"] = "text/plain"
            return

        if "text" in content_type or "json" in content_type or "html" in content_type or "xml" in content_type:
            request = mitm_to_wapiti_request(flow.request)
            if request is None:
                return

            decoded_headers = decode_key_value_dict(flow.response.headers)

            response = Response(
                httpx.Response(
                    status_code=flow.response.status_code,
                    headers=decoded_headers,
                    # httpx expect the raw content (not decompressed)
                    content=flow.response.raw_content,
                ),
                url=flow.request.url
            )

            await self._queue.put(
                (request, response)
            )


async def launch_proxy(
        port: int,
        data_queue: asyncio.Queue,
        headers: httpx.Headers,
        cookies: CookieJar,
        scope: Scope,
        proxy: Optional[str] = None,
        drop_cookies: bool = False,
):
    log_blue(
        f"Launching MitmProxy on port {port}. Configure your browser to use it, press ctrl+c when you are done."
    )
    opt = Options()
    # We can use an upstream proxy that way but socks is not supported
    if proxy:
        log_blue(f"Using upstream proxy {proxy}")
        opt.update(mode=f"upstream:{proxy}")

    opt.update(listen_port=port,  http2=False, ssl_insecure=True)
    master = Master(opt)
    master.addons.add(addons.core.Core())
    master.addons.add(addons.proxyserver.Proxyserver())
    master.addons.add(addons.next_layer.NextLayer())
    # mitmproxy will generate an authority cert in the ~/.mitmproxy directory. Load it in your browser.
    master.addons.add(addons.tlsconfig.TlsConfig())
    # If ever we want to have both the interception proxy and an automated crawler then we need to sync cookies
    # This mitmproxy module will do that and also load init cookies in the internal jar
    master.addons.add(AsyncStickyCookie(cookies))
    # Finally here is our custom addon that will generate Wapiti Request and Response objects and push them to the queue
    master.addons.add(MitmFlowToWapitiRequests(data_queue, headers, scope, drop_cookies))
    try:
        await master.run()
    except asyncio.CancelledError:
        log_blue("Stopping mitmproxy")
    return master.addons.get("asyncstickycookie").jar


def extract_requests(html: Html, request: Request):
    candidates = html.links + html.js_redirections + html.html_redirections + list(html.extra_urls)

    for link in candidates:
        url_parts = urlparse(link)
        if url_parts.path.endswith((".css", ".js")):
            continue

        if not url_parts.query and url_parts.path.endswith(EXCLUDED_MEDIA_EXTENSIONS):
            continue

        next_request = Request(link, link_depth=request.link_depth + 1)
        yield next_request

        if "?" in link:
            next_request = Request(link.split("?")[0], link_depth=request.link_depth + 1)
            yield next_request

    for form in html.iter_forms():
        for form_request in form.to_requests():
            # if scope.check(form) and form not in to_explore and form not in excluded_requests:
            form_request.link_depth = request.link_depth + 1
            yield form_request


async def click_in_webpage(page, request: Request, wait_time: float, timeout: float):
    # We are using XPath because CSS selectors doesn't allow to combine nth-of-type with other cool stuff
    for xpath_selector in (".//button", ".//*[@role=\"button\" and not(@href)]"):
        try:
            buttons = await page.query_selector_all(xpath_selector)
        except PlaywrightError:
            continue

        for button in buttons:
            try:
                await button.click(timeout=timeout * 1000)
            except PlaywrightError:
                continue
            else:
                await page.wait_for_timeout(wait_time * 1000)
                if page.url != request.url_with_fragment:
                    await page.goto(
                        request.url_with_fragment,
                        timeout=timeout * 1000,
                        wait_until="domcontentloaded"
                    )


# pylint: disable=too-many-positional-arguments
async def launch_headless_explorer(
        stop_event: asyncio.Event,
        crawler: AsyncCrawler,
        to_explore: Deque[Request],
        scope: Scope,
        proxy_port: int,
        excluded_requests: List[Request],
        exclusion_regexes: List[re.Pattern],
        visibility: str = "hidden",
        wait_time: float = 2.,
        max_depth: int = 20,
):
    # The headless browser will be configured to use the MITM proxy
    # The intercepting will be in charge of generating Request objects.
    # This is the only way as a headless browser can't provide us response headers.
    proxy_settings = {
        "server": f"http://127.0.0.1:{proxy_port}"
    }

    # We need to make a copy of this list otherwise requests won't make their way into async_explore (because list is
    # shared). Also, we want our own list here because we will see URLs with anchors that the proxy can't catch.
    excluded_requests = list(excluded_requests)

    try:
        async with async_playwright() as p:
            browser = await p.firefox.launch(
                headless=visibility == "hidden",
                proxy=proxy_settings,
                firefox_user_prefs={
                    "network.proxy.allow_hijacking_localhost": True,
                    "devtools.jsonview.enabled": False,
                }
            )
            context = await browser.new_context(ignore_https_errors=True)
            page = await context.new_page()

            try:
                while to_explore and not stop_event.is_set():
                    request = to_explore.popleft()
                    excluded_requests.append(request)

                    if request.method == "GET":
                        try:
                            await page.goto(
                                request.url_with_fragment,
                                timeout=crawler.timeout.connect * 1000,
                                wait_until="domcontentloaded"
                            )
                            await page.wait_for_timeout(wait_time * 1000)
                            # We may be redirected outside our target so let's check the URL first
                            if not scope.check(page.url):
                                continue

                            page_source = await page.content()
                            await click_in_webpage(page, request, wait_time, timeout=crawler.timeout.connect)
                        except (PlaywrightError, asyncio.TimeoutError) as exception:
                            if isinstance(exception, TargetClosedError):
                                break

                            logging.error("%s generated an exception: %s", request, type(exception).__name__)
                            continue
                    else:
                        try:
                            response = await crawler.async_send(request, timeout=crawler.timeout.connect)
                        except httpx.RequestError as exception:
                            logging.error("%s generated an exception: %s", request, type(exception).__name__)
                            continue

                        page_source = response.content

                    if request.link_depth == max_depth:
                        continue

                    html = Html(page_source, request.url, allow_fragments=True)

                    for next_request in extract_requests(html, request):
                        if not scope.check(next_request):
                            continue

                        if any(regex.match(next_request.url) for regex in exclusion_regexes):
                            continue

                        if next_request not in to_explore and next_request not in excluded_requests:
                            to_explore.append(next_request)
            finally:
                await browser.close()

    except Exception as exception:  # pylint: disable=broad-except
        if isinstance(exception, TargetClosedError):
            return

        exception_traceback = sys.exc_info()[2]
        print_tb(exception_traceback)
        frm = inspect.trace()[-1]
        mod = inspect.getmodule(frm[0])
        logging.error(
            "Headless browser stopped prematurely due to exception: %s.%s", mod.__name__, type(exception).__name__
        )

    await asyncio.sleep(1)
    stop_event.set()


class InterceptingExplorer(Explorer):
    # pylint: disable=too-many-positional-arguments
    def __init__(
            self,
            crawler_configuration: CrawlerConfiguration,
            scope: Scope,
            stop_event: asyncio.Event,
            parallelism: int = 8,
            mitm_port: int = 8080,
            proxy: Optional[str] = None,
            drop_cookies: bool = False,
            headless: str = "no",
            cookies: Optional[CookieJar] = None,
            wait_time: float = 2.
    ):
        super().__init__(crawler_configuration, scope, stop_event, parallelism)
        self._mitm_port = mitm_port
        self._proxy = proxy
        self._drop_cookies = drop_cookies
        self._headless = headless
        self._final_cookies = None
        self._cookies = cookies or CookieJar()
        self._wait_time = wait_time
        self._headless_task = None
        self._mitm_task = None
        self._queue = None

    async def process_requests(self, excluded_requests, exclusion_regexes):
        while True:
            try:
                request, response = self._queue.get_nowait()
            except asyncio.QueueEmpty:
                if self._stopped.is_set():
                    break
                await asyncio.sleep(.1)
                continue
            except KeyboardInterrupt:
                break

            self._queue.task_done()

            # Scope check and deduplication are made here
            if not self._scope.check(request) or request in self._processed_requests:
                continue

            # Check for exclusion here because we don't have full control over the headless browser
            if request in excluded_requests or any(regex.match(request.url) for regex in exclusion_regexes):
                continue

            dir_name = request.dir_name
            if self._max_files_per_dir and self._file_counts[dir_name] >= self._max_files_per_dir:
                continue

            self._file_counts[dir_name] += 1

            if self.has_too_many_parameters(request):
                continue

            if self._qs_limit and request.parameters_count:
                self._pattern_counts[request.pattern] += 1

            yield request, response
            self._processed_requests.append(request)
            log_verbose(f"[+] {request}")

    async def async_explore(
            self,
            to_explore: Deque[Request],
            excluded_urls: list = None
    ) -> AsyncIterator[Tuple[Request, Response]]:
        self._queue = asyncio.Queue()

        exclusion_regexes = []
        excluded_requests = []

        if isinstance(excluded_urls, list):
            for bad_request in excluded_urls:
                if isinstance(bad_request, str):
                    exclusion_regexes.append(wildcard_translate(bad_request))
                elif isinstance(bad_request, Request):
                    excluded_requests.append(bad_request)

        # Launch proxy as asyncio task
        self._mitm_task = asyncio.create_task(
            launch_proxy(
                self._mitm_port,
                self._queue,
                self._crawler.headers,
                self._cookies,
                self._scope,
                proxy=self._proxy,
                drop_cookies=self._drop_cookies,
            )
        )

        if self._headless == "no":
            # No headless crawler, just intercepting mode so no starting URLs
            to_explore.clear()
        else:
            self._headless_task = asyncio.create_task(
                launch_headless_explorer(
                    self._stopped,
                    self._crawler,
                    to_explore,
                    scope=self._scope,
                    proxy_port=self._mitm_port,
                    excluded_requests=excluded_requests,
                    exclusion_regexes=exclusion_regexes,
                    visibility=self._headless,
                    wait_time=self._wait_time,
                    max_depth=self._max_depth,
                )
            )

        async for request, response in self.process_requests(excluded_requests, exclusion_regexes):
            yield request, response
            if self._stopped.is_set():
                break

    def empty_queue(self):
        while not self._queue.empty():
            self._queue.get_nowait()
            self._queue.task_done()

    async def clean(self):
        self.empty_queue()

        # The headless crawler must stop when the stop event is set, let's just wait for it
        if self._headless_task:
            await self._headless_task

        # We are canceling the mitm proxy, but we could have used a special request to shut down the master to.
        # https://docs.mitmproxy.org/stable/addons-examples/#shutdown
        self._mitm_task.cancel()
        self._final_cookies = await self._mitm_task
        await self._crawler.close()

    @property
    def cookie_jar(self):
        return mitm_jar_to_cookiejar(self._final_cookies)
