Update notifications.py

Author: webxos

main/api/mcp/lib/notifications.py

@@ -1,63 +1,89 @@
-import json
 import logging
-from typing import Dict, Any
-from fastapi import WebSocket
-from config.config import DatabaseConfig
+from typing import Dict, Any, Optional
+from config.config import DatabaseConfig, SMTPConfig
+from email.mime.text import MIMEText
+import smtplib
+import asyncio
+from datetime import datetime
+from pydantic import BaseModel
 
 logger = logging.getLogger("mcp.notifications")
 logger.setLevel(logging.INFO)
 
-class NotificationHandler:
-    def __init__(self):
-        self.active_connections: Dict[str, WebSocket] = {}
-
-    async def connect(self, websocket: WebSocket, client_id: str):
-        await websocket.accept()
-        self.active_connections[client_id] = websocket
-        await self.notify(client_id, {
-            "method": "connection",
-            "params": {"status": "connected", "client_id": client_id}
-        })
-        logger.info(f"WebSocket connected for client {client_id}")
-
-    async def disconnect(self, client_id: str):
-        if client_id in self.active_connections:
-            await self.active_connections[client_id].close()
-            del self.active_connections[client_id]
-            logger.info(f"WebSocket disconnected for client {client_id}")
+class NotificationInput(BaseModel):
+    user_id: str
+    notification_type: str
+    details: Dict[str, Any]
 
-    async def notify(self, client_id: str, message: Dict[str, Any]):
-        if client_id in self.active_connections:
-            try:
-                await self.active_connections[client_id].send_text(json.dumps(message))
-            except Exception as e:
-                logger.error(f"Error sending notification to {client_id}: {str(e)}")
-                await self.disconnect(client_id)
-
-    async def notify_auth(self, user_id: str, access_token: str):
-        await self.notify(user_id, {
-            "method": "auth.success",
-            "params": {"user_id": user_id, "access_token": access_token}
-        })
-        logger.info(f"Notified auth success for {user_id}")
-
-    async def notify_git_push(self, user_id: str, vial_id: str, commit_hash: str, balance: float):
-        await self.notify(user_id, {
-            "method": "vial_management.gitPush",
-            "params": {"vial_id": vial_id, "commit_hash": commit_hash, "balance": balance}
-        })
-        logger.info(f"Notified Git push for {user_id}, vial {vial_id}")
+class NotificationHandler:
+    def __init__(self, db: DatabaseConfig):
+        self.db = db
+        self.smtp_config = SMTPConfig()
 
-    async def notify_cash_out(self, user_id: str, transaction_id: str, amount: float, new_balance: float):
-        await self.notify(user_id, {
-            "method": "wallet.cashOut",
-            "params": {"transaction_id": transaction_id, "amount": amount, "new_balance": new_balance}
-        })
-        logger.info(f"Notified cash-out for {user_id}, transaction {transaction_id}")
+    async def send_notification(self, input: NotificationInput):
+        try:
+            # Store notification in database
+            await self.db.query(
+                """
+                INSERT INTO notifications (user_id, notification_type, details, created_at)
+                VALUES ($1, $2, $3, $4)
+                """,
+                [
+                    input.user_id,
+                    input.notification_type,
+                    json.dumps(input.details),
+                    datetime.utcnow()
+                ]
+            )
+            
+            # Send email notification for critical events
+            if input.notification_type in ["anomaly_detected", "data_erasure"]:
+                subject = f"Vial MCP: {input.notification_type.replace('_', ' ').title()}"
+                body = f"Notification for user {input.user_id}:\n\nType: {input.notification_type}\nDetails: {json.dumps(input.details, indent=2)}"
+                
+                msg = MIMEText(body)
+                msg["Subject"] = subject
+                msg["From"] = self.smtp_config.smtp_user
+                msg["To"] = self.smtp_config.alert_email
+                
+                with smtplib.SMTP(self.smtp_config.smtp_server, self.smtp_config.smtp_port) as server:
+                    server.starttls()
+                    server.login(self.smtp_config.smtp_user, self.smtp_config.smtp_password)
+                    server.send_message(msg)
+                
+                logger.info(f"Sent notification email for {input.notification_type} to user {input.user_id}")
+            
+            logger.info(f"Stored notification for user {input.user_id}: {input.notification_type}")
+        except Exception as e:
+            logger.error(f"Error sending notification: {str(e)}")
+            raise
 
-    async def notify_wallet_update(self, user_id: str, vial_id: str, balance: float):
-        await self.notify(user_id, {
-            "method": "wallet.update",
-            "params": {"vial_id": vial_id, "balance": balance}
-        })
-        logger.info(f"Notified wallet update for {user_id}, vial {vial_id}")
+    async def monitor_audit_logs(self):
+        """Monitor audit logs for critical events and send notifications."""
+        try:
+            while True:
+                # Check for recent critical audit log entries
+                critical_actions = await self.db.query(
+                    """
+                    SELECT user_id, action, details FROM audit_logs 
+                    WHERE action IN ('anomaly_detected', 'data_erasure') 
+                    AND created_at > $1
+                    AND notified IS NULL
+                    """,
+                    [datetime.utcnow() - timedelta(minutes=5)]
+                )
+                
+                for row in critical_actions.rows:
+                    await self.send_notification(NotificationInput(
+                        user_id=row["user_id"],
+                        notification_type=row["action"],
+                        details=json.loads(row["details"])
+                    ))
+                    await self.db.query(
+                        "UPDATE audit_logs SET notified = TRUE WHERE user_id = $1 AND action = $2 AND created_at = $3",
+                        [row["user_id"], row["action"], row["created_at"]]
+                    )
+                
+                await asyncio.sleep(60)  # Check every minute
+        except Exception as e:
+            logger.error(f"Error monitoring audit logs: {str(e)}")