156 lines
7.5 KiB
Python
156 lines
7.5 KiB
Python
from __future__ import annotations
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import numpy as np
|
|
import pandas as pd
|
|
import sys as _sys
|
|
from pathlib import Path as _Path
|
|
_sys.path.insert(0, str(_Path(__file__).resolve().parents[1]))
|
|
from common.data_contract import PACKET_FEATURE_NAMES, PACKET_CONTINUOUS_CHANNEL_IDX, PACKET_BINARY_CHANNEL_IDX, fit_packet_stats as _fit_packet_stats, zscore as _zscore
|
|
import importlib.util as _ilu
|
|
_UDATA_NAME = 'unified_cfm_data'
|
|
if _UDATA_NAME not in _sys.modules:
|
|
_udata_spec = _ilu.spec_from_file_location(_UDATA_NAME, _Path(__file__).resolve().parents[1] / 'Unified_CFM' / 'data.py')
|
|
_udata = _ilu.module_from_spec(_udata_spec)
|
|
_sys.modules[_UDATA_NAME] = _udata
|
|
_udata_spec.loader.exec_module(_udata)
|
|
else:
|
|
_udata = _sys.modules[_UDATA_NAME]
|
|
DEFAULT_FLOW_META_COLUMNS = _udata.DEFAULT_FLOW_META_COLUMNS
|
|
_read_aligned_flow_features = _udata._read_aligned_flow_features
|
|
_preprocess_flow = _udata._preprocess_flow
|
|
|
|
@dataclass
|
|
class MixedData:
|
|
train_cont: np.ndarray
|
|
val_cont: np.ndarray
|
|
attack_cont: np.ndarray
|
|
train_disc: np.ndarray
|
|
val_disc: np.ndarray
|
|
attack_disc: np.ndarray
|
|
train_flow: np.ndarray
|
|
val_flow: np.ndarray
|
|
attack_flow: np.ndarray
|
|
train_len: np.ndarray
|
|
val_len: np.ndarray
|
|
attack_len: np.ndarray
|
|
attack_labels: np.ndarray
|
|
cont_mean: np.ndarray
|
|
cont_std: np.ndarray
|
|
flow_mean: np.ndarray
|
|
flow_std: np.ndarray
|
|
flow_feature_names: tuple[str, ...]
|
|
packet_feature_names: tuple[str, ...] = PACKET_FEATURE_NAMES
|
|
|
|
@property
|
|
def T(self) -> int:
|
|
return int(self.train_cont.shape[1])
|
|
|
|
@property
|
|
def n_cont(self) -> int:
|
|
return int(self.train_cont.shape[2])
|
|
|
|
@property
|
|
def n_disc(self) -> int:
|
|
return int(self.train_disc.shape[2])
|
|
|
|
@property
|
|
def flow_dim(self) -> int:
|
|
return int(self.train_flow.shape[1])
|
|
|
|
def _zscore_cont(train_x: np.ndarray, val_x: np.ndarray, attack_x: np.ndarray, train_l: np.ndarray, val_l: np.ndarray, attack_l: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
(mean, std) = _fit_packet_stats(train_x, train_l)
|
|
|
|
def prep(x: np.ndarray, l: np.ndarray) -> np.ndarray:
|
|
z = _zscore(x, mean, std)
|
|
T = x.shape[1]
|
|
m = np.arange(T)[None, :] < l[:, None]
|
|
return (z * m[:, :, None]).astype(np.float32)
|
|
return (prep(train_x, train_l), prep(val_x, val_l), prep(attack_x, attack_l), mean, std)
|
|
|
|
def load_mixed_data(*, packets_npz: Path | None=None, source_store: Path | None=None, flows_parquet: Path, flow_features_path: Path, flow_feature_columns: Optional[list[str]]=None, flow_features_align: str='auto', T: int=64, split_seed: int=42, train_ratio: float=0.8, benign_label: str='normal', min_len: int=2, attack_cap: int | None=None, val_cap: int | None=None) -> MixedData:
|
|
if (packets_npz is None) == (source_store is None):
|
|
raise ValueError('pass exactly one of packets_npz or source_store')
|
|
flows_parquet = Path(flows_parquet)
|
|
print(f'[data] flows={flows_parquet} packets={(packets_npz if packets_npz else source_store)}')
|
|
flow_cols = ['flow_id', 'label', 'src_ip', 'src_port', 'dst_ip', 'dst_port', 'protocol']
|
|
flows = pd.read_parquet(flows_parquet, columns=flow_cols)
|
|
labels_full = flows['label'].to_numpy().astype(str)
|
|
flow_id = flows['flow_id'].to_numpy()
|
|
tokens_full: np.ndarray | None = None
|
|
store = None
|
|
if packets_npz is not None:
|
|
pz = np.load(Path(packets_npz))
|
|
tokens_full = pz['packet_tokens'].astype(np.float32)
|
|
lens_full = pz['packet_lengths'].astype(np.int32)
|
|
if T > tokens_full.shape[1]:
|
|
raise ValueError(f'requested T={T} > stored {tokens_full.shape[1]}')
|
|
tokens_full = tokens_full[:, :T].copy()
|
|
lens_full = np.minimum(lens_full, T).astype(np.int32)
|
|
if 'flow_id' in pz.files and (not np.array_equal(pz['flow_id'], flow_id)):
|
|
raise ValueError('packets_npz / flows_parquet not row-aligned')
|
|
else:
|
|
from common.packet_store import PacketShardStore
|
|
store = PacketShardStore.open(Path(source_store))
|
|
store_id = store.read_flows(columns=['flow_id'])['flow_id'].to_numpy()
|
|
if not np.array_equal(store_id, flow_id):
|
|
raise ValueError('source_store / flows_parquet not row-aligned')
|
|
lens_full = np.minimum(store.manifest['packet_length'].to_numpy(dtype=np.int32), T)
|
|
(flow_features, flow_names) = _read_aligned_flow_features(Path(flow_features_path), flows, feature_columns=flow_feature_columns, align=flow_features_align)
|
|
keep = lens_full >= min_len
|
|
labels = labels_full[keep]
|
|
flow_features = flow_features[keep]
|
|
lens = lens_full[keep]
|
|
global_idx = np.flatnonzero(keep).astype(np.int64)
|
|
materialized = tokens_full[keep] if tokens_full is not None else None
|
|
print(f'[data] kept {keep.sum():,} of {len(keep):,} (min_len={min_len})')
|
|
benign = np.where(labels == benign_label)[0]
|
|
attack = np.where(labels != benign_label)[0]
|
|
rng = np.random.default_rng(split_seed)
|
|
rng.shuffle(benign)
|
|
n_train = int(len(benign) * train_ratio)
|
|
train_local = benign[:n_train]
|
|
val_local = benign[n_train:]
|
|
if val_cap is not None and len(val_local) > val_cap:
|
|
val_local = np.sort(rng.choice(val_local, size=val_cap, replace=False))
|
|
if attack_cap is not None and len(attack) > attack_cap:
|
|
attack = np.sort(rng.choice(attack, size=attack_cap, replace=False))
|
|
print(f'[data] train={len(train_local):,} val={len(val_local):,} attack={len(attack):,}')
|
|
|
|
def _materialize(idx_local: np.ndarray) -> np.ndarray:
|
|
if materialized is not None:
|
|
return materialized[idx_local].astype(np.float32, copy=False)
|
|
assert store is not None
|
|
g = global_idx[idx_local]
|
|
(tok, _) = store.read_packets(g.astype(np.int64), T=T)
|
|
return tok.astype(np.float32, copy=False)
|
|
tr_p = _materialize(train_local)
|
|
va_p = _materialize(val_local)
|
|
at_p = _materialize(attack)
|
|
tr_l = lens[train_local]
|
|
va_l = lens[val_local]
|
|
at_l = lens[attack]
|
|
tr_f = flow_features[train_local]
|
|
va_f = flow_features[val_local]
|
|
at_f = flow_features[attack]
|
|
cont_idx = list(PACKET_CONTINUOUS_CHANNEL_IDX)
|
|
disc_idx = list(PACKET_BINARY_CHANNEL_IDX)
|
|
tr_cont = tr_p[..., cont_idx]
|
|
va_cont = va_p[..., cont_idx]
|
|
at_cont = at_p[..., cont_idx]
|
|
tr_disc = tr_p[..., disc_idx].astype(np.int8)
|
|
va_disc = va_p[..., disc_idx].astype(np.int8)
|
|
at_disc = at_p[..., disc_idx].astype(np.int8)
|
|
(tr_cont, va_cont, at_cont, c_mean, c_std) = _zscore_cont(tr_cont, va_cont, at_cont, tr_l, va_l, at_l)
|
|
(tr_flow, va_flow, at_flow, f_mean, f_std) = _preprocess_flow(tr_f, va_f, at_f)
|
|
return MixedData(train_cont=tr_cont, val_cont=va_cont, attack_cont=at_cont, train_disc=tr_disc, val_disc=va_disc, attack_disc=at_disc, train_flow=tr_flow, val_flow=va_flow, attack_flow=at_flow, train_len=tr_l, val_len=va_l, attack_len=at_l, attack_labels=labels[attack], cont_mean=c_mean, cont_std=c_std, flow_mean=f_mean, flow_std=f_std, flow_feature_names=tuple(flow_names))
|
|
|
|
def subsample_train(data: MixedData, n: int, seed: int) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
if n <= 0 or n >= len(data.train_cont):
|
|
return (data.train_flow, data.train_cont, data.train_disc, data.train_len)
|
|
rng = np.random.default_rng(seed)
|
|
idx = rng.choice(len(data.train_cont), n, replace=False)
|
|
idx.sort()
|
|
return (data.train_flow[idx], data.train_cont[idx], data.train_disc[idx], data.train_len[idx])
|