SentinelAI / src /data /loaders /generator.py
ACA050's picture
Upload 114 files
5a3b9db verified
Raw
History Blame Contribute Delete
10.7 kB
"""
Implements: 02_Data/06_SYNTHETIC_DATA_GENERATION.md
"""
import pandas as pd
import numpy as np
from faker import Faker
from datetime import datetime, timedelta
import random
from pathlib import Path
import logging
from typing import Dict, List, Any
logger = logging.getLogger(__name__)
class DepartmentGenerator:
def __init__(self, fake: Faker):
self.fake = fake
def generate(self, config: dict) -> pd.DataFrame:
num_departments = config.get("num_departments", 8)
data = []
departments = [
"HR", "Finance", "Engineering", "Sales", "IT Support",
"Marketing", "Legal", "Operations", "Contractors", "Service Accounts"
]
for i in range(1, num_departments + 1):
data.append({
"department_id": i,
"department_name": departments[i % len(departments)],
"work_start_hour": random.choice([7, 8, 9, 20]), # Including shift workers
"work_end_hour": random.choice([16, 17, 18, 6])
})
return pd.DataFrame(data)
class UserGenerator:
def __init__(self, fake: Faker):
self.fake = fake
def generate(self, config: dict, departments: pd.DataFrame) -> pd.DataFrame:
num_users = config.get("num_users", 1000)
data = []
dept_ids = departments["department_id"].tolist()
for i in range(1, num_users + 1):
role_type = random.choice(["Office Worker", "Remote Employee", "Hybrid Worker", "Contractor", "Service Account"])
data.append({
"user_id": i,
"employee_id": f"EMP{i:05d}" if role_type != "Service Account" else f"SVC{i:05d}",
"name": self.fake.name(),
"email": self.fake.email(),
"department_id": random.choice(dept_ids),
"role": role_type,
"privilege_level": random.randint(1, 5) if role_type != "Service Account" else 5,
"office_location": self.fake.city() if "Remote" not in role_type else "Remote",
"employment_status": "Active",
"created_at": datetime.now() - timedelta(days=random.randint(30, 365))
})
return pd.DataFrame(data)
class DeviceGenerator:
def __init__(self, fake: Faker):
self.fake = fake
def generate(self, config: dict, users: pd.DataFrame) -> pd.DataFrame:
num_devices = config.get("num_devices", 1800)
data = []
user_ids = users["user_id"].tolist()
for i in range(1, num_devices + 1):
data.append({
"device_id": i,
"user_id": random.choice(user_ids),
"device_type": random.choice(["Laptop", "Mobile", "Desktop"]),
"operating_system": random.choice(["Windows", "macOS", "Linux"]),
"browser": random.choice(["Chrome", "Firefox", "Edge"]),
"trust_status": random.choice([True, False]),
"fingerprint": self.fake.sha256(),
"registration_date": datetime.now() - timedelta(days=180)
})
return pd.DataFrame(data)
class ResourceGenerator:
def __init__(self, fake: Faker):
self.fake = fake
def generate(self, config: dict, departments: pd.DataFrame) -> pd.DataFrame:
num_resources = config.get("num_resources", 30)
data = []
dept_ids = departments["department_id"].tolist()
for i in range(1, num_resources + 1):
data.append({
"resource_id": i,
"resource_name": f"App_{self.fake.word()}",
"department_id": random.choice(dept_ids),
"sensitivity_level": random.randint(1, 5),
"requires_mfa": random.choice([True, False])
})
return pd.DataFrame(data)
class AuthenticationGenerator:
def __init__(self, fake: Faker):
self.fake = fake
def generate(self, config: dict, users: pd.DataFrame, devices: pd.DataFrame, resources: pd.DataFrame) -> pd.DataFrame:
num_events = config.get("num_events", 10000)
user_ids = users["user_id"].tolist()
device_ids = devices["device_id"].tolist()
resource_ids = resources["resource_id"].tolist()
data = []
base_time = datetime.now() - timedelta(days=config.get("simulation_days", 90))
for i in range(1, num_events + 1):
data.append({
"event_id": i,
"timestamp": base_time + timedelta(minutes=random.randint(1, 120000)),
"user_id": random.choice(user_ids),
"device_id": random.choice(device_ids),
"session_id": random.randint(1, max(1, num_events // 5)),
"resource_id": random.choice(resource_ids),
"ip_address": self.fake.ipv4_private(),
"country": "USA",
"city": self.fake.city(),
"latitude": float(self.fake.latitude()),
"longitude": float(self.fake.longitude()),
"authentication_method": random.choice(["Password", "MFA", "SSO"]),
"authentication_result": "Success" if random.random() > 0.05 else "Failure",
"failure_reason": None,
"attack_label": 0,
"is_attack": False
})
df = pd.DataFrame(data)
# Ensure chronological order
df = df.sort_values("timestamp").reset_index(drop=True)
df["event_id"] = df.index + 1
return df
class SessionGenerator:
def generate(self, events: pd.DataFrame) -> pd.DataFrame:
data = []
for session_id, group in events.groupby("session_id"):
login_time = group["timestamp"].min()
logout_time = group["timestamp"].max() + timedelta(minutes=random.randint(5, 120))
data.append({
"session_id": session_id,
"user_id": group["user_id"].iloc[0],
"login_time": login_time,
"logout_time": logout_time,
"duration_minutes": (logout_time - login_time).total_seconds() / 60,
"resource_count": group["resource_id"].nunique(),
"total_events": len(group)
})
return pd.DataFrame(data)
class AttackInjector:
def __init__(self, fake: Faker):
self.fake = fake
def inject(self, config: dict, events: pd.DataFrame) -> pd.DataFrame:
scenario = config.get("scenario", "Mixed Enterprise Attack")
attack_ratio = config.get("attack_percentage", 0.03)
if scenario == "Normal Activity":
attack_ratio = 0.001 # Keep a tiny bit of noise or 0
num_attacks = int(len(events) * attack_ratio)
attack_indices = random.sample(range(len(events)), num_attacks)
label_map = {
"Credential Stuffing": 1,
"Password Spray": 2,
"Brute Force": 3,
"Impossible Travel": 4,
"Insider Threat": 5,
"Privilege Escalation": 6,
"Lateral Movement": 7,
"Low-and-Slow Exfiltration": 7
}
for idx in attack_indices:
if scenario == "Mixed Enterprise Attack" or scenario == "Normal Activity":
attack_type = random.choice([
"Brute Force", "Credential Stuffing", "Password Spray",
"Impossible Travel", "Insider Threat", "Lateral Movement",
"Privilege Escalation", "Low-and-Slow Exfiltration"
])
else:
# Map scenario name to label name
scenario_map = {
"Brute Force Attack": "Brute Force",
"Credential Stuffing": "Credential Stuffing",
"Password Spray": "Password Spray",
"Impossible Travel": "Impossible Travel",
"Insider Threat": "Insider Threat",
"Suspicious Device": "Lateral Movement"
}
attack_type = scenario_map.get(scenario, "Brute Force")
is_failure = attack_type in ["Brute Force", "Password Spray"]
events.at[idx, "is_attack"] = True
events.at[idx, "attack_label"] = label_map[attack_type]
events.at[idx, "authentication_result"] = "Failure" if is_failure else "Success"
events.at[idx, "failure_reason"] = "Invalid Password" if is_failure else None
events.at[idx, "country"] = self.fake.country()
events.at[idx, "ip_address"] = self.fake.ipv4_public()
return events
class SyntheticDataGenerator:
"""
Implements: 02_Data/06_SYNTHETIC_DATA_GENERATION.md
Orchestrates the sequential generation pipeline.
"""
def __init__(self, config_overrides=None):
self.config = config_overrides or {}
self.seed = self.config.get("seed", 42)
self.fake = Faker()
Faker.seed(self.seed)
np.random.seed(self.seed)
random.seed(self.seed)
self.dept_gen = DepartmentGenerator(self.fake)
self.user_gen = UserGenerator(self.fake)
self.device_gen = DeviceGenerator(self.fake)
self.res_gen = ResourceGenerator(self.fake)
self.auth_gen = AuthenticationGenerator(self.fake)
self.session_gen = SessionGenerator()
self.injector = AttackInjector(self.fake)
def generate(self) -> Dict[str, pd.DataFrame]:
logger.info("Stage 2: Generate organization")
departments = self.dept_gen.generate(self.config)
logger.info("Stage 3: Generate users")
users = self.user_gen.generate(self.config, departments)
logger.info("Stage 4: Assign devices")
devices = self.device_gen.generate(self.config, users)
logger.info("Stage 5: Generate resources")
resources = self.res_gen.generate(self.config, departments)
logger.info("Stage 6: Generate authentication events")
events = self.auth_gen.generate(self.config, users, devices, resources)
logger.info("Stage 7: Generate sessions")
sessions = self.session_gen.generate(events)
logger.info("Stage 8: Inject attacks")
events = self.injector.inject(self.config, events)
return {
"departments": departments,
"users": users,
"devices": devices,
"resources": resources,
"authentication_events": events,
"sessions": sessions
}