144 lines
5.3 KiB
Python
144 lines
5.3 KiB
Python
import json
|
|
import logging
|
|
from typing import Dict, Any, Optional, Tuple
|
|
try:
|
|
import consul
|
|
CONSUL_AVAILABLE = True
|
|
except ImportError:
|
|
consul = None
|
|
CONSUL_AVAILABLE = False
|
|
|
|
class ConsulPersistence:
|
|
"""Handles Consul-based state persistence for connection monitoring"""
|
|
|
|
def __init__(self, consul_url: str = 'http://consul.service.dc1.consul:8500'):
|
|
"""
|
|
Initialize Consul persistence
|
|
|
|
Args:
|
|
consul_url: Consul server URL
|
|
"""
|
|
self.consul_url = consul_url
|
|
self.logger = logging.getLogger(__name__)
|
|
self.consul_client = None
|
|
self.base_key = "qbitcheck/connection_monitor/"
|
|
|
|
if CONSUL_AVAILABLE:
|
|
self._initialize_consul_client()
|
|
else:
|
|
self.logger.warning("python-consul package not available. State persistence disabled.")
|
|
|
|
def _initialize_consul_client(self) -> bool:
|
|
"""Initialize Consul client if available"""
|
|
if not CONSUL_AVAILABLE:
|
|
return False
|
|
|
|
try:
|
|
# Parse URL to extract host and port
|
|
url_parts = self.consul_url.split('://')
|
|
if len(url_parts) < 2:
|
|
raise ValueError(f"Invalid Consul URL format: {self.consul_url}")
|
|
|
|
host_port = url_parts[1].split(':')
|
|
host = host_port[0]
|
|
port = int(host_port[1]) if len(host_port) > 1 else 8500
|
|
|
|
self.consul_client = consul.Consul(host=host, port=port)
|
|
self.logger.info(f"Consul client initialized for {self.consul_url}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to initialize Consul client: {e}")
|
|
self.consul_client = None
|
|
return False
|
|
|
|
def save_state(self, state_data: Dict[str, Any],
|
|
remediation_data: Dict[str, Any],
|
|
stability_data: Dict[str, Any],
|
|
vpn_data: Dict[str, Any]) -> bool:
|
|
"""
|
|
Save state to Consul KV store
|
|
|
|
Args:
|
|
state_data: Connection state data
|
|
remediation_data: Remediation state data
|
|
stability_data: Stability tracking data
|
|
vpn_data: VPN status and IP data
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
if not self.consul_client:
|
|
return False
|
|
|
|
try:
|
|
# Save each section to Consul
|
|
self.consul_client.kv.put(f"{self.base_key}state", json.dumps(state_data))
|
|
self.consul_client.kv.put(f"{self.base_key}remediation", json.dumps(remediation_data))
|
|
self.consul_client.kv.put(f"{self.base_key}stability", json.dumps(stability_data))
|
|
self.consul_client.kv.put(f"{self.base_key}vpn", json.dumps(vpn_data))
|
|
|
|
self.logger.debug("State successfully saved to Consul")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to save state to Consul: {e}")
|
|
return False
|
|
|
|
def load_state(self) -> Tuple[Optional[Dict[str, Any]],
|
|
Optional[Dict[str, Any]],
|
|
Optional[Dict[str, Any]],
|
|
Optional[Dict[str, Any]]]:
|
|
"""
|
|
Load state from Consul KV store
|
|
|
|
Returns:
|
|
Tuple of (state_data, remediation_data, stability_data, vpn_data)
|
|
"""
|
|
if not self.consul_client:
|
|
return None, None, None, None
|
|
|
|
try:
|
|
state_data = None
|
|
remediation_data = None
|
|
stability_data = None
|
|
vpn_data = None
|
|
|
|
# Load connection state
|
|
_, state_kv = self.consul_client.kv.get(f"{self.base_key}state")
|
|
if state_kv:
|
|
state_data = json.loads(state_kv['Value'].decode('utf-8'))
|
|
|
|
# Load remediation state
|
|
_, remediation_kv = self.consul_client.kv.get(f"{self.base_key}remediation")
|
|
if remediation_kv:
|
|
remediation_data = json.loads(remediation_kv['Value'].decode('utf-8'))
|
|
|
|
# Load stability tracking
|
|
_, stability_kv = self.consul_client.kv.get(f"{self.base_key}stability")
|
|
if stability_kv:
|
|
stability_data = json.loads(stability_kv['Value'].decode('utf-8'))
|
|
|
|
# Load VPN state
|
|
_, vpn_kv = self.consul_client.kv.get(f"{self.base_key}vpn")
|
|
if vpn_kv:
|
|
vpn_data = json.loads(vpn_kv['Value'].decode('utf-8'))
|
|
|
|
self.logger.info("State successfully loaded from Consul")
|
|
return state_data, remediation_data, stability_data, vpn_data
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to load state from Consul: {e}")
|
|
return None, None, None, None
|
|
|
|
def is_available(self) -> bool:
|
|
"""Check if Consul persistence is available"""
|
|
return self.consul_client is not None and CONSUL_AVAILABLE
|
|
|
|
def get_consul_status(self) -> Dict[str, Any]:
|
|
"""Get Consul connection status"""
|
|
return {
|
|
'available': self.is_available(),
|
|
'url': self.consul_url,
|
|
'client_initialized': self.consul_client is not None
|
|
} |