def AI(analyze: str) -> dict[str, any]:
    # Prompt about what the query is all about
    prompt = f"""
        Do a vulnerability analysis report on the following JSON data and
        follow the following rules:
        1) Calculate the criticality score.
        2) Return all the open ports within the open_ports list.
        3) Return all the closed ports within the closed_ports list.
        4) Return all the filtered ports within the filtered_ports list.

        output format: {{
            "open_ports": [],
            "closed_ports": [],
            "filtered_ports": [],
            "criticality_score": ""
            }}

        data = {analize}
    """
    try:
        # A structure for the request
        completion = openai.Completion.create(
            engine=model_engine,
            prompt=prompt,
            max_tokens=1024,
            n=1,
            stop=None,
        )
        response = completion.choices[0].text

        # Assuming extract_ai_output returns a dictionary
        extracted_data = extract_ai_output(response)
    except KeyboardInterrupt:
        print("Bye")
        quit()

    # Store outputs in a dictionary
    ai_output = {
        "open_ports": extracted_data.get("open_ports"),
        "closed_ports": extracted_data.get("closed_ports"),
        "filtered_ports": extracted_data.get("filtered_ports"),
        "criticality_score": extracted_data.get("criticality_score")
    }

    return ai_output

    prompt = f"""
        Do a vulnerability analysis report on the following JSON data provided.
        It's the data extracted from my network scanner.
        follow the following rules for analysis:
        1) Calculate the criticality score based on the service or CVE.
        2) Return all the open ports within the open_ports list.
        3) Return all the closed ports within the closed_ports list.
        4) Return all the filtered ports within the filtered_ports list.
        6) Keep the highest possible accuracy.
        7) Do not provide unwanted explanations.
        8) Only provide details in the output_format provided.

        output_format: {{
            "open_ports": [],
            "closed_ports": [],
            "filtered_ports": [],
            "criticality_score": ""
            }}

        data = {analize}
    """

def extract_ai_output(ai_output: str) -> dict[str, Any]:
    result = {
        "open_ports": [],
        "closed_ports": [],
        "filtered_ports": [],
        "criticality_score": ""
    }

    # Match and extract ports
    open_ports_match = re.search(r'"open_ports": \[([^\]]*)\]', ai_output)
    closed_ports_match = re.search(r'"closed_ports": \[([^\]]*)\]', ai_output)
    filtered_ports_match = re.search(
        r'"filtered_ports": \[([^\]]*)\]', ai_output)

    # If found, convert string of ports to list
    if open_ports_match:
        result["open_ports"] = list(
            map(cast(Callable[[Any], str], int),
                open_ports_match.group(1).split(',')))
    if closed_ports_match:
        result["closed_ports"] = list(
            map(cast(Callable[[Any], str], int),
                closed_ports_match.group(1).split(',')))
    if filtered_ports_match:
        result["filtered_ports"] = list(
            map(cast(Callable[[Any], str], int),
                filtered_ports_match.group(1).split(',')))

    # Match and extract criticality score
    criticality_score_match = re.search(
        r'"criticality_score": "([^"]*)"', ai_output)
    if criticality_score_match:
        result["criticality_score"] = criticality_score_match.group(1)

    return result

openai.api_key = '__API__KEY__'

from nmap_api import app

app.openai.api_key = '__API__KEY__'
app.start_api()