348 lines
12 KiB
Python
348 lines
12 KiB
Python
"""
|
|
Real-Time WebSocket Alerts System
|
|
==================================
|
|
|
|
Provides real-time notification via WebSocket for:
|
|
- Transaction monitoring
|
|
- Security alerts
|
|
- Entity changes
|
|
- New wallet discoveries
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from collections import defaultdict
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ─── REAL-TIME ALERT TYPES ──────────────────────────────────────────────
|
|
|
|
|
|
class AlertType:
|
|
"""Supported alert types."""
|
|
|
|
NEW_TRANSACTION = "new_transaction"
|
|
SECURITY_ALERT = "security_alert"
|
|
ENTITY_CHANGE = "entity_change"
|
|
NEW_WALLET = "new_wallet"
|
|
CROSS_CHAIN = "cross_chain_movement"
|
|
RATE_LIMIT = "rate_limit_exceeded"
|
|
SYSTEM_STATUS = "system_status"
|
|
|
|
|
|
class AlertSeverity:
|
|
"""Alert severity levels."""
|
|
|
|
INFO = "info"
|
|
WARNING = "warning"
|
|
HIGH = "high"
|
|
CRITICAL = "critical"
|
|
|
|
|
|
class RealTimeAlert:
|
|
"""Represents a real-time alert."""
|
|
|
|
def __init__(
|
|
self,
|
|
alert_type: str,
|
|
message: str,
|
|
severity: str,
|
|
data: dict[str, Any] | None = None,
|
|
chain: str = "ethereum",
|
|
wallet_address: str = "",
|
|
related_entities: list[str] | None = None,
|
|
):
|
|
self.alert_type = alert_type
|
|
self.message = message
|
|
self.severity = severity
|
|
self.data = data or {}
|
|
self.chain = chain
|
|
self.wallet_address = wallet_address
|
|
self.related_entities = related_entities or []
|
|
self.timestamp = datetime.utcnow().isoformat()
|
|
self.id = f"{alert_type}_{datetime.utcnow().timestamp()}"
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""Convert to dictionary."""
|
|
return {
|
|
"id": self.id,
|
|
"type": self.alert_type,
|
|
"message": self.message,
|
|
"severity": self.severity,
|
|
"chain": self.chain,
|
|
"wallet_address": self.wallet_address,
|
|
"related_entities": self.related_entities,
|
|
"data": self.data,
|
|
"timestamp": self.timestamp,
|
|
}
|
|
|
|
def to_json(self) -> str:
|
|
"""Convert to JSON string."""
|
|
return json.dumps(self.to_dict())
|
|
|
|
|
|
# ─── ALERT BROADCASTER ────────────────────────────────────────────────
|
|
|
|
|
|
class AlertBroadcaster:
|
|
"""Broadcasts alerts to connected WebSocket clients."""
|
|
|
|
def __init__(self):
|
|
self.clients: dict[str, WebSocket] = {}
|
|
self.subscriptions: dict[str, set[str]] = defaultdict(set) # wallet -> client_id
|
|
self.alert_queue: list[RealTimeAlert] = []
|
|
self.max_queue_size = 1000
|
|
|
|
async def connect(self, websocket: WebSocket, client_id: str | None = None):
|
|
"""Accept a WebSocket connection."""
|
|
await websocket.accept()
|
|
if not client_id:
|
|
client_id = f"client_{id(websocket)}"
|
|
self.clients[client_id] = websocket
|
|
logger.info(f"WebSocket client connected: {client_id}")
|
|
return client_id
|
|
|
|
async def disconnect(self, client_id: str):
|
|
"""Remove a client."""
|
|
if client_id in self.clients:
|
|
del self.clients[client_id]
|
|
# Remove subscriptions
|
|
for wallet in self.subscriptions:
|
|
self.subscriptions[wallet].discard(client_id)
|
|
logger.info(f"WebSocket client disconnected: {client_id}")
|
|
|
|
def subscribe(self, client_id: str, wallet_address: str):
|
|
"""Subscribe a client to wallet alerts."""
|
|
self.subscriptions[wallet_address].add(client_id)
|
|
|
|
def unsubscribe(self, client_id: str, wallet_address: str):
|
|
"""Unsubscribe a client from wallet alerts."""
|
|
self.subscriptions[wallet_address].discard(client_id)
|
|
|
|
async def broadcast(self, alert: RealTimeAlert):
|
|
"""Broadcast an alert to relevant clients."""
|
|
self.alert_queue.append(alert)
|
|
|
|
# Keep queue size manageable
|
|
if len(self.alert_queue) > self.max_queue_size:
|
|
self.alert_queue = self.alert_queue[-self.max_queue_size :]
|
|
|
|
# Find interested clients
|
|
interested_clients = set()
|
|
interested_clients.update(self.clients.keys()) # Broadcast to all
|
|
|
|
# Also add clients subscribed to related wallets
|
|
if alert.wallet_address:
|
|
interested_clients.update(self.subscriptions.get(alert.wallet_address, set()))
|
|
|
|
if alert.related_entities:
|
|
for entity in alert.related_entities:
|
|
interested_clients.update(self.subscriptions.get(entity, set()))
|
|
|
|
# Send to interested clients
|
|
for client_id in list(interested_clients):
|
|
if client_id in self.clients:
|
|
try:
|
|
await self.clients[client_id].send_text(alert.to_json())
|
|
except Exception as e:
|
|
logger.error(f"Failed to send alert to {client_id}: {e}")
|
|
await self.disconnect(client_id)
|
|
|
|
async def broadcast_system_status(self, status: str, details: dict[str, Any] | None = None):
|
|
"""Broadcast system status update."""
|
|
alert = RealTimeAlert(
|
|
alert_type=AlertType.SYSTEM_STATUS,
|
|
message=f"System {status}",
|
|
severity=AlertSeverity.INFO,
|
|
data=details or {},
|
|
)
|
|
await self.broadcast(alert)
|
|
|
|
async def broadcast_security_alert(self, wallet: str, severity: str, message: str, chain: str = "ethereum"):
|
|
"""Broadcast a security alert."""
|
|
alert = RealTimeAlert(
|
|
alert_type=AlertType.SECURITY_ALERT,
|
|
message=message,
|
|
severity=severity,
|
|
chain=chain,
|
|
wallet_address=wallet,
|
|
data={"severity": severity, "category": "security"},
|
|
)
|
|
await self.broadcast(alert)
|
|
|
|
async def broadcast_transaction(self, wallet: str, tx_hash: str, chain: str = "ethereum"):
|
|
"""Broadcast a new transaction."""
|
|
alert = RealTimeAlert(
|
|
alert_type=AlertType.NEW_TRANSACTION,
|
|
message=f"New transaction for {wallet[:20]}... on {chain}",
|
|
severity=AlertSeverity.INFO,
|
|
chain=chain,
|
|
wallet_address=wallet,
|
|
data={"tx_hash": tx_hash, "category": "transaction"},
|
|
)
|
|
await self.broadcast(alert)
|
|
|
|
async def broadcast_cross_chain(self, from_chain: str, to_chain: str, wallet: str):
|
|
"""Broadcast cross-chain movement."""
|
|
alert = RealTimeAlert(
|
|
alert_type=AlertType.CROSS_CHAIN,
|
|
message=f"Cross-chain movement: {from_chain} -> {to_chain}",
|
|
severity=AlertSeverity.WARNING,
|
|
chain=to_chain,
|
|
wallet_address=wallet,
|
|
data={"from_chain": from_chain, "to_chain": to_chain, "category": "cross_chain"},
|
|
)
|
|
await self.broadcast(alert)
|
|
|
|
def get_alerts_by_wallet(self, wallet: str) -> list[dict[str, Any]]:
|
|
"""Get recent alerts for a specific wallet."""
|
|
return [alert.to_dict() for alert in self.alert_queue if alert.wallet_address == wallet]
|
|
|
|
def get_alerts_by_type(self, alert_type: str) -> list[dict[str, Any]]:
|
|
"""Get alerts of a specific type."""
|
|
return [alert.to_dict() for alert in self.alert_queue if alert.alert_type == alert_type]
|
|
|
|
|
|
# ─── WEBSOCKET HANDLER ────────────────────────────────────────────────
|
|
|
|
|
|
class WebSocketManager:
|
|
"""Manages WebSocket connections and events."""
|
|
|
|
def __init__(self):
|
|
self.broadcaster = AlertBroadcaster()
|
|
self.active_connections: dict[str, WebSocket] = {}
|
|
self.connection_lock = asyncio.Lock()
|
|
|
|
async def connect(self, websocket: WebSocket, client_id: str | None = None):
|
|
"""Connect a new WebSocket client."""
|
|
async with self.connection_lock:
|
|
client_id = await self.broadcaster.connect(websocket, client_id)
|
|
self.active_connections[client_id] = websocket
|
|
await websocket.send_json({"type": "connected", "client_id": client_id})
|
|
|
|
async def disconnect(self, client_id: str):
|
|
"""Disconnect a WebSocket client."""
|
|
async with self.connection_lock:
|
|
await self.broadcaster.disconnect(client_id)
|
|
if client_id in self.active_connections:
|
|
del self.active_connections[client_id]
|
|
|
|
async def receive_message(self, client_id: str, message: str):
|
|
"""Handle incoming message from client."""
|
|
try:
|
|
data = json.loads(message)
|
|
msg_type = data.get("type")
|
|
|
|
if msg_type == "subscribe":
|
|
wallet = data.get("wallet")
|
|
if wallet:
|
|
self.broadcaster.subscribe(client_id, wallet)
|
|
await self.active_connections[client_id].send_json({"type": "subscribed", "wallet": wallet})
|
|
|
|
elif msg_type == "unsubscribe":
|
|
wallet = data.get("wallet")
|
|
if wallet:
|
|
self.broadcaster.unsubscribe(client_id, wallet)
|
|
|
|
elif msg_type == "ping":
|
|
await self.active_connections[client_id].send_json(
|
|
{"type": "pong", "timestamp": datetime.utcnow().isoformat()}
|
|
)
|
|
|
|
else:
|
|
logger.warning(f"Unknown message type: {msg_type}")
|
|
|
|
except json.JSONDecodeError:
|
|
logger.error(f"Invalid JSON from {client_id}")
|
|
|
|
async def broadcast(self, alert: RealTimeAlert):
|
|
"""Broadcast an alert to all connected clients."""
|
|
await self.broadcaster.broadcast(alert)
|
|
|
|
async def broadcast_security_alert(self, wallet: str, severity: str, message: str, chain: str = "ethereum"):
|
|
"""Broadcast a security alert."""
|
|
await self.broadcaster.broadcast_security_alert(wallet, severity, message, chain)
|
|
|
|
async def broadcast_transaction(self, wallet: str, tx_hash: str, chain: str = "ethereum"):
|
|
"""Broadcast a new transaction."""
|
|
await self.broadcaster.broadcast_transaction(wallet, tx_hash, chain)
|
|
|
|
async def broadcast_cross_chain(self, from_chain: str, to_chain: str, wallet: str):
|
|
"""Broadcast cross-chain movement."""
|
|
await self.broadcaster.broadcast_cross_chain(from_chain, to_chain, wallet)
|
|
|
|
|
|
# ─── FastAPI Integration ──────────────────────────────────────────────
|
|
|
|
# Global instance
|
|
ws_manager = WebSocketManager()
|
|
|
|
|
|
async def start_websocket_handler():
|
|
"""Start the WebSocket manager."""
|
|
logger.info("WebSocket manager started")
|
|
return ws_manager
|
|
|
|
|
|
def get_websocket_manager() -> WebSocketManager:
|
|
"""Get the global WebSocket manager instance."""
|
|
return ws_manager
|
|
|
|
|
|
async def websocket_alert_endpoint(websocket: WebSocket, client_id: str = "default"):
|
|
"""WebSocket endpoint for real-time alerts."""
|
|
await ws_manager.connect(websocket, client_id)
|
|
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_text()
|
|
await ws_manager.receive_message(client_id, data)
|
|
except WebSocketDisconnect:
|
|
await ws_manager.disconnect(client_id)
|
|
except Exception as e:
|
|
logger.error(f"WebSocket error: {e}")
|
|
await ws_manager.disconnect(client_id)
|
|
|
|
|
|
# ─── TEST HELPER ──────────────────────────────────────────────────────
|
|
|
|
|
|
async def test_websocket_alerts():
|
|
"""Test the WebSocket alert system."""
|
|
import asyncio
|
|
|
|
async def mock_client():
|
|
"""Simulate a WebSocket client."""
|
|
import websockets
|
|
|
|
uri = "ws://localhost:8000/ws/alerts"
|
|
|
|
async with websockets.connect(uri) as ws:
|
|
# Subscribe to a wallet
|
|
await ws.send(json.dumps({"type": "subscribe", "wallet": "0x1234..."}))
|
|
response = await ws.recv()
|
|
print(f"Subscribed: {response}")
|
|
|
|
# Wait for alerts
|
|
for _ in range(3):
|
|
try:
|
|
alert = await asyncio.wait_for(ws.recv(), timeout=5)
|
|
print(f"Alert received: {alert}")
|
|
except TimeoutError:
|
|
print("No alert received")
|
|
|
|
# Run test
|
|
await mock_client()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
|
|
asyncio.run(test_websocket_alerts())
|