Version_4/pipeline/graph_engine.py

import torch
from typing import Dict, List, Tuple

class TopologicalGraphEngine:
    """Translates incoming raw telemetry logs into PyTorch Geometric HeteroData formats."""
    def __init__(self):
        self.ip_map, self.domain_map, self.asn_map = {}, {}, {}

    def extract_and_build(self, telemetries: List[dict]) -> Tuple[Dict[str, torch.Tensor], Dict[Tuple, torch.Tensor]]:
        # Map dynamic components to index pointers
        for log in telemetries:
            if log.get('ip') and log['ip'] not in self.ip_map:
                self.ip_map[log['ip']] = len(self.ip_map)
            if log.get('domain') and log['domain'] not in self.domain_map:
                self.domain_map[log['domain']] = len(self.domain_map)
            if log.get('asn') and log['asn'] not in self.asn_map:
                self.asn_map[log['asn']] = len(self.asn_map)

        # High-dimensional hidden state initializers (using structural placeholders)
        # Note: These sizes must match in_channels_dict in app.py and train.py
        x_dict = {
            'ip': torch.randn((max(1, len(self.ip_map)), 16)),
            'domain': torch.randn((max(1, len(self.domain_map)), 32)),
            'asn': torch.randn((max(1, len(self.asn_map)), 8)),
            'cert': torch.randn((1, 16))
        }

        domain_to_ip = [[], []]
        ip_to_asn = [[], []]

        for log in telemetries:
            if log.get('domain') and log.get('ip'):
                domain_to_ip[0].append(self.domain_map[log['domain']])
                domain_to_ip[1].append(self.ip_map[log['ip']])
            if log.get('ip') and log.get('asn'):
                ip_to_asn[0].append(self.ip_map[log['ip']])
                ip_to_asn[1].append(self.asn_map[log['asn']])

        edge_index_dict = {
            ('domain', 'resolves_to', 'ip'): torch.tensor(domain_to_ip, dtype=torch.long),
            ('ip', 'hosted_on', 'asn'): torch.tensor(ip_to_asn, dtype=torch.long),
            ('domain', 'secured_by', 'cert'): torch.empty((2, 0), dtype=torch.long),
            ('cert', 'issued_to', 'ip'): torch.empty((2, 0), dtype=torch.long),
            ('domain', 'redirects_to', 'domain'): torch.empty((2, 0), dtype=torch.long)
        }

        return x_dict, edge_index_dict