neptune/core/lib/nicktrace.py

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