import logging from datetime import datetime from pprint import pprint import requests from django.conf import settings from core.db import StorageBackend, add_defaults, dedup_list from core.db.processing import annotate_results, parse_results logger = logging.getLogger(__name__) class ManticoreBackend(StorageBackend): def __init__(self): super().__init__("Manticore") def initialise(self, **kwargs): """ Initialise the Manticore client """ pass # we use requests def construct_query(self, query, size, index, blank=False): """ Accept some query parameters and construct an OpenSearch query. """ if not size: size = 5 query_base = { "index": index, "limit": size, "query": {"bool": {"must": []}}, } query_string = { "query_string": query, } if not blank: query_base["query"]["bool"]["must"].append(query_string) return query_base def run_query(self, client, user, search_query): response = requests.post( f"{settings.MANTICORE_URL}/json/search", json=search_query ) return response def query_results( self, request, query_params, size=None, annotate=True, custom_query=False, reverse=False, dedup=False, dedup_fields=None, tags=None, ): query = None message = None message_class = None add_bool = [] add_top = [] add_top_negative = [] sort = None query_created = False source = None add_defaults(query_params) # Check size if request.user.is_anonymous: sizes = settings.MANTICORE_MAIN_SIZES_ANON else: sizes = settings.MANTICORE_MAIN_SIZES if not size: if "size" in query_params: size = query_params["size"] if size not in sizes: message = "Size is not permitted" message_class = "danger" return {"message": message, "class": message_class} size = int(size) else: size = 20 # Check index if "index" in query_params: index = query_params["index"] if index == "main": index = settings.MANTICORE_INDEX_MAIN else: if not request.user.has_perm(f"core.index_{index}"): message = "Not permitted to search by this index" message_class = "danger" return { "message": message, "class": message_class, } if index == "meta": index = settings.MANTICORE_INDEX_META elif index == "internal": index = settings.MANTICORE_INDEX_INT else: message = "Index is not valid." message_class = "danger" return { "message": message, "class": message_class, } else: index = settings.MANTICORE_INDEX_MAIN # Create the search query if "query" in query_params: query = query_params["query"] search_query = self.construct_query(query, size, index) query_created = True else: if custom_query: search_query = custom_query if tags: # Get a blank search query if not query_created: search_query = self.construct_query(None, size, index, blank=True) query_created = True for tagname, tagvalue in tags.items(): add_bool.append({tagname: tagvalue}) required_any = ["query_full", "query", "tags"] if not any([field in query_params.keys() for field in required_any]): if not custom_query: message = "Empty query!" message_class = "warning" return {"message": message, "class": message_class} # Check for a source if "source" in query_params: source = query_params["source"] if source in settings.SOURCES_RESTRICTED: if not request.user.has_perm("core.restricted_sources"): message = "Access denied" message_class = "danger" return {"message": message, "class": message_class} elif source not in settings.MAIN_SOURCES: message = "Invalid source" message_class = "danger" return {"message": message, "class": message_class} if source == "all": source = None # the next block will populate it if source: sources = [source] else: sources = list(settings.MAIN_SOURCES) if request.user.has_perm("core.restricted_sources"): for source_iter in settings.SOURCES_RESTRICTED: sources.append(source_iter) add_top_tmp = {"bool": {"should": []}} total_count = 0 for source_iter in sources: add_top_tmp["bool"]["should"].append({"equals": {"src": source_iter}}) total_count += 1 total_sources = len(settings.MAIN_SOURCES) + len(settings.SOURCES_RESTRICTED) if not total_count == total_sources: add_top.append(add_top_tmp) # Date/time range if set({"from_date", "to_date", "from_time", "to_time"}).issubset( query_params.keys() ): from_ts = f"{query_params['from_date']}T{query_params['from_time']}Z" to_ts = f"{query_params['to_date']}T{query_params['to_time']}Z" from_ts = datetime.strptime(from_ts, "%Y-%m-%dT%H:%MZ") to_ts = datetime.strptime(to_ts, "%Y-%m-%dT%H:%MZ") from_ts = int(from_ts.timestamp()) to_ts = int(to_ts.timestamp()) range_query = { "range": { "ts": { "gt": from_ts, "lt": to_ts, } } } add_top.append(range_query) # Sorting if "sorting" in query_params: sorting = query_params["sorting"] if sorting not in ("asc", "desc", "none"): message = "Invalid sort" message_class = "danger" return {"message": message, "class": message_class} if sorting in ("asc", "desc"): sort = [ { "ts": { "order": sorting, } } ] # Sentiment handling if "check_sentiment" in query_params: if "sentiment_method" not in query_params: message = "No sentiment method" message_class = "danger" return {"message": message, "class": message_class} if "sentiment" in query_params: sentiment = query_params["sentiment"] try: sentiment = float(sentiment) except ValueError: message = "Sentiment is not a float" message_class = "danger" return {"message": message, "class": message_class} sentiment_method = query_params["sentiment_method"] range_query_compare = {"range": {"sentiment": {}}} range_query_precise = { "match": { "sentiment": None, } } if sentiment_method == "below": range_query_compare["range"]["sentiment"]["lt"] = sentiment add_top.append(range_query_compare) elif sentiment_method == "above": range_query_compare["range"]["sentiment"]["gt"] = sentiment add_top.append(range_query_compare) elif sentiment_method == "exact": range_query_precise["match"]["sentiment"] = sentiment add_top.append(range_query_precise) elif sentiment_method == "nonzero": range_query_precise["match"]["sentiment"] = 0 add_top_negative.append(range_query_precise) if add_bool: # if "bool" not in search_query["query"]: # search_query["query"]["bool"] = {} # if "must" not in search_query["query"]["bool"]: # search_query["query"]["bool"] = {"must": []} for item in add_bool: search_query["query"]["bool"]["must"].append({"match": item}) if add_top: for item in add_top: search_query["query"]["bool"]["must"].append(item) if add_top_negative: for item in add_top_negative: if "must_not" in search_query["query"]["bool"]: search_query["query"]["bool"]["must_not"].append(item) else: search_query["query"]["bool"]["must_not"] = [item] if sort: search_query["sort"] = sort pprint(search_query) results = self.run_query( self.client, request.user, # passed through run_main_query to filter_blacklisted search_query, ) if not results: message = "Error running query" message_class = "danger" return {"message": message, "class": message_class} # results = results.to_dict() if "error" in results: message = results["error"] message_class = "danger" return {"message": message, "class": message_class} results_parsed = parse_results(results) if annotate: annotate_results(results_parsed) if "dedup" in query_params: if query_params["dedup"] == "on": dedup = True else: dedup = False else: dedup = False if reverse: results_parsed = results_parsed[::-1] if dedup: if not dedup_fields: dedup_fields = ["msg", "nick", "ident", "host", "net", "channel"] results_parsed = dedup_list(results_parsed, dedup_fields) context = { "object_list": results_parsed, "card": results["hits"]["total"], "took": results["took"], } if "cache" in results: context["cache"] = results["cache"] return context