node_process.py

# Recovered from chat history
import os
import json
import re

from sqlalchemy.sql.sqltypes import NULLTYPE
### NEW: Global feature utilities
import numpy as np
from collections import defaultdict, Counter
from math import log


# Define multiple input folders and their corresponding output folders
input_output_folders = [
    # {'input': 'dataset/train_dataset/hyperGraph-73normalTX/nodes',
    #  'output': 'dataset/train_dataset/hyperGraph-73normalTX/new_nodes'},
    # {'input': 'dataset/train_dataset/hyperGraph-90attacks/nodes',
    #  'output': 'dataset/train_dataset/hyperGraph-90attacks/new_nodes'},
    # {'input': 'dataset/test_dataset/hyperGraph-73normalTX/nodes',
    #  'output': 'dataset/test_dataset/hyperGraph-73normalTX/new_nodes'},
    # {'input': 'dataset/test_dataset/hyperGraph-90attacks/nodes',
    #  'output': 'dataset/test_dataset/hyperGraph-90attacks/new_nodes'},
    # {'input': 'dataset/hyperGraph_only_in_REKT_1102/nodes',
    #  'output': 'dataset/hyperGraph_only_in_REKT_1102/new_nodes'},
    # {'input': 'dataset/hyperGraph_REKT_1102/nodes',
    #  'output': 'dataset/hyperGraph_REKT_1102/new_nodes'},
    # {'input': 'dataset/test1_dataset/hyperGraph-100normalTX/nodes',
    #  'output': 'dataset/test1_dataset/hyperGraph-100normalTX/new_nodes'},
    # {'input': 'dataset/test1_dataset/hyperGraph-100attacks/nodes',
    #  'output': 'dataset/test1_dataset/hyperGraph-100attacks/new_nodes'},
    # {'input': 'dataset/hyperGraph_only_in_REKT_1107/nodes',
    #  'output': 'dataset/hyperGraph_only_in_REKT_1107/new_nodes'},
    # {'input': 'dataset/hyperGraph_REKT_1107/nodes',
    #  'output': 'dataset/hyperGraph_REKT_1107/new_nodes'},

    # Modified folders
    # {'input': 'dataset/train_dataset/hyperGraph-attacks/nodes',
    #  'output': 'dataset/train_dataset/hyperGraph-attacks/new_nodes'},
    # {'input': 'dataset/train_dataset/hyperGraph-normalTX/nodes',# nodes_1710,nodes_1710+1000,nodes_1710+D3,nodes_1710+few
    #  'output': 'dataset/train_dataset/hyperGraph-normalTX/new_nodes'},
    # {'input': 'dataset/test1_dataset/hyperGraph-attacks/nodes',
    #  'output': 'dataset/test1_dataset/hyperGraph-attacks/new_nodes'},
    {'input': 'dataset/test1_dataset/hyperGraph-normalTX/nodes',# nodes_7000,nodes_2years
     'output': 'dataset/test1_dataset/hyperGraph-normalTX/new_nodes'},
    # {'input': 'dataset/test1_dataset/hyperGraph-normalTX/nodes_hard_neg_remaining',# nodes_7000,nodes_2years,nodes_hard_neg_remaining
    #  'output': 'dataset/test1_dataset/hyperGraph-normalTX/new_nodes'},

    # {'input': 'dataset/train_dataset/hyperGraph-normalTX/nodes_96',
    #  'output': 'dataset/train_dataset/hyperGraph-normalTX/new_nodes'},

    # {'input': 'dataset/hyperGraph_D3/nodes',
    #  'output': 'dataset/hyperGraph_D3/new_nodes'},
]
train_attack_dir="dataset/train_dataset/hyperGraph-attacks/new_nodes"
train_normal_dir="dataset/train_dataset/hyperGraph-normalTX/new_nodes"
test_attack_dir="dataset/test1_dataset/hyperGraph-attacks/new_nodes"
test_normal_dir="dataset/test1_dataset/hyperGraph-normalTX/new_nodes"
# Global processing input/output folders
# train_global_feature_input_folders = ["dataset/train_dataset/hyperGraph-normalTX/nodes_1710", "dataset/train_dataset/hyperGraph-attacks/nodes"]
# train_global_feature_input_folders = ["dataset/train_dataset/hyperGraph-normalTX/nodes_1710+few", "dataset/train_dataset/hyperGraph-attacks/nodes"]
train_global_feature_input_folders = ["dataset/train_dataset/hyperGraph-normalTX/nodes", "dataset/train_dataset/hyperGraph-attacks/nodes"]
train_global_feature_output_path = "dataset/train_dataset/processed_data/global_features.json"

