graph/relationships.py

from __future__ import annotations

from typing import Dict, List

from shield.utils.hashing import hash_entity


def build_relationship_edges(account_hash: str, entities: Dict[str, str]) -> List[Dict[str, str]]:
    edge_names = {
        "device": "USES_DEVICE",
        "ip": "USES_IP",
        "email": "OWNS_EMAIL",
        "phone": "OWNS_PHONE",
        "username": "USES_USERNAME",
    }
    edges: List[Dict[str, str]] = []
    for entity_type, value in entities.items():
        edge_type = edge_names.get(entity_type)
        if not edge_type:
            continue
        edges.append({
            "from": account_hash,
            "to": hash_entity(entity_type, value),
            "entity_type": entity_type,
            "edge_type": edge_type,
        })
    return edges