449 lines
16 KiB
Python
449 lines
16 KiB
Python
"""
|
|
Historical Data Storage & Analytics Module
|
|
==========================================
|
|
|
|
Provides persistent storage for:
|
|
- Transaction history
|
|
- Entity relationships
|
|
- Alert history
|
|
- Analytics queries
|
|
|
|
Uses Redis for fast-access cache and optional PostgreSQL for long-term storage.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
import redis
|
|
from pydantic import BaseModel, Field
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ─── PERSISTENT MODELS ────────────────────────────────────────────────
|
|
|
|
|
|
class TransactionRecord(BaseModel):
|
|
"""Represents a transaction record."""
|
|
|
|
tx_hash: str
|
|
chain: str = "ethereum"
|
|
from_address: str
|
|
to_address: str
|
|
value: float = 0.0
|
|
gas_used: int = 0
|
|
gas_price: int = 0
|
|
block_number: int = 0
|
|
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
|
status: int = 1 # 1 = success, 0 = failed
|
|
function_name: str = ""
|
|
token_transfers: list[dict[str, Any]] = Field(default_factory=list)
|
|
|
|
|
|
class WalletHistory(BaseModel):
|
|
"""Complete transaction history for a wallet."""
|
|
|
|
wallet_address: str
|
|
chain: str = "ethereum"
|
|
transactions: list[TransactionRecord] = Field(default_factory=list)
|
|
first_seen: datetime = Field(default_factory=datetime.utcnow)
|
|
last_seen: datetime = Field(default_factory=datetime.utcnow)
|
|
total_tx_count: int = 0
|
|
total_volume: float = 0.0
|
|
unique_interactions: int = 0
|
|
|
|
|
|
class EntityAlertRecord(BaseModel):
|
|
"""Record of an alert for an entity."""
|
|
|
|
alert_id: str
|
|
entity_id: str
|
|
alert_type: str
|
|
severity: str
|
|
message: str
|
|
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
|
metadata: dict[str, Any] = Field(default_factory=dict)
|
|
resolved: bool = False
|
|
|
|
|
|
# ─── REDIS STORAGE ────────────────────────────────────────────────────
|
|
|
|
|
|
class RedisStorage:
|
|
"""Redis-backed storage for historical data."""
|
|
|
|
def __init__(self, host: str = "localhost", port: int = 6379, db: int = 0):
|
|
self.host = host
|
|
self.port = port
|
|
self.db = db
|
|
self.client = redis.Redis(host=host, port=port, db=db, decode_responses=True)
|
|
logger.info(f"Connected to Redis at {host}:{port}/{db}")
|
|
|
|
def key_prefix(self, category: str, identifier: str) -> str:
|
|
"""Generate a Redis key."""
|
|
return f"rmi:{category}:{identifier}"
|
|
|
|
def save_transaction(self, tx: TransactionRecord) -> bool:
|
|
"""Save a transaction record."""
|
|
try:
|
|
key = self.key_prefix("transaction", tx.tx_hash)
|
|
self.client.setex(
|
|
key,
|
|
86400 * 30, # 30 days TTL
|
|
tx.json(),
|
|
)
|
|
|
|
# Index by wallet
|
|
wallet_key = self.key_prefix("wallet", f"{tx.from_address}_{tx.chain}")
|
|
self.client.zadd(wallet_key, {tx.tx_hash: tx.timestamp.timestamp()})
|
|
|
|
# Update wallet history
|
|
self._update_wallet_history(tx.wallet_address if hasattr(tx, "wallet_address") else tx.from_address, tx)
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to save transaction: {e}")
|
|
return False
|
|
|
|
def get_transaction(self, tx_hash: str) -> dict[str, Any] | None:
|
|
"""Get a transaction record."""
|
|
try:
|
|
key = self.key_prefix("transaction", tx_hash)
|
|
data = self.client.get(key)
|
|
return json.loads(data) if data else None
|
|
except Exception as e:
|
|
logger.error(f"Failed to get transaction: {e}")
|
|
return None
|
|
|
|
def get_wallet_history(self, wallet_address: str, chain: str = "ethereum") -> WalletHistory | None:
|
|
"""Get complete wallet history."""
|
|
try:
|
|
key = self.key_prefix("wallet", f"{wallet_address}_{chain}")
|
|
|
|
if not self.client.exists(key):
|
|
return None
|
|
|
|
tx_hashes = self.client.zrange(key, 0, -1, withscores=True)
|
|
|
|
transactions = []
|
|
for tx_hash, score in tx_hashes:
|
|
tx = self.get_transaction(tx_hash)
|
|
if tx:
|
|
tx_record = TransactionRecord(**tx)
|
|
tx_record.timestamp = datetime.fromtimestamp(score)
|
|
transactions.append(tx_record)
|
|
|
|
# Sort by timestamp
|
|
transactions.sort(key=lambda x: x.timestamp)
|
|
|
|
return WalletHistory(
|
|
wallet_address=wallet_address,
|
|
chain=chain,
|
|
transactions=transactions,
|
|
total_tx_count=len(transactions),
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to get wallet history: {e}")
|
|
return None
|
|
|
|
def save_alert(self, alert: EntityAlertRecord) -> bool:
|
|
"""Save an alert record."""
|
|
try:
|
|
key = self.key_prefix("alert", alert.alert_id)
|
|
self.client.setex(
|
|
key,
|
|
86400 * 90, # 90 days TTL
|
|
alert.json(),
|
|
)
|
|
|
|
# Index by entity
|
|
entity_key = self.key_prefix("entity_alert", alert.entity_id)
|
|
self.client.zadd(entity_key, {alert.alert_id: alert.timestamp.timestamp()})
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to save alert: {e}")
|
|
return False
|
|
|
|
def get_entity_alerts(self, entity_id: str, limit: int = 100) -> list[dict[str, Any]]:
|
|
"""Get alerts for an entity."""
|
|
try:
|
|
key = self.key_prefix("entity_alert", entity_id)
|
|
|
|
if not self.client.exists(key):
|
|
return []
|
|
|
|
alert_ids = self.client.zrange(key, 0, limit - 1)
|
|
|
|
alerts = []
|
|
for alert_id in alert_ids:
|
|
alert = self.get_alert(alert_id)
|
|
if alert:
|
|
alerts.append(alert)
|
|
|
|
return alerts
|
|
except Exception as e:
|
|
logger.error(f"Failed to get entity alerts: {e}")
|
|
return []
|
|
|
|
def get_alert(self, alert_id: str) -> dict[str, Any] | None:
|
|
"""Get an alert record."""
|
|
try:
|
|
key = self.key_prefix("alert", alert_id)
|
|
data = self.client.get(key)
|
|
return json.loads(data) if data else None
|
|
except Exception as e:
|
|
logger.error(f"Failed to get alert: {e}")
|
|
return None
|
|
|
|
def save_entity_relation(self, from_entity: str, to_entity: str, relation: str):
|
|
"""Save an entity relationship."""
|
|
try:
|
|
key = self.key_prefix("entity_relation", from_entity)
|
|
self.client.sadd(key, json.dumps({"to": to_entity, "relation": relation}))
|
|
except Exception as e:
|
|
logger.error(f"Failed to save entity relation: {e}")
|
|
|
|
def get_entity_relations(self, entity_id: str) -> list[dict[str, str]]:
|
|
"""Get relations for an entity."""
|
|
try:
|
|
key = self.key_prefix("entity_relation", entity_id)
|
|
relations = self.client.smembers(key)
|
|
return [json.loads(r) for r in relations]
|
|
except Exception as e:
|
|
logger.error(f"Failed to get entity relations: {e}")
|
|
return []
|
|
|
|
def save_wallet_cluster(self, cluster_id: str, members: list[str], labels: list[str]):
|
|
"""Save a wallet cluster."""
|
|
try:
|
|
key = self.key_prefix("cluster", cluster_id)
|
|
data = {
|
|
"cluster_id": cluster_id,
|
|
"members": members,
|
|
"labels": labels,
|
|
"created_at": datetime.utcnow().isoformat(),
|
|
}
|
|
self.client.setex(key, 86400 * 365, json.dumps(data)) # 1 year TTL
|
|
except Exception as e:
|
|
logger.error(f"Failed to save cluster: {e}")
|
|
|
|
def get_wallet_cluster(self, cluster_id: str) -> dict[str, Any] | None:
|
|
"""Get a wallet cluster."""
|
|
try:
|
|
key = self.key_prefix("cluster", cluster_id)
|
|
data = self.client.get(key)
|
|
return json.loads(data) if data else None
|
|
except Exception as e:
|
|
logger.error(f"Failed to get cluster: {e}")
|
|
return None
|
|
|
|
def get_or_create_wallet_history(self, wallet_address: str, chain: str = "ethereum") -> WalletHistory:
|
|
"""Get or create wallet history."""
|
|
history = self.get_wallet_history(wallet_address, chain)
|
|
if history is None:
|
|
history = WalletHistory(wallet_address=wallet_address, chain=chain)
|
|
return history
|
|
|
|
def _update_wallet_history(self, wallet_address: str, tx: TransactionRecord):
|
|
"""Update wallet history metadata."""
|
|
history = self.get_or_create_wallet_history(wallet_address, tx.chain)
|
|
|
|
# Update last seen
|
|
history.last_seen = tx.timestamp
|
|
|
|
# Update total volume
|
|
history.total_volume += tx.value
|
|
|
|
# Update interaction count
|
|
if tx.to_address not in [t.to_address for t in history.transactions]:
|
|
history.unique_interactions += 1
|
|
|
|
# Update transaction count
|
|
history.total_tx_count = len(history.transactions) + 1
|
|
|
|
|
|
# ─── DATABASE WRAPPER ─────────────────────────────────────────────────
|
|
|
|
|
|
class AnalyticsDatabase:
|
|
"""Database wrapper for analytics queries."""
|
|
|
|
def __init__(self, redis_host: str = "localhost", redis_port: int = 6379, redis_db: int = 0):
|
|
self.redis = RedisStorage(host=redis_host, port=redis_port, db=redis_db)
|
|
|
|
def store_transaction(self, tx: TransactionRecord) -> bool:
|
|
"""Store a transaction."""
|
|
return self.redis.save_transaction(tx)
|
|
|
|
def store_entity_alert(self, alert: EntityAlertRecord) -> bool:
|
|
"""Store an entity alert."""
|
|
return self.redis.save_alert(alert)
|
|
|
|
def store_entity_relation(self, from_entity: str, to_entity: str, relation: str):
|
|
"""Store an entity relationship."""
|
|
self.redis.save_entity_relation(from_entity, to_entity, relation)
|
|
|
|
def store_wallet_cluster(self, cluster_id: str, members: list[str], labels: list[str]):
|
|
"""Store a wallet cluster."""
|
|
self.redis.save_wallet_cluster(cluster_id, members, labels)
|
|
|
|
def get_wallet_history(self, wallet_address: str, chain: str = "ethereum") -> WalletHistory | None:
|
|
"""Get wallet history."""
|
|
return self.redis.get_wallet_history(wallet_address, chain)
|
|
|
|
def get_entity_alerts(self, entity_id: str, limit: int = 100) -> list[dict[str, Any]]:
|
|
"""Get entity alerts."""
|
|
return self.redis.get_entity_alerts(entity_id, limit)
|
|
|
|
def get_entity_relations(self, entity_id: str) -> list[dict[str, str]]:
|
|
"""Get entity relations."""
|
|
return self.redis.get_entity_relations(entity_id)
|
|
|
|
def get_wallet_cluster(self, cluster_id: str) -> dict[str, Any] | None:
|
|
"""Get wallet cluster."""
|
|
return self.redis.get_wallet_cluster(cluster_id)
|
|
|
|
def get_transaction(self, tx_hash: str) -> dict[str, Any] | None:
|
|
"""Get transaction by hash."""
|
|
return self.redis.get_transaction(tx_hash)
|
|
|
|
# ─── ANALYTICS QUERIES ────────────────────────────────────────────
|
|
|
|
def get_wallet_activity_summary(self, wallet_address: str, chain: str = "ethereum") -> dict[str, Any]:
|
|
"""Get activity summary for a wallet."""
|
|
history = self.redis.get_wallet_history(wallet_address, chain)
|
|
|
|
if history is None or not history.transactions:
|
|
return {
|
|
"wallet_address": wallet_address,
|
|
"chain": chain,
|
|
"total_transactions": 0,
|
|
"total_volume": 0,
|
|
"first_seen": None,
|
|
"last_seen": None,
|
|
"unique_contracts": 0,
|
|
}
|
|
|
|
return {
|
|
"wallet_address": wallet_address,
|
|
"chain": chain,
|
|
"total_transactions": len(history.transactions),
|
|
"total_volume": history.total_volume,
|
|
"first_seen": history.first_seen.isoformat(),
|
|
"last_seen": history.last_seen.isoformat(),
|
|
"unique_contracts": history.unique_interactions,
|
|
}
|
|
|
|
def get_wallet_similarity(self, address1: str, address2: str) -> dict[str, Any]:
|
|
"""Calculate similarity between two wallets based on interactions."""
|
|
history1 = self.redis.get_wallet_history(address1, "ethereum")
|
|
history2 = self.redis.get_wallet_history(address2, "ethereum")
|
|
|
|
if not history1 or not history2 or not history1.transactions or not history2.transactions:
|
|
return {"similarity": 0, "shared_contracts": [], "reason": "Insufficient data"}
|
|
|
|
# Get unique contract interactions
|
|
contracts1 = {t.to_address for t in history1.transactions}
|
|
contracts2 = {t.to_address for t in history2.transactions}
|
|
|
|
# Calculate Jaccard similarity
|
|
intersection = contracts1 & contracts2
|
|
union = contracts1 | contracts2
|
|
|
|
jaccard = len(intersection) / len(union) if union else 0
|
|
|
|
return {
|
|
"similarity": round(jaccard, 4),
|
|
"shared_contracts": list(intersection)[:10], # Top 10
|
|
"total_shared": len(intersection),
|
|
"unique_1": len(contracts1 - contracts2),
|
|
"unique_2": len(contracts2 - contracts1),
|
|
}
|
|
|
|
def get_entity_network(self, entity_id: str, depth: int = 2) -> dict[str, Any]:
|
|
"""Get entity's network of connected entities."""
|
|
relations = self.redis.get_entity_relations(entity_id)
|
|
|
|
network = {"entity_id": entity_id, "direct_relations": relations, "depth": depth}
|
|
|
|
# If depth > 0, get relations of related entities
|
|
if depth > 0:
|
|
related_entities = [r["to"] for r in relations]
|
|
network["related_entities"] = related_entities
|
|
|
|
if depth >= 2:
|
|
network["second_degree"] = []
|
|
for related in related_entities:
|
|
second_degree = self.redis.get_entity_relations(related)
|
|
network["second_degree"].extend(second_degree)
|
|
|
|
return network
|
|
|
|
|
|
# ─── SINGLETON INSTANCE ───────────────────────────────────────────────
|
|
|
|
_db_instance: AnalyticsDatabase | None = None
|
|
|
|
|
|
def get_analytics_database(
|
|
redis_host: str | None = None, redis_port: int | None = None, redis_db: int | None = None
|
|
) -> AnalyticsDatabase:
|
|
"""Get the analytics database instance."""
|
|
global _db_instance
|
|
|
|
if _db_instance is None:
|
|
_db_instance = AnalyticsDatabase(
|
|
redis_host=redis_host or "localhost",
|
|
redis_port=redis_port or 6379,
|
|
redis_db=redis_db or 0,
|
|
)
|
|
|
|
return _db_instance
|
|
|
|
|
|
# ─── INITIAL DATA IMPORT ──────────────────────────────────────────────
|
|
|
|
|
|
def initialize_analytics():
|
|
"""Initialize analytics storage with default data."""
|
|
get_analytics_database()
|
|
|
|
# Clear old data (optional - for fresh starts)
|
|
# db.redis.client.flushdb()
|
|
|
|
logger.info("Analytics database initialized")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Test the analytics database
|
|
db = get_analytics_database()
|
|
|
|
# Create a test transaction
|
|
tx = TransactionRecord(
|
|
tx_hash="0x" + "a" * 64,
|
|
chain="ethereum",
|
|
from_address="0x1234567890123456789012345678901234567890",
|
|
to_address="0xabcdef1234567890abcdef1234567890abcdef12",
|
|
value=1.5,
|
|
gas_used=21000,
|
|
block_number=1000000,
|
|
function_name="transfer",
|
|
)
|
|
|
|
db.store_transaction(tx)
|
|
|
|
# Get the transaction back
|
|
stored = db.get_transaction(tx.tx_hash)
|
|
print(f"Stored transaction: {stored}")
|
|
|
|
# Get wallet history
|
|
history = db.get_wallet_history(tx.from_address, "ethereum")
|
|
if history:
|
|
print(f"Wallet history: {history.total_tx_count} transactions")
|
|
|
|
# Get activity summary
|
|
summary = db.get_wallet_activity_summary(tx.from_address, "ethereum")
|
|
print(f"Activity summary: {summary}")
|