Query.AI

import urllib.parse
from typing import Any

import demistomock as demisto
import urllib3
from CommonServerPython import *

from CommonServerUserPython import *

# Disable insecure warnings
urllib3.disable_warnings()

DEFAULT_TIMEOUT = 60  # in seconds


class Client(BaseClient):  # type: ignore
    def __init__(
        self,
        base_url,
        verify=True,
        proxy=False,
        ok_codes=(),
        headers=None,
        auth=None,
        api_token=None,
        connection_params="{}",
        alias=None,
        timeout=DEFAULT_TIMEOUT,
    ):
        super().__init__(base_url, verify, proxy, ok_codes, headers, auth)
        self._api_token = api_token
        self.alias = alias
        self.connection_params = safe_load_json(connection_params) if connection_params else {}
        self.timeout = timeout
        self._headers = {"Content-Type": "application/json", "Accept": "application/json", "Authorization": self._api_token}

    def queryai_http_request(self, method, url_suffix, json_data=None, params=None, data=None, **kwargs):
        if not json_data:
            json_data = {}
        return self._http_request(
            method=method, url_suffix=url_suffix, params=params, json_data=json_data, data=data, timeout=self.timeout, **kwargs
        )

    def run_query(
        self, query: str, connection_params: str = None, workflow_params: str = None, time_text: str = None, alias: str = None
    ) -> dict[str, Any]:
        """Gets result from QueryAI for the given query

        :type query: ``str``
        :param query: Query to be run

        :type connection_params: ``str``
        :param connection_params: Connection parameters as Stringified JSON

        :type workflow_params: ``str``
        :param workflow_params: Workflow parameters as Stringified JSON

        :type time_text: ``str``
        :param time_text: Search time-period

        :type alias: ``str``
        :param alias: Platform alias name

        :return: dict containing the response from the API
        :rtype: ``Dict[str, Any]``
        """

        if workflow_params:
            workflow_params = safe_load_json(workflow_params)

        if connection_params:
            connection_params = safe_load_json(connection_params)
        else:
            connection_params = self.connection_params  # type: ignore[assignment]

        if not alias:
            alias = self.alias

        return self.queryai_http_request(
            method="POST",
            url_suffix="/query",
            json_data={
                "query": query,
                "connection_params": connection_params,
                "workflow_params": workflow_params,
                "time_text": time_text,
                "alias": alias,
            },
        )


def generate_drilldown_url(query, alias, time_text=None, workflow_params=None):
    base_drilldown = "https://app.query.ai/login;"
    questions_url_param = f"questions={urllib.parse.quote(query)};"
    alias_url_param = f"alias={urllib.parse.quote(alias)};"
    drilldown_url = base_drilldown + questions_url_param + alias_url_param

    if time_text:
        query_duration_url_param = f"queryDuration={urllib.parse.quote(time_text)};"
        drilldown_url += query_duration_url_param
    if workflow_params:
        workflow_params_url_param = f"params={urllib.parse.quote(workflow_params)};"
        drilldown_url += workflow_params_url_param

    return drilldown_url


def queryai_run_query_command(client: Client, args: dict[str, Any]) -> CommandResults:  # type: ignore
    """queryai-run-query command: Returns response for the query being run on QueryAI

    :type client: ``Client``
    :param client: QueryAI client to use

    :type args: ``Dict[str, Any]``
    :param args:
        all command arguments, usually passed from ``demisto.args()``.
        ``args['query']`` query to run

    :return:
        A ``CommandResults`` object that is then passed to ``return_results``,
        that contains an alert

    :rtype: ``CommandResults``
    """

    query = args.get("query", "")
    connection_params = args.get("connection_params", client.connection_params)
    workflow_params = args.get("workflow_params", None)
    time_text = args.get("time_text", None)
    alias = args.get("alias", client.alias)

    if not query:
        raise ValueError('Missing argument: "query"')

    try:
        result = client.run_query(
            query=query, connection_params=connection_params, workflow_params=workflow_params, time_text=time_text, alias=alias
        )

        drilldown_url = f"### Click here to [see details]({generate_drilldown_url(query, alias, time_text, workflow_params)})"
        readable_output = tableToMarkdown(f"Query.AI Result for the query: {query}", result["data"]) if result.get("data") else ""
        readable_output = readable_output + "### " + result["reply"] + "\n" if result.get("reply") else readable_output
        readable_output += drilldown_url
        reply = {"result": result["data"] if result.get("data") else result["reply"], "markdown_string": readable_output}
        return CommandResults(
            readable_output=readable_output, outputs_prefix="QueryAI.query", outputs_key_field="", outputs=reply
        )
    except DemistoException as e:
        return_error(str(e))


def test_module(client: Client) -> str:
    """Tests API connectivity and authentication'

    Returning 'ok' indicates that the integration works like it is supposed to.
    Connection to the service is successful.
    Raises exceptions if something goes wrong.

    :type client: ``Client``
    :param client: QueryAI client to use

    :return: 'ok' if test passed, anything else will fail the test.
    :rtype: ``str``
    """

    try:
        client.run_query(query="hello")
    except DemistoException as e:
        if "Forbidden" in str(e):
            return "Authorization Error: make sure API token is correctly set"
        elif "requests.exceptions.ConnectionError" in str(e) or "Error in API call" in str(e):
            return "Connection Error - Check that the Query.AI Proxy URL parameter is correct."
        else:
            raise e
    return "ok"


def main() -> None:
    """main function, parses params and runs command functions"""
    params = demisto.params()
    api_token = params.get("api_token")
    base_url = urljoin(params["url"], "/api/v1")
    alias = params.get("alias")
    connection_params = params.get("connection_params", "{}")
    timeout = int(params.get("timeout", DEFAULT_TIMEOUT))

    verify_certificate = not params.get("insecure", False)
    proxy = params.get("proxy", False)

    demisto.debug(f"Command being called is {demisto.command()}")
    try:
        headers = {}  # type:  Dict[str, str]
        client = Client(
            base_url=base_url,
            verify=verify_certificate,
            headers=headers,
            proxy=proxy,
            api_token=api_token,
            alias=alias,
            connection_params=connection_params,
            timeout=timeout,
        )

        if demisto.command() == "test-module":
            # This is the call made when pressing the integration Test button.
            result = test_module(client)
            return_results(result)

        elif demisto.command() == "queryai-run-query":
            return_results(queryai_run_query_command(client, demisto.args()))

        else:
            return_error(f"Unsupported Command: {demisto.command()}.\n")

    # Log exceptions and return errors
    except Exception as e:
        demisto.error(traceback.format_exc())  # print the traceback
        return_error(f"Failed to execute {demisto.command()} command.\nError:\n{e!s}")


if __name__ in ("__main__", "__builtin__", "builtins"):
    main()