mirror of
https://github.com/infinition/Bjorn.git
synced 2025-12-13 08:04:59 +00:00
332 lines
14 KiB
Python
332 lines
14 KiB
Python
# db_utils/studio.py
|
|
# Actions Studio visual editor operations
|
|
|
|
import json
|
|
from typing import Dict, List, Optional
|
|
import logging
|
|
|
|
from logger import Logger
|
|
|
|
logger = Logger(name="db_utils.studio", level=logging.DEBUG)
|
|
|
|
|
|
class StudioOps:
|
|
"""Actions Studio visual editor and workflow operations"""
|
|
|
|
def __init__(self, base):
|
|
self.base = base
|
|
|
|
def create_tables(self):
|
|
"""Create Actions Studio tables"""
|
|
# Studio actions (extended action metadata for visual editor)
|
|
self.base.execute("""
|
|
CREATE TABLE IF NOT EXISTS actions_studio (
|
|
b_class TEXT PRIMARY KEY,
|
|
studio_x REAL,
|
|
studio_y REAL,
|
|
studio_locked INTEGER DEFAULT 0,
|
|
studio_color TEXT,
|
|
studio_metadata TEXT,
|
|
updated_at TEXT DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
""")
|
|
|
|
# Studio edges (relationships between actions)
|
|
self.base.execute("""
|
|
CREATE TABLE IF NOT EXISTS studio_edges (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
from_action TEXT NOT NULL,
|
|
to_action TEXT NOT NULL,
|
|
edge_type TEXT DEFAULT 'requires',
|
|
edge_label TEXT,
|
|
edge_metadata TEXT,
|
|
created_at TEXT DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (from_action) REFERENCES actions_studio(b_class) ON DELETE CASCADE,
|
|
FOREIGN KEY (to_action) REFERENCES actions_studio(b_class) ON DELETE CASCADE,
|
|
UNIQUE(from_action, to_action, edge_type)
|
|
);
|
|
""")
|
|
|
|
# Studio hosts (hosts for test mode)
|
|
self.base.execute("""
|
|
CREATE TABLE IF NOT EXISTS studio_hosts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
mac_address TEXT UNIQUE NOT NULL,
|
|
ips TEXT,
|
|
hostnames TEXT,
|
|
alive INTEGER DEFAULT 1,
|
|
ports TEXT,
|
|
services TEXT,
|
|
vulns TEXT,
|
|
creds TEXT,
|
|
studio_x REAL,
|
|
studio_y REAL,
|
|
is_simulated INTEGER DEFAULT 1,
|
|
metadata TEXT,
|
|
created_at TEXT DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
""")
|
|
|
|
# Studio layouts (saved layout snapshots)
|
|
self.base.execute("""
|
|
CREATE TABLE IF NOT EXISTS studio_layouts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT UNIQUE NOT NULL,
|
|
description TEXT,
|
|
layout_data TEXT NOT NULL,
|
|
screenshot BLOB,
|
|
is_active INTEGER DEFAULT 0,
|
|
created_at TEXT DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TEXT DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
""")
|
|
|
|
logger.debug("Actions Studio tables created/verified")
|
|
|
|
# =========================================================================
|
|
# STUDIO ACTION OPERATIONS
|
|
# =========================================================================
|
|
|
|
def get_studio_actions(self):
|
|
"""Retrieve all studio actions with their positions"""
|
|
return self.base.query("""
|
|
SELECT * FROM actions_studio
|
|
ORDER BY b_priority DESC, b_class
|
|
""")
|
|
|
|
def get_db_actions(self):
|
|
"""Retrieve all actions from the main actions table"""
|
|
return self.base.query("""
|
|
SELECT * FROM actions
|
|
ORDER BY b_priority DESC, b_class
|
|
""")
|
|
|
|
def update_studio_action(self, b_class: str, updates: dict):
|
|
"""Update a studio action"""
|
|
sets = []
|
|
params = []
|
|
for key, value in updates.items():
|
|
sets.append(f"{key} = ?")
|
|
params.append(value)
|
|
params.append(b_class)
|
|
|
|
self.base.execute(f"""
|
|
UPDATE actions_studio
|
|
SET {', '.join(sets)}, updated_at = CURRENT_TIMESTAMP
|
|
WHERE b_class = ?
|
|
""", params)
|
|
|
|
# =========================================================================
|
|
# STUDIO EDGE OPERATIONS
|
|
# =========================================================================
|
|
|
|
def get_studio_edges(self):
|
|
"""Retrieve all studio edges"""
|
|
return self.base.query("SELECT * FROM studio_edges")
|
|
|
|
def upsert_studio_edge(self, from_action: str, to_action: str, edge_type: str, metadata: dict = None):
|
|
"""Create or update a studio edge"""
|
|
meta_json = json.dumps(metadata) if metadata else None
|
|
# Try UPDATE first
|
|
updated = self.base.execute("""
|
|
UPDATE studio_edges
|
|
SET edge_metadata = ?
|
|
WHERE from_action = ? AND to_action = ? AND edge_type = ?
|
|
""", (meta_json, from_action, to_action, edge_type))
|
|
if not updated:
|
|
# If no rows updated, INSERT
|
|
self.base.execute("""
|
|
INSERT OR IGNORE INTO studio_edges(from_action, to_action, edge_type, edge_metadata)
|
|
VALUES(?,?,?,?)
|
|
""", (from_action, to_action, edge_type, meta_json))
|
|
|
|
def delete_studio_edge(self, edge_id: int):
|
|
"""Delete a studio edge"""
|
|
self.base.execute("DELETE FROM studio_edges WHERE id = ?", (edge_id,))
|
|
|
|
# =========================================================================
|
|
# STUDIO HOST OPERATIONS
|
|
# =========================================================================
|
|
|
|
def get_studio_hosts(self, include_real: bool = True):
|
|
"""Retrieve studio hosts"""
|
|
if include_real:
|
|
# Combine real and simulated hosts
|
|
return self.base.query("""
|
|
SELECT mac_address, ips, hostnames, alive, ports,
|
|
NULL as services, NULL as vulns, NULL as creds,
|
|
NULL as studio_x, NULL as studio_y, 0 as is_simulated
|
|
FROM hosts
|
|
UNION ALL
|
|
SELECT mac_address, ips, hostnames, alive, ports,
|
|
services, vulns, creds, studio_x, studio_y, is_simulated
|
|
FROM studio_hosts
|
|
""")
|
|
else:
|
|
return self.base.query("SELECT * FROM studio_hosts WHERE is_simulated = 1")
|
|
|
|
def upsert_studio_host(self, mac_address: str, data: dict):
|
|
"""Create or update a simulated host"""
|
|
self.base.execute("""
|
|
INSERT INTO studio_hosts (
|
|
mac_address, ips, hostnames, alive, ports, services,
|
|
vulns, creds, studio_x, studio_y, is_simulated, metadata
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(mac_address) DO UPDATE SET
|
|
ips = excluded.ips,
|
|
hostnames = excluded.hostnames,
|
|
alive = excluded.alive,
|
|
ports = excluded.ports,
|
|
services = excluded.services,
|
|
vulns = excluded.vulns,
|
|
creds = excluded.creds,
|
|
studio_x = excluded.studio_x,
|
|
studio_y = excluded.studio_y,
|
|
metadata = excluded.metadata
|
|
""", (
|
|
mac_address,
|
|
data.get('ips'),
|
|
data.get('hostnames'),
|
|
data.get('alive', 1),
|
|
data.get('ports'),
|
|
json.dumps(data.get('services', [])),
|
|
json.dumps(data.get('vulns', [])),
|
|
json.dumps(data.get('creds', [])),
|
|
data.get('studio_x'),
|
|
data.get('studio_y'),
|
|
1, # is_simulated
|
|
json.dumps(data.get('metadata', {}))
|
|
))
|
|
|
|
def delete_studio_host(self, mac: str):
|
|
"""Delete a studio host"""
|
|
self.base.execute("DELETE FROM studio_hosts WHERE mac_address = ?", (mac,))
|
|
|
|
# =========================================================================
|
|
# STUDIO LAYOUT OPERATIONS
|
|
# =========================================================================
|
|
|
|
def save_studio_layout(self, name: str, layout_data: dict, description: str = None):
|
|
"""Save a complete layout"""
|
|
self.base.execute("""
|
|
INSERT INTO studio_layouts (name, description, layout_data)
|
|
VALUES (?, ?, ?)
|
|
ON CONFLICT(name) DO UPDATE SET
|
|
description = excluded.description,
|
|
layout_data = excluded.layout_data,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
""", (name, description, json.dumps(layout_data)))
|
|
|
|
def load_studio_layout(self, name: str):
|
|
"""Load a saved layout"""
|
|
row = self.base.query_one("SELECT * FROM studio_layouts WHERE name = ?", (name,))
|
|
if row:
|
|
row['layout_data'] = json.loads(row['layout_data'])
|
|
return row
|
|
|
|
# =========================================================================
|
|
# STUDIO SYNC OPERATIONS
|
|
# =========================================================================
|
|
|
|
def apply_studio_to_runtime(self):
|
|
"""Apply studio configurations to the main actions table"""
|
|
self.base.execute("""
|
|
UPDATE actions
|
|
SET
|
|
b_trigger = (SELECT b_trigger FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_requires = (SELECT b_requires FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_priority = (SELECT b_priority FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_enabled = (SELECT b_enabled FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_timeout = (SELECT b_timeout FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_max_retries = (SELECT b_max_retries FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_cooldown = (SELECT b_cooldown FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_rate_limit = (SELECT b_rate_limit FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_service = (SELECT b_service FROM actions_studio WHERE actions_studio.b_class = actions.b_class),
|
|
b_port = (SELECT b_port FROM actions_studio WHERE actions_studio.b_class = actions.b_class)
|
|
WHERE b_class IN (SELECT b_class FROM actions_studio)
|
|
""")
|
|
|
|
def _replace_actions_studio_with_actions(self, vacuum: bool = False):
|
|
"""
|
|
Reset actions_studio (delete all rows) then resync from actions via _sync_actions_studio_schema_and_rows().
|
|
Optionally run VACUUM.
|
|
"""
|
|
# Ensure table exists so DELETE doesn't fail
|
|
self.base.execute("""
|
|
CREATE TABLE IF NOT EXISTS actions_studio (
|
|
b_class TEXT PRIMARY KEY,
|
|
studio_x REAL,
|
|
studio_y REAL,
|
|
studio_locked INTEGER DEFAULT 0,
|
|
studio_color TEXT,
|
|
studio_metadata TEXT,
|
|
updated_at TEXT DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
""")
|
|
|
|
# Total purge
|
|
self.base.execute("DELETE FROM actions_studio;")
|
|
|
|
# Optional compaction
|
|
if vacuum:
|
|
self.base.execute("VACUUM;")
|
|
|
|
# Non-destructive resynchronization from actions
|
|
self._sync_actions_studio_schema_and_rows()
|
|
|
|
def _sync_actions_studio_schema_and_rows(self):
|
|
"""
|
|
Sync actions_studio with actions table:
|
|
- Create minimal table if needed
|
|
- Add missing columns from actions
|
|
- Insert missing b_class entries
|
|
- Update NULL fields only (non-destructive)
|
|
"""
|
|
# 1) Minimal table: PK + studio_* columns
|
|
self.base.execute("""
|
|
CREATE TABLE IF NOT EXISTS actions_studio (
|
|
b_class TEXT PRIMARY KEY,
|
|
studio_x REAL,
|
|
studio_y REAL,
|
|
studio_locked INTEGER DEFAULT 0,
|
|
studio_color TEXT,
|
|
studio_metadata TEXT,
|
|
updated_at TEXT DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
""")
|
|
|
|
# 2) Dynamically add all columns from actions that are missing in actions_studio
|
|
act_cols = [r["name"] for r in self.base.query("PRAGMA table_info(actions);")]
|
|
stu_cols = [r["name"] for r in self.base.query("PRAGMA table_info(actions_studio);")]
|
|
|
|
# Get column types from actions
|
|
act_col_defs = {r["name"]: r["type"] for r in self.base.query("PRAGMA table_info(actions);")}
|
|
|
|
for col in act_cols:
|
|
if col == "b_class":
|
|
continue
|
|
if col not in stu_cols:
|
|
col_type = act_col_defs.get(col, "TEXT") or "TEXT"
|
|
self.base.execute(f"ALTER TABLE actions_studio ADD COLUMN {col} {col_type};")
|
|
|
|
# 3) Insert missing b_class entries, non-destructive
|
|
self.base.execute("""
|
|
INSERT OR IGNORE INTO actions_studio (b_class)
|
|
SELECT b_class FROM actions;
|
|
""")
|
|
|
|
# 4) Pre-fill only NULL fields from actions (without overwriting)
|
|
for col in act_cols:
|
|
if col == "b_class":
|
|
continue
|
|
# Only update if the studio value is NULL
|
|
self.base.execute(f"""
|
|
UPDATE actions_studio
|
|
SET {col} = (SELECT a.{col} FROM actions a
|
|
WHERE a.b_class = actions_studio.b_class)
|
|
WHERE {col} IS NULL
|
|
AND EXISTS (SELECT 1 FROM actions a WHERE a.b_class = actions_studio.b_class);
|
|
""")
|
|
|
|
# 5) Touch timestamp
|
|
self.base.execute("UPDATE actions_studio SET updated_at = CURRENT_TIMESTAMP;") |