This commit is contained in:
2025-10-17 09:26:56 -04:00
parent 6b68ea19dd
commit 766657da8b
28 changed files with 5302 additions and 15 deletions
+232
View File
@@ -0,0 +1,232 @@
import sqlite3
import logging
from models.agent import Agent
from models.policy import Policy, Allowlist
from datetime import datetime
from typing import List
logger = logging.getLogger(__name__)
# ------------------ Database Initialization ------------------
def initialize_db(db_path: str, schema_path: str, schema_version: str = "1.0"):
try:
with open(schema_path, 'r') as f:
schema_sql = f.read()
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.executescript(schema_sql)
cursor.execute("""
CREATE TABLE IF NOT EXISTS schema_version (
version TEXT,
applied_on TEXT
)
""")
cursor.execute("""
INSERT INTO schema_version (version, applied_on)
VALUES (?, ?)
""", (schema_version, datetime.now().isoformat()))
conn.commit()
conn.close()
logger.info(f"Tracking database initialized at {db_path} with schema version {schema_version}")
except Exception as e:
logger.error(f"Error initializing database: {e}")
# ------------------ Change Tracking ------------------
def compare_and_track_changes(
conn,
current_agents: List[Agent],
current_policies: List[Policy],
current_allowlists: List[Allowlist]
):
try:
cursor = conn.cursor()
# --- Agents ---
cursor.execute("SELECT * FROM agents")
agent_rows = cursor.fetchall()
agent_columns = [desc[0] for desc in cursor.description]
previous_agents = {
row[agent_columns.index("agentid")]: Agent(**dict(zip(agent_columns, row)))
for row in agent_rows
}
for agent in current_agents:
if agent.agentid not in previous_agents:
insert_new_agent(conn, agent)
else:
log_agent_changes(conn, previous_agents[agent.agentid], agent)
# --- Policies ---
cursor.execute("SELECT * FROM policies")
policy_rows = cursor.fetchall()
policy_columns = [desc[0] for desc in cursor.description]
previous_policies = {
row[policy_columns.index("groupid")]: Policy(**dict(zip(policy_columns, row)))
for row in policy_rows
}
for policy in current_policies:
if policy.groupid not in previous_policies:
insert_new_policy(conn, policy)
else:
log_policy_changes(conn, previous_policies[policy.groupid], policy)
# --- Allowlists ---
cursor.execute("SELECT * FROM allowlists")
allowlist_rows = cursor.fetchall()
allowlist_columns = [desc[0] for desc in cursor.description]
previous_allowlists = {
row[allowlist_columns.index("applicationid")]: Allowlist(**dict(zip(allowlist_columns, row)))
for row in allowlist_rows
}
for allowlist in current_allowlists:
if allowlist.applicationid not in previous_allowlists:
insert_new_allowlist(conn, allowlist)
else:
log_allowlist_changes(conn, previous_allowlists[allowlist.applicationid], allowlist)
conn.commit()
logger.info("Change tracking completed successfully.")
except Exception as e:
logger.error(f"Error during change tracking: {e}")
# ------------------ Insert and Log Functions ------------------
def insert_new_agent(conn, agent: Agent):
try:
timestamp = datetime.now().isoformat()
cursor = conn.cursor()
cursor.execute("""
INSERT INTO agents (
agentid, clientversion, domain, freespace, groupid, hostname, ip, localip,
lastcheckin, os, policyversion, status, username, groupname, status_text, firstseen
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(agentid) DO NOTHING
""", (
agent.agentid, agent.clientversion, agent.domain, agent.freespace, agent.groupid,
agent.hostname, agent.ip, agent.localip, agent.lastcheckin, agent.os,
agent.policyversion, agent.status, agent.username, agent.groupname, agent.status_text,
timestamp
))
for field, value in [
("clientversion", agent.clientversion),
("ip", agent.ip),
("localip", agent.localip),
("policyversion", agent.policyversion),
("hostname", agent.hostname),
("status", agent.status)
]:
cursor.execute(f"""
INSERT INTO agent_{field}_history (agentid, {field}, timestamp)
VALUES (?, ?, ?)
""", (agent.agentid, value, timestamp))
conn.commit()
logger.info(f"Inserted new agent: {agent.agentid}")
except Exception as e:
logger.error(f"Error inserting new agent {agent.agentid}: {e}")
def log_agent_changes(conn, previous: Agent, current: Agent):
try:
timestamp = datetime.now().isoformat()
cursor = conn.cursor()
for field in ["clientversion", "ip", "localip", "policyversion", "hostname", "status"]:
if getattr(previous, field) != getattr(current, field):
cursor.execute(f"""
INSERT INTO agent_{field}_history (agentid, {field}, timestamp)
VALUES (?, ?, ?)
""", (current.agentid, getattr(current, field), timestamp))
cursor.execute("""
UPDATE agents SET lastcheckin = ? WHERE agentid = ?
""", (current.lastcheckin, current.agentid))
conn.commit()
logger.info(f"Logged changes for agent: {current.agentid}")
except Exception as e:
logger.error(f"Error logger changes for agent {current.agentid}: {e}")
def insert_new_policy(conn, policy: Policy):
try:
timestamp = datetime.now().isoformat()
cursor = conn.cursor()
cursor.execute("""
INSERT INTO policies (groupid, hidden, name, parent, firstseen)
VALUES (?, ?, ?, ?, ?)
ON CONFLICT(groupid) DO NOTHING
""", (policy.groupid, policy.hidden, policy.name, policy.parent, timestamp))
cursor.execute("""
INSERT INTO policy_name_history (groupid, name, timestamp)
VALUES (?, ?, ?)
""", (policy.groupid, policy.name, timestamp))
conn.commit()
logger.info(f"Inserted new policy: {policy.groupid}")
except Exception as e:
logger.error(f"Error inserting new policy {policy.groupid}: {e}")
def log_policy_changes(conn, previous: Policy, current: Policy):
try:
timestamp = datetime.now().isoformat()
cursor = conn.cursor()
if previous.name != current.name:
cursor.execute("""
INSERT INTO policy_name_history (groupid, name, timestamp)
VALUES (?, ?, ?)
""", (current.groupid, current.name, timestamp))
conn.commit()
logger.info(f"Logged changes for policy: {current.groupid}")
except Exception as e:
logger.error(f"Error logger changes for policy {current.groupid}: {e}")
def insert_new_allowlist(conn, allowlist: Allowlist):
try:
timestamp = datetime.now().isoformat()
cursor = conn.cursor()
cursor.execute("""
INSERT INTO allowlists (applicationid, name, version, firstseen)
VALUES (?, ?, ?, ?)
ON CONFLICT(applicationid) DO NOTHING
""", (allowlist.applicationid, allowlist.name, allowlist.version, timestamp))
for field, value in [("name", allowlist.name), ("version", allowlist.version)]:
cursor.execute(f"""
INSERT INTO allowlist_{field}_history (applicationid, {field}, timestamp)
VALUES (?, ?, ?)
""", (allowlist.applicationid, value, timestamp))
conn.commit()
logger.info(f"Inserted new allowlist: {allowlist.applicationid}")
except Exception as e:
logger.error(f"Error inserting new allowlist {allowlist.applicationid}: {e}")
def log_allowlist_changes(conn, previous: Allowlist, current: Allowlist):
try:
timestamp = datetime.now().isoformat()
cursor = conn.cursor()
for field in ["name", "version"]:
if getattr(previous, field) != getattr(current, field):
cursor.execute(f"""
INSERT INTO allowlist_{field}_history (applicationid, {field}, timestamp)
VALUES (?, ?, ?)
""", (current.applicationid, getattr(current, field), timestamp))
conn.commit()
logger.info(f"Logged changes for allowlist: {current.applicationid}")
except Exception as e:
logger.error(f"Error logger changes for allowlist {current.applicationid}: {e}")