mirror of
https://github.com/velocitatem/PHANTOM.git
synced 2026-05-31 08:33:36 +00:00
84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
import os
|
|
from pydantic import BaseModel as Base
|
|
import json
|
|
|
|
class PayloadModel(Base):
|
|
sessionId: str
|
|
experimentId: str | None
|
|
eventName: str
|
|
page: str | None
|
|
productId: str | None
|
|
metadata: dict
|
|
storeMode: str
|
|
userAgent: str
|
|
ts: str
|
|
|
|
class ValueModel(Base):
|
|
payload: PayloadModel
|
|
encoding: str
|
|
isPayloadNull: bool
|
|
schemaId: int
|
|
size: int
|
|
|
|
class InteractionModel(Base):
|
|
partitionID: int
|
|
offset: int
|
|
timestamp: int
|
|
compression: str
|
|
isTransactional: bool
|
|
headers: list
|
|
key: dict
|
|
value: ValueModel
|
|
|
|
class Loader:
|
|
def __init__(self, src_dir: str):
|
|
self.src_dir = src_dir
|
|
self.entries = os.listdir(src_dir)
|
|
if not self.entries: raise ValueError("empty directory")
|
|
self.data = self._load_sessions()
|
|
|
|
def _is_admin_page(self, interaction: InteractionModel) -> bool:
|
|
page = interaction.value.payload.page
|
|
return page and page.startswith("/admin/")
|
|
|
|
def _load_sessions(self) -> dict:
|
|
sessions = {}
|
|
for entry in self.entries:
|
|
int_path = f"{self.src_dir}/{entry}/int.json"
|
|
raw = json.load(open(int_path))
|
|
ints = [InteractionModel(**i) for i in raw]
|
|
sessions[entry] = [i for i in ints if not self._is_admin_page(i)]
|
|
return sessions
|
|
|
|
def get_data(self) -> dict:
|
|
return self.data
|
|
|
|
def get_entries(self) -> tuple[list[str], int]:
|
|
return self.entries, len(self.entries)
|
|
|
|
class AgentLoader(Loader):
|
|
"""Loader for agent interaction data with simplified schema (direct PayloadModel format)"""
|
|
|
|
def _is_admin_page_simple(self, interaction: PayloadModel) -> bool:
|
|
return interaction.page and interaction.page.startswith("/admin/")
|
|
|
|
def _load_sessions(self) -> dict:
|
|
sessions = {}
|
|
for entry in self.entries:
|
|
int_path = f"{self.src_dir}/{entry}/int.json"
|
|
raw = json.load(open(int_path))
|
|
ints = [PayloadModel(**i) for i in raw]
|
|
sessions[entry] = [i for i in ints if not self._is_admin_page_simple(i)]
|
|
return sessions
|
|
|
|
if __name__ == "__main__":
|
|
DIR = "/home/velocitatem/Documents/Projects/PHANTOM/experiments/agents/collected_data/"
|
|
loader = AgentLoader(DIR)
|
|
_, n = loader.get_entries()
|
|
print(f"Loaded {n} sessions from {DIR}")
|
|
|
|
DIR = "/home/velocitatem/Documents/Projects/PHANTOM/experiments/collected_data/"
|
|
loader = Loader(DIR)
|
|
_, n = loader.get_entries()
|
|
print(f"Loaded {n} sessions from {DIR}")
|