# test1_global_feature_input_folders = ["dataset/test1_dataset/hyperGraph-normalTX/nodes_7000", "dataset/test1_dataset/hyperGraph-attacks/nodes"]
# test1_global_feature_input_folders = ["dataset/test1_dataset/hyperGraph-normalTX/nodes_2years", "dataset/test1_dataset/hyperGraph-attacks/nodes"]
test1_global_feature_input_folders = ["dataset/test1_dataset/hyperGraph-normalTX/nodes", "dataset/test1_dataset/hyperGraph-attacks/nodes"]
# test1_global_feature_input_folders = ["dataset/test1_dataset/hyperGraph-normalTX/nodes_hard_neg_remaining", "dataset/test1_dataset/hyperGraph-attacks/nodes"]
test1_global_feature_output_path = "dataset/test1_dataset/processed_data/global_features.json"

# Process addresses
def process_node_json(input_output_folders):
    # Address pattern: starts with "0x" followed by 40 hex characters
    address_pattern = re.compile(r"^0x[0-9a-fA-F]{40}$")
    # Specific address labels (use lowercase for comparison)
    ADDRESS_LABELS = {
        "0x1838de7d4e4e42c8eb7b204a91e28e9fad14f536".lower(): "ProtocolFeeRecipient",
        "0x66e1f353d8c19de6107981bc017df85df60e8199".lower(): "UNI-V2",
        "0xa373f8a0c75e2315444079c3745bcc97653b0bd5".lower(): "ChainTools: Yield Vault",
    }

    for folder_pair in input_output_folders:
        input_folder = folder_pair['input']
        output_folder = folder_pair['output']
        os.makedirs(output_folder, exist_ok=True)

        for file_name in os.listdir(input_folder):
            if not file_name.endswith('.json'):
                continue

            # Read JSON
            input_file_path = os.path.join(input_folder, file_name)
            with open(input_file_path, 'r') as file:
                data = json.load(file)

            # ==== 1) First build Token/address mappings (preserve existing logic) ====
            token_mapping, from_mapping, to_mapping = {}, {}, {}
            token_counter = from_counter = to_counter = 1

            for node in data.get('Nodes', []):
                # TokenName
                token_name = node.get('TokenName', '')
                if token_name not in token_mapping:
                    token_mapping[token_name] = f"Token{token_counter}"
                    token_counter += 1
                node['TokenName'] = token_mapping[token_name]

                # From
                from_name = node.get("From", "")
                if address_pattern.match(from_name):
                    addr_lc = from_name.lower()
                    if addr_lc in ADDRESS_LABELS:
                        # Matched a specific address -> use fixed label
                        node['From'] = ADDRESS_LABELS[addr_lc]
                        # Also record in the mapping table to ensure consistency for subsequent same addresses
                        from_mapping[from_name] = ADDRESS_LABELS[addr_lc]
                    else:
                        # Regular address -> use existing From1/From2... incremental mapping
                        if from_name not in from_mapping:
                            from_mapping[from_name] = f"From{from_counter}"
                            from_counter += 1
                        node['From'] = from_mapping[from_name]
                else:
                    from_mapping[from_name] = from_name
                    node['From'] = from_mapping[from_name]

                # To
                to_name = node.get("To", "")
                if address_pattern.match(to_name):
                    addr_lc = to_name.lower()
                    if addr_lc in ADDRESS_LABELS:
                        node['To'] = ADDRESS_LABELS[addr_lc]
                        to_mapping[to_name] = ADDRESS_LABELS[addr_lc]
                    else:
                        if to_name not in to_mapping:
                            to_mapping[to_name] = f"To{to_counter}"
                            to_counter += 1
                        node['To'] = to_mapping[to_name]
                else:
                    to_mapping[to_name] = to_name
                    node['To'] = to_mapping[to_name]

            # ==== 2) Compute relative amounts and write back (aligned with compute_features_from_nodes) ====
            nodes = data.get('Nodes', [])
            n_nodes = len(nodes)

            # Collect amounts (prefer 'Amount', otherwise use 'amount_raw', default to 0)
            amounts = []
            amount_keys = []  # Record the field name used for each node, for writing back later
            for nd in nodes:
                if 'Amount' in nd:
                    amt = safe_float(nd.get('Amount', 0))
                    amount_keys.append('Amount')
                elif 'amount_raw' in nd:
                    amt = safe_float(nd.get('amount_raw', 0))
                    amount_keys.append('amount_raw')
                else:
                    amt = 0.0
                    amount_keys.append('Amount')  # If no field exists, write back to 'Amount' later
                # Clip to non-negative
                amounts.append(max(0.0, amt))

            # log1p -> min-max (within file)
            if n_nodes == 0:
                a_rel = np.array([], dtype=np.float64)
            else:
                a_log = np.log1p(np.asarray(amounts, dtype=np.float64))
                if n_nodes > 1:
                    a_min = float(np.min(a_log))
                    a_max = float(np.max(a_log))
                    denom = max(1e-9, a_max - a_min)
                    a_rel = (a_log - a_min) / denom
                else:
                    a_rel = np.array([0.0], dtype=np.float64)

            # Write relative amounts back to original field (or default to 'Amount')
            for i, nd in enumerate(nodes):
                key = amount_keys[i]
                nd[key] = float(a_rel[i]) if a_rel.size else 0.0

            # ==== 3) Save ====
            output_file_path = os.path.join(output_folder, file_name)
            with open(output_file_path, 'w') as output_file:
                json.dump(data, output_file, indent=4, ensure_ascii=False)

    print("nodes have been processed into new_nodes (amounts converted to relative values).")
    return 0


