Update utils/audit.py

utils/audit.py

@@ -1,37 +1,40 @@
+# utils/audit.py
 import streamlit as st
-from datetime import datetime
-import uuid
-import json
 import logging
+import json
+import uuid
+from datetime import datetime
 from typing import Optional, Any
 
 class AuditLogger:
     def __init__(self):
+        self.conn = None
         self._setup_logging()
-    
-    def initialize_with_connection(self, conn):
-        self.conn = conn
 
-audit_logger = AuditLogger()
-    
     def _setup_logging(self):
         logging.basicConfig(
             filename='audit.log',
             level=logging.INFO,
             format='%(asctime)s - %(levelname)s - %(message)s'
         )
-    
+
+    def initialize_with_connection(self, conn):
+        self.conn = conn
+
     def log_activity(
-        self, 
-        user_id: str, 
-        action: str, 
+        self,
+        user_id: str,
+        action: str,
         ip_address: Optional[str] = None,
         details: Any = None,
         severity: str = 'info'
     ):
-        cursor = self.conn.cursor()
-        
+        if not self.conn:
+            logging.error("Database connection not initialized")
+            return
+
         try:
+            cursor = self.conn.cursor()
             cursor.execute("""
                 INSERT INTO audit_logs (
                     id, user_id, action, ip_address, timestamp, details, severity
@@ -46,23 +49,27 @@ audit_logger = AuditLogger()
                 severity
             ))
             self.conn.commit()
-            
             logging.info(f"Audit log: {action} by user {user_id}")
-            
         except Exception as e:
             logging.error(f"Failed to create audit log: {str(e)}")
-    
+
     def get_user_activity(self, user_id: str, limit: int = 50):
-        cursor = self.conn.cursor()
-        cursor.execute("""
-            SELECT action, timestamp, details, severity
-            FROM audit_logs 
-            WHERE user_id = ?
-            ORDER BY timestamp DESC
-            LIMIT ?
-        """, (user_id, limit))
-        return cursor.fetchall()
+        if not self.conn:
+            return []
+            
+        try:
+            cursor = self.conn.cursor()
+            cursor.execute("""
+                SELECT action, timestamp, details, severity
+                FROM audit_logs 
+                WHERE user_id = ?
+                ORDER BY timestamp DESC
+                LIMIT ?
+            """, (user_id, limit))
+            return cursor.fetchall()
+        except Exception as e:
+            logging.error(f"Failed to fetch user activity: {str(e)}")
+            return []
 
-audit_logger = AuditLogger(st.session_state.digital_heir.db.conn)
-# Create a global instance without connection
+# Create a single instance
 audit_logger = AuditLogger()