232 lines
8.7 KiB
Python
232 lines
8.7 KiB
Python
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}") |