G_KEYS = [
    "n_nodes","n_addrs","n_tokens",
    "avg_amount_rel","max_amount_rel","std_amount_rel",
    "in_deg_max","in_deg_mean","out_deg_max","out_deg_mean",
    "amount_hhi","amount_gini","cycle_len_min",
    "token_entropy","addr_entropy"
]

# ---------- Compute global features ----------

def _safe_str(x):
    return str(x) if x is not None else ""

def safe_float(x):
    try:
        if x is None: return 0.0
        if isinstance(x, (int, float)): return float(x)
        s = str(x).strip()
        if s.lower().startswith("0x"):  # Hexadecimal number
            return float(int(s, 16))
        return float(s.replace(",", ""))  # Remove thousands-separator commas
    except Exception:
        return 0.0

def _entropy_from_counter(cn: Counter) -> float:
    if not cn: return 0.0
    arr = np.array(list(cn.values()), dtype=np.float64)
    p = arr / np.sum(arr)
    return float(-np.sum(p * np.log(p + 1e-12)))

def _gini_from_weights(w) -> float:
    w = np.asarray([x for x in w if x > 0], dtype=np.float64)
    if w.size == 0: return 0.0
    w = np.sort(w); n = w.size
    idx = np.arange(1, n+1)
    g = 1 - 2 * np.sum((n - idx + 0.5) * w) / (n * np.sum(w))
    return float(max(0.0, min(1.0, g)))

def _shortest_directed_cycle_len(edges, max_depth=6) -> int:
    adj = defaultdict(list)
    nodes = set()
    for u, v in edges:
        adj[u].append(v)
        nodes.add(u); nodes.add(v)
    best = None
    for s in nodes:
        q = [(s, 0)]
        seen = {s: 0}
        while q:
            cur, d = q.pop(0)
            if d >= max_depth: continue
            for nb in adj[cur]:
                if nb == s and d+1 > 1:
                    best = d+1 if best is None else min(best, d+1)
                if nb not in seen or seen[nb] > d+1:
                    seen[nb] = d+1
                    q.append((nb, d+1))
    return 0 if best is None else best

def compute_features_from_nodes(nodes_list: list) -> dict:
    """Compute global features directly from data['Nodes']"""
    n_nodes = len(nodes_list)
    addrs, tokens, amounts, edges = set(), set(), [], []
    in_deg, out_deg, token_cnt, addr_cnt = Counter(), Counter(), Counter(), Counter()

    for nd in nodes_list:
        u = _safe_str(nd.get("From", nd.get("from")))
        v = _safe_str(nd.get("To", nd.get("to")))
        tok = _safe_str(nd.get("TokenName", nd.get("token")))
        amt = safe_float(nd.get("Amount", nd.get("amount_raw")))

        amounts.append(max(amt, 0.0))
        if u: addrs.add(u)
        if v: addrs.add(v)
        if tok: tokens.add(tok)
        if u and v:
            edges.append((u, v))
            out_deg[u] += 1; in_deg[v] += 1
        if tok: token_cnt[tok] += 1
        if u: addr_cnt[u] += 1
        if v: addr_cnt[v] += 1

    # Within-transaction relative amounts: log1p -> min-max
    a_log = np.log1p(np.asarray(amounts, dtype=np.float64)) if n_nodes > 0 else np.array([0.0])
    if n_nodes > 1:
        a_min, a_max = float(np.min(a_log)), float(np.max(a_log))
        denom = max(1e-9, a_max - a_min)
        a_rel = (a_log - a_min) / denom
    else:
        a_rel = np.array([0.0])

    s = float(np.sum(a_rel))
    w = (a_rel / s).tolist() if s > 0 else [1.0]
    hhi = float(np.sum(np.square(w)))
    gini = _gini_from_weights(w)

    in_vals, out_vals = list(in_deg.values()), list(out_deg.values())
    feats = {
        "n_nodes": n_nodes,
        "n_addrs": len(addrs),
        "n_tokens": len(tokens),
        "avg_amount_rel": float(np.mean(a_rel)),
        "max_amount_rel": float(np.max(a_rel)) if a_rel.size else 0.0,
        "std_amount_rel": float(np.std(a_rel)),
        "in_deg_max": max(in_vals) if in_vals else 0,
        "in_deg_mean": float(np.mean(in_vals)) if in_vals else 0.0,
        "out_deg_max": max(out_vals) if out_vals else 0,
        "out_deg_mean": float(np.mean(out_vals)) if out_vals else 0.0,
        "amount_hhi": hhi,
        "amount_gini": gini,
        "cycle_len_min": _shortest_directed_cycle_len(edges, max_depth=6),
        "token_entropy": _entropy_from_counter(token_cnt),
        "addr_entropy": _entropy_from_counter(addr_cnt)
    }
    return feats
def compute_global_features(input_folders, save_path="global_features.json"):
    """
    Iterate over each directory in input_folders, read *.json files,
    extract TxHash as key, compute global feature dict as value.
    Return {tx_hash: feats_dict} and save to save_path file.
    """
    all_feats = {}

    for input_folder in input_folders:
        for file_name in os.listdir(input_folder):
            if not file_name.endswith('.json'):
                continue
            path = os.path.join(input_folder, file_name)
            with open(path, 'r', encoding='utf-8') as f:
                data = json.load(f)

            tx_hash = data.get("TxHash")
            if not tx_hash:  # Skip if no TxHash
                continue

            feats = compute_features_from_nodes(data.get("Nodes", []))
            all_feats[tx_hash] = feats

    # Save to file
    with open(save_path, "w", encoding="utf-8") as f:
        json.dump(all_feats, f, ensure_ascii=False, indent=4)

    print(f"Processed {len(all_feats)} transactions, global features saved to {save_path}")
    return all_feats


import os
from typing import Dict, List


# Dataset deduplication
def resolve_dataset_duplicates(
        train_attack_dir: str,
        train_normal_dir: str,
        test_attack_dir: str,
        test_normal_dir: str,
        extensions=(".json",),
) -> Dict:
    """
    Check for duplicate filenames across four data directories (train/test x attack/normal) and deduplicate.
    Retention priority: train-attack > train-normal > test-attack > test-normal.
    When duplicates are found, keep the highest-priority file, delete the rest, and print conflict info.

    Args:
        train_attack_dir: Training set - attack transaction folder
        train_normal_dir: Training set - normal transaction folder
        test_attack_dir:  Test set - attack transaction folder
        test_normal_dir:  Test set - normal transaction folder
        extensions:       Tuple of file extensions considered as transaction samples, default (".json",)

    Returns:
        {
            "deleted": List[{"file": filename, "kept_in": retention directory, "deleted_from": deletion directory}],
            "remaining_counts": {
                "train_attack": remaining json file count,
                "train_normal": remaining json file count,
                "test_attack":  remaining json file count,
                "test_normal":  remaining json file count
            }
        }
    """
    # Directory labels and priority (lower number = higher priority)
    folders = [
        ("train_attack", train_attack_dir, 0),
        ("train_normal", train_normal_dir, 1),
        ("test_attack", test_attack_dir, 2),
        ("test_normal", test_normal_dir, 3),
    ]
    label_to_path = {lbl: p for lbl, p, _ in folders}
    prio = {lbl: rank for lbl, _, rank in folders}

    def is_valid_sample(fn: str) -> bool:
        return any(fn.endswith(ext) for ext in extensions)

    def list_names(folder: str) -> set:
        if not folder or not os.path.isdir(folder):
            return set()
        return {fn for fn in os.listdir(folder) if is_valid_sample(fn)}

    # Collect which directories each filename appears in
    name_to_locs: Dict[str, List[str]] = {}
    for lbl, path, _ in folders:
        for fn in list_names(path):
            name_to_locs.setdefault(fn, []).append(lbl)

    deletions: List[Dict] = []

    # For filenames appearing >= 2 times, keep one by priority and delete the rest
    for fn, locs in sorted(name_to_locs.items()):
        if len(locs) <= 1:
            continue
        # Select the location to keep (lowest priority number = highest priority)
        keep_lbl = min(locs, key=lambda x: prio[x])
        keep_path = label_to_path[keep_lbl]

        # Delete from remaining locations and print conflict info
        for lose_lbl in sorted(locs, key=lambda x: prio[x]):
            if lose_lbl == keep_lbl:
                continue
            lose_path = label_to_path[lose_lbl]
            lose_file = os.path.join(lose_path, fn)

            # Print per-conflict error message
            print(
                f"[Conflict] Filename: {fn} | Conflicting dirs: {keep_lbl} <-> {lose_lbl} | "
                f"Keep: {keep_lbl} ({os.path.join(keep_path, fn)}) | Delete: {lose_lbl} ({lose_file})"
            )
            # Perform deletion
            try:
                if os.path.exists(lose_file):
                    os.remove(lose_file)
                    deletions.append({
                        "file": fn,
                        "kept_in": keep_lbl,
                        "deleted_from": lose_lbl
                    })
            except Exception as e:
                print(f"[WARN] Deletion failed: {lose_file} | Error: {e}")

    # Count remaining files in the four folders
    remaining_counts = {
        "train_attack": len(list_names(train_attack_dir)),
        "train_normal": len(list_names(train_normal_dir)),
        "test_attack": len(list_names(test_attack_dir)),
        "test_normal": len(list_names(test_normal_dir)),
    }

    # Print summary
    if deletions:
        print("\n===== Deduplication Summary =====")
        print(f"Deleted {len(deletions)} duplicate file instances (possibly from multiple directories).")
        for d in deletions:
            print(f" - {d['file']}: Kept in [{d['kept_in']}], deleted from [{d['deleted_from']}]")
    else:
        print("No cross-directory duplicate files found.")

    print("\n===== Remaining File Counts (only extension-matched sample files) =====")
    for k, v in remaining_counts.items():
        print(f"{k}: {v}")

    return {"deleted": deletions, "remaining_counts": remaining_counts}


if __name__ == "__main__":
    # Process addresses in data
    process_node_json(input_output_folders)

    # Dataset deduplication; for safety, operates on new_nodes folders (output of process_node_json) rather than original nodes folders
    # Deduplication retention priority: train-attack > train-normal > test-attack > test-normal
    resolve_dataset_duplicates(train_attack_dir,train_normal_dir,test_attack_dir,test_normal_dir)

    # Compute global features for training and testing
    # compute_global_features(train_global_feature_input_folders,train_global_feature_output_path)
    compute_global_features(test1_global_feature_input_folders,test1_global_feature_output_path)