#!/usr/bin/env python3

"""APIClient.py: QRadar Ariel DB searcher"""
__author__      = "Radek Hranicky"

import logging
import time
import urllib.parse

logger = logging.getLogger('domainradar')

class ArielSearcher:
    def __init__(self, config, api):
        """
        ! Constructor of the ArielSearcher class
        @param config Program config object
        @param api QRadar APIClient object
        """
        self.config = config
        self.api = api

    def query_ok(self, ariel_query):
        """
        ! Checks whether the given Ariel query is correct
        @param ariel_query Query in Ariel Query Language (AQL)
        @return True is the query is OK, False otherwise
        """
        params = {"query_expression": ariel_query}
        resp_json = self.api.request_post("/ariel/validators/aql", params)
        if resp_json == None:
            return False

        error_messages = resp_json["error_messages"]
        if len(error_messages) == 0:
            return True
        else:
            return False

    def make_search(self, ariel_query):
        """
        ! Launches an AQL search
        @param ariel_query Query in Ariel Query Language (AQL)
        @return Unique ID that identifies the search
        """
        if self.query_ok(ariel_query):
            params = {"query_expression": ariel_query}
            resp_json = self.api.request_post("/ariel/searches", params)
            if resp_json is None:
                logger.error("Ariel search request `" + ariel_query + "` failed!")
                return None
            search_id = resp_json["search_id"]
            return search_id
        else:
            logger.error("Incorrect query: " + ariel_query)
            return None

    def is_completed(self, search_id):
        """
        ! Checks whether the given Ariel search is completed
        @param search_id Query ID
        @return True if completed, False is not, None in case of an error
        """
        if search_id is None:
            logger.error("Search ID not specified!")
            return None
        endpoint = "/ariel/searches/%s" % search_id
        resp_json = self.api.request_get(endpoint)
        if resp_json is None:
            logger.error("Lookup for search with ID " + str(search_id) + " failed!")
            return None
        if resp_json["status"] == "COMPLETED":
            return True
        else:
            return False

    def results(self, search_id):
        """
        ! Returns the result of an AQL query
        @param search_id Query ID
        @return Result or None in case of an error
        """
        if search_id is None:
            logger.error("Search ID not specified!")
            return None
        endpoint = "/ariel/searches/%s/results" % search_id
        resp_json = self.api.request_get(endpoint)
        if resp_json is None:
            logger.error("Unable to get results for query with ID " + str(search_id) + " !")
            return None
        return resp_json
