331 lines
9.8 KiB
Python
331 lines
9.8 KiB
Python
import logging
|
|
from datetime import datetime
|
|
from pprint import pprint
|
|
import httpx
|
|
import orjson
|
|
|
|
import requests
|
|
from django.conf import settings
|
|
|
|
from core.db import StorageBackend, add_defaults, dedup_list
|
|
from core.db.processing import parse_results
|
|
from core.lib.parsing import (
|
|
QueryError,
|
|
parse_date_time,
|
|
parse_index,
|
|
parse_rule,
|
|
parse_sentiment,
|
|
parse_size,
|
|
parse_sort,
|
|
parse_source,
|
|
)
|
|
|
|
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
|
|
|
|
async def async_initialise(self, **kwargs):
|
|
"""
|
|
Initialise the Manticore client in async mode
|
|
"""
|
|
pass # we use requests
|
|
|
|
def delete_rule_entries(self, rule_id):
|
|
"""
|
|
Delete all entries for a given rule.
|
|
:param rule_id: The rule ID to delete.
|
|
"""
|
|
# TODO
|
|
|
|
def construct_query(self, query, size=None, blank=False, **kwargs):
|
|
"""
|
|
Accept some query parameters and construct an OpenSearch query.
|
|
"""
|
|
if not size:
|
|
size = 5
|
|
query_base = {
|
|
"index": kwargs.get("index"),
|
|
"limit": size,
|
|
"query": {"bool": {"must": []}},
|
|
}
|
|
print("BASE", query_base)
|
|
query_string = {
|
|
"query_string": query,
|
|
}
|
|
if not blank:
|
|
query_base["query"]["bool"]["must"].append(query_string)
|
|
return query_base
|
|
|
|
def parse(self, response, **kwargs):
|
|
parsed = parse_results(response, **kwargs)
|
|
return parsed
|
|
|
|
def run_query(self, user, search_query, **kwargs):
|
|
"""
|
|
Low level helper to run Manticore query.
|
|
"""
|
|
index = kwargs.get("index")
|
|
raw = kwargs.get("raw")
|
|
if search_query and not raw:
|
|
search_query["index"] = index
|
|
pprint(search_query)
|
|
|
|
|
|
path = kwargs.get("path", "json/search")
|
|
if raw:
|
|
response = requests.post(
|
|
f"{settings.MANTICORE_URL}/{path}", search_query
|
|
)
|
|
else:
|
|
response = requests.post(
|
|
f"{settings.MANTICORE_URL}/{path}", json=search_query
|
|
)
|
|
|
|
return orjson.loads(response.text)
|
|
|
|
async def async_run_query(self, user, search_query, **kwargs):
|
|
"""
|
|
Low level helper to run Manticore query asynchronously.
|
|
"""
|
|
index = kwargs.get("index")
|
|
search_query["index"] = index
|
|
pprint(search_query)
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.post(
|
|
f"{settings.MANTICORE_URL}/json/search", json=search_query
|
|
)
|
|
return orjson.loads(response.text)
|
|
|
|
async def async_store_matches(self, matches):
|
|
"""
|
|
Store a list of matches in Manticore.
|
|
:param index: The index to store the matches in.
|
|
:param matches: A list of matches to store.
|
|
"""
|
|
# TODO
|
|
|
|
def store_matches(self, matches):
|
|
"""
|
|
Store a list of matches in Manticore.
|
|
:param index: The index to store the matches in.
|
|
:param matches: A list of matches to store.
|
|
"""
|
|
# TODO
|
|
|
|
def prepare_schedule_query(self, rule_object):
|
|
"""
|
|
Helper to run a scheduled query with reduced functionality.
|
|
"""
|
|
# TODO
|
|
|
|
def schedule_query_results_test_sync(self, rule_object):
|
|
"""
|
|
Helper to run a scheduled query test with reduced functionality.
|
|
Sync version for running from Django forms.
|
|
Does not return results.
|
|
"""
|
|
# TODO
|
|
|
|
async def schedule_query_results(self, rule_object):
|
|
"""
|
|
Helper to run a scheduled query with reduced functionality and async.
|
|
"""
|
|
# TODO
|
|
|
|
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)
|
|
|
|
# Now, run the helpers for SIQTSRSS/ADR
|
|
# S - Size
|
|
# I - Index
|
|
# Q - Query
|
|
# T - Tags
|
|
# S - Source
|
|
# R - Ranges
|
|
# S - Sort
|
|
# S - Sentiment
|
|
# A - Annotate
|
|
# D - Dedup
|
|
# R - Reverse
|
|
|
|
# S - Size
|
|
if request.user.is_anonymous:
|
|
sizes = settings.MAIN_SIZES_ANON
|
|
else:
|
|
sizes = settings.MAIN_SIZES
|
|
if not size:
|
|
size = parse_size(query_params, sizes)
|
|
if isinstance(size, dict):
|
|
return size
|
|
|
|
rule_object = parse_rule(request.user, query_params)
|
|
if isinstance(rule_object, dict):
|
|
return rule_object
|
|
|
|
if rule_object is not None:
|
|
index = settings.INDEX_RULE_STORAGE
|
|
add_bool.append({"rule_id": str(rule_object.id)})
|
|
else:
|
|
# I - Index
|
|
index = parse_index(request.user, query_params)
|
|
if isinstance(index, dict):
|
|
return index
|
|
|
|
# Q/T - Query/Tags
|
|
search_query = self.parse_query(
|
|
query_params, tags, size, custom_query, add_bool
|
|
)
|
|
# Query should be a dict, so check if it contains message here
|
|
if "message" in search_query:
|
|
return search_query
|
|
|
|
# S - Sources
|
|
sources = parse_source(request.user, query_params)
|
|
if isinstance(sources, dict):
|
|
return sources
|
|
total_count = len(sources)
|
|
# Total is -1 due to the "all" source
|
|
total_sources = (
|
|
len(settings.MAIN_SOURCES) - 1 + len(settings.SOURCES_RESTRICTED)
|
|
)
|
|
|
|
# If the sources the user has access to are equal to all
|
|
# possible sources, then we don't need to add the source
|
|
# filter to the query.
|
|
if total_count != total_sources:
|
|
add_top_tmp = {"bool": {"should": []}}
|
|
for source_iter in sources:
|
|
add_top_tmp["bool"]["should"].append(
|
|
{"match_phrase": {"src": source_iter}}
|
|
)
|
|
if query_params["source"] != "all":
|
|
add_top.append(add_top_tmp)
|
|
|
|
# R - Ranges
|
|
# date_query = False
|
|
from_ts, to_ts = parse_date_time(query_params)
|
|
if from_ts:
|
|
range_query = {
|
|
"range": {
|
|
"ts": {
|
|
"gt": from_ts,
|
|
"lt": to_ts,
|
|
}
|
|
}
|
|
}
|
|
add_top.append(range_query)
|
|
|
|
# S - Sort
|
|
sort = parse_sort(query_params)
|
|
if isinstance(sort, dict):
|
|
return sort
|
|
|
|
if rule_object is not None:
|
|
field = "match_ts"
|
|
else:
|
|
field = "ts"
|
|
if sort:
|
|
# For Druid compatibility
|
|
sort_map = {"ascending": "asc", "descending": "desc"}
|
|
sorting = [
|
|
{
|
|
field: {
|
|
"order": sort_map[sort],
|
|
}
|
|
}
|
|
]
|
|
search_query["sort"] = sorting
|
|
|
|
# S - Sentiment
|
|
sentiment_r = parse_sentiment(query_params)
|
|
if isinstance(sentiment_r, dict):
|
|
return sentiment_r
|
|
if sentiment_r:
|
|
if rule_object is not None:
|
|
sentiment_index = "meta.aggs.avg_sentiment.value"
|
|
else:
|
|
sentiment_index = "sentiment"
|
|
sentiment_method, sentiment = sentiment_r
|
|
range_query_compare = {"range": {sentiment_index: {}}}
|
|
range_query_precise = {
|
|
"match": {
|
|
sentiment_index: None,
|
|
}
|
|
}
|
|
if sentiment_method == "below":
|
|
range_query_compare["range"][sentiment_index]["lt"] = sentiment
|
|
add_top.append(range_query_compare)
|
|
elif sentiment_method == "above":
|
|
range_query_compare["range"][sentiment_index]["gt"] = sentiment
|
|
add_top.append(range_query_compare)
|
|
elif sentiment_method == "exact":
|
|
range_query_precise["match"][sentiment_index] = sentiment
|
|
add_top.append(range_query_precise)
|
|
elif sentiment_method == "nonzero":
|
|
range_query_precise["match"][sentiment_index] = 0
|
|
add_top_negative.append(range_query_precise)
|
|
|
|
# Add in the additional information we already populated
|
|
self.add_bool(search_query, add_bool)
|
|
self.add_top(search_query, add_top)
|
|
self.add_top(search_query, add_top_negative, negative=True)
|
|
|
|
response = self.query(
|
|
request.user,
|
|
search_query,
|
|
index=index,
|
|
)
|
|
if not response:
|
|
message = "Error running query"
|
|
message_class = "danger"
|
|
return {"message": message, "class": message_class}
|
|
|
|
# results = results.to_dict()
|
|
if "error" in response:
|
|
message = response["error"]
|
|
message_class = "danger"
|
|
return {"message": message, "class": message_class}
|
|
if "message" in response:
|
|
return response
|
|
|
|
# A/D/R - Annotate/Dedup/Reverse
|
|
response["object_list"] = self.process_results(
|
|
response["object_list"],
|
|
annotate=annotate,
|
|
dedup=dedup,
|
|
dedup_fields=dedup_fields,
|
|
reverse=reverse,
|
|
)
|
|
|
|
context = response
|
|
return context
|