neptune/core/lib/nicktrace.py

93 lines
3.0 KiB
Python

from math import ceil
from django.conf import settings
from numpy import array_split
from core.lib.opensearch import client, run_main_query
def construct_query(net, nicks):
# Construct the query
query_nicks = [{"match": {"nick": x}} for x in nicks]
query_users = [{"match": {"user": x}} for x in nicks]
query_should = query_nicks + query_users
# Get the initial query
query = {
"size": settings.NICKTRACE_QUERY_SIZE,
"query": {
"bool": {
"must": [
{"match": {"net": net}},
{"match": {"type": "nick"}},
{
"bool": {
"should": query_should,
}
},
]
}
},
}
return query
def get_nicks(request, net, nicks, iter=True):
"""
Get all related nicknames of the given nickname by tracking nickname changes.
"""
# Split query into chunks
split_nicks = array_split(
nicks, ceil(len(nicks) / settings.NICKTRACE_MAX_CHUNK_SIZE)
)
nicks = [*nicks]
for nicks_chunked in split_nicks:
if len(nicks_chunked) == 0:
break
query = construct_query(net, nicks_chunked)
results = run_main_query(client, request.user, query, custom_query=True)
if "hits" in results.keys():
if "hits" in results["hits"]:
for item in results["hits"]["hits"]:
element = item["_source"]
element["id"] = item["_id"]
# Split the timestamp into date and time
ts = element["ts"]
ts_spl = ts.split("T")
date = ts_spl[0]
time = ts_spl[1]
element["date"] = date
element["time"] = time
if element["nick"] not in nicks:
nicks.append(element["nick"])
if element["user"] not in nicks:
nicks.append(element["user"])
# Run the search again, passing in all the users we found
# Nicknames we find from the repeated search
nicks_searched = []
if iter:
nicks_l2 = []
loop = 0
while loop < settings.NICKTRACE_MAX_ITERATIONS:
loop += 1
nicks_not_searched = [x for x in nicks if x not in nicks_searched]
if not nicks_not_searched:
break
nicks_l2 = get_nicks(request, net, nicks, False)
# Add all the nicks we just searched for to the list
for x in nicks_not_searched:
if x not in nicks_not_searched:
nicks_searched.append(x)
# If all of the nicks we received now, we already know about
if set(nicks_l2).issubset(set(nicks)):
break
for x in nicks_l2:
if x not in nicks:
nicks.append(x)
return nicks