263 lines
11 KiB
Python
263 lines
11 KiB
Python
import time
|
|
import logging
|
|
from typing import Dict, Any, Optional, Tuple
|
|
from .consul_persistence import ConsulPersistence
|
|
|
|
class StateManager:
|
|
"""Manages connection monitoring state with optional Consul persistence"""
|
|
|
|
def __init__(self, consul_url: Optional[str] = None):
|
|
"""
|
|
Initialize state manager
|
|
|
|
Args:
|
|
consul_url: Optional Consul server URL for persistence
|
|
"""
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
# Initialize Consul persistence if URL provided
|
|
self.consul_persistence = ConsulPersistence(consul_url) if consul_url else None
|
|
|
|
# Default state values
|
|
self.connection_state = 'stable'
|
|
self.last_state_change_time = time.time()
|
|
self.consecutive_failures = 0
|
|
self.consecutive_stable_checks = 0
|
|
self.last_failure_time = None
|
|
|
|
# Remediation state
|
|
self.remediation_state = None
|
|
self.remediation_start_time = None
|
|
self.stabilization_checks = 0
|
|
|
|
# Stability tracking
|
|
self.stability_start_time = None
|
|
|
|
# VPN status tracking
|
|
self.vpn_status = 'unknown'
|
|
self.last_vpn_status_change = time.time()
|
|
self.public_ip = None
|
|
self.last_public_ip_change = None
|
|
self.public_ip_details = {}
|
|
|
|
# Load state from Consul if available
|
|
self._load_state()
|
|
|
|
def _load_state(self):
|
|
"""Load state from persistence if available"""
|
|
if self.consul_persistence and self.consul_persistence.is_available():
|
|
state_data, remediation_data, stability_data, vpn_data = self.consul_persistence.load_state()
|
|
|
|
if state_data:
|
|
self.connection_state = state_data.get('connection_state', 'stable')
|
|
self.last_state_change_time = state_data.get('last_state_change_time', time.time())
|
|
self.consecutive_failures = state_data.get('consecutive_failures', 0)
|
|
self.consecutive_stable_checks = state_data.get('consecutive_stable_checks', 0)
|
|
self.last_failure_time = state_data.get('last_failure_time')
|
|
|
|
if remediation_data:
|
|
self.remediation_state = remediation_data.get('state')
|
|
self.remediation_start_time = remediation_data.get('start_time')
|
|
self.stabilization_checks = remediation_data.get('stabilization_checks', 0)
|
|
|
|
if stability_data:
|
|
self.stability_start_time = stability_data.get('start_time')
|
|
|
|
if vpn_data:
|
|
self.vpn_status = vpn_data.get('vpn_status', 'unknown')
|
|
self.last_vpn_status_change = vpn_data.get('last_vpn_status_change', time.time())
|
|
self.public_ip = vpn_data.get('public_ip')
|
|
self.last_public_ip_change = vpn_data.get('last_public_ip_change')
|
|
self.public_ip_details = vpn_data.get('public_ip_details', {})
|
|
|
|
self.logger.debug(f"Loaded state: connection={self.connection_state}, "
|
|
f"remediation={self.remediation_state}, "
|
|
f"failures={self.consecutive_failures}, "
|
|
f"vpn={self.vpn_status}")
|
|
|
|
def save_state(self):
|
|
"""Save current state to persistence if available"""
|
|
if self.consul_persistence and self.consul_persistence.is_available():
|
|
state_data = {
|
|
'connection_state': self.connection_state,
|
|
'last_state_change_time': self.last_state_change_time,
|
|
'consecutive_failures': self.consecutive_failures,
|
|
'consecutive_stable_checks': self.consecutive_stable_checks,
|
|
'last_failure_time': self.last_failure_time
|
|
}
|
|
|
|
remediation_data = {
|
|
'state': self.remediation_state,
|
|
'start_time': self.remediation_start_time,
|
|
'stabilization_checks': self.stabilization_checks
|
|
}
|
|
|
|
stability_data = {
|
|
'start_time': self.stability_start_time
|
|
}
|
|
|
|
# VPN state data
|
|
vpn_data = {
|
|
'vpn_status': self.vpn_status,
|
|
'last_vpn_status_change': self.last_vpn_status_change,
|
|
'public_ip': self.public_ip,
|
|
'last_public_ip_change': self.last_public_ip_change,
|
|
'public_ip_details': self.public_ip_details
|
|
}
|
|
|
|
return self.consul_persistence.save_state(state_data, remediation_data, stability_data, vpn_data)
|
|
return False
|
|
|
|
def reset_remediation_state(self):
|
|
"""Reset remediation state to initial values"""
|
|
self.remediation_state = None
|
|
self.remediation_start_time = None
|
|
self.stabilization_checks = 0
|
|
self.stability_start_time = None
|
|
self.save_state()
|
|
|
|
def start_remediation(self):
|
|
"""Start remediation process"""
|
|
self.remediation_state = 'stopping_torrents'
|
|
self.remediation_start_time = time.time()
|
|
self.stabilization_checks = 0
|
|
self.save_state()
|
|
|
|
def update_remediation_state(self, new_state: str):
|
|
"""Update remediation state"""
|
|
self.remediation_state = new_state
|
|
self.save_state()
|
|
|
|
def increment_stabilization_checks(self):
|
|
"""Increment stabilization check counter"""
|
|
self.stabilization_checks += 1
|
|
self.save_state()
|
|
|
|
def start_stability_timer(self):
|
|
"""Start stability timer"""
|
|
self.stability_start_time = time.time()
|
|
self.save_state()
|
|
|
|
def reset_stability_timer(self):
|
|
"""Reset stability timer"""
|
|
self.stability_start_time = None
|
|
self.save_state()
|
|
|
|
def update_vpn_status(self, vpn_status: str):
|
|
"""Update VPN status and track changes"""
|
|
if vpn_status != self.vpn_status:
|
|
self.vpn_status = vpn_status
|
|
self.last_vpn_status_change = time.time()
|
|
self.save_state()
|
|
|
|
def update_public_ip(self, public_ip: str, ip_details: Dict[str, Any]):
|
|
"""Update public IP and track changes"""
|
|
if public_ip != self.public_ip:
|
|
self.public_ip = public_ip
|
|
self.public_ip_details = ip_details
|
|
self.last_public_ip_change = time.time()
|
|
self.save_state()
|
|
|
|
def get_state_summary(self) -> Dict[str, Any]:
|
|
"""Get summary of current state"""
|
|
return {
|
|
'connection_state': self.connection_state,
|
|
'consecutive_failures': self.consecutive_failures,
|
|
'remediation_state': self.remediation_state,
|
|
'stability_start_time': self.stability_start_time,
|
|
'consul_available': self.consul_persistence.is_available() if self.consul_persistence else False
|
|
}
|
|
|
|
def get_debug_metrics(self) -> Dict[str, Any]:
|
|
"""Get debug metrics for logging in improved format with color support"""
|
|
current_time = time.time()
|
|
|
|
# ANSI color codes
|
|
GREEN = '\033[92m'
|
|
YELLOW = '\033[93m'
|
|
RED = '\033[91m'
|
|
BLUE = '\033[94m'
|
|
MAGENTA = '\033[95m'
|
|
CYAN = '\033[96m'
|
|
RESET = '\033[0m'
|
|
BOLD = '\033[1m'
|
|
|
|
# Helper function for compact duration formatting
|
|
def format_compact_duration(seconds: float) -> str:
|
|
if seconds == 0:
|
|
return "0s"
|
|
hours = int(seconds // 3600)
|
|
minutes = int((seconds % 3600) // 60)
|
|
secs = int(seconds % 60)
|
|
if hours > 0:
|
|
return f"{hours}h{minutes:02d}m{secs:02d}s"
|
|
elif minutes > 0:
|
|
return f"{minutes}m{secs:02d}s"
|
|
else:
|
|
return f"{secs}s"
|
|
|
|
# Build formatted metrics with colors
|
|
metrics = []
|
|
colored_metrics = []
|
|
|
|
# Connection state
|
|
state_duration = current_time - self.last_state_change_time
|
|
state_color = GREEN if self.connection_state == 'stable' else RED
|
|
metrics.append(f"Connection: {self.connection_state} ({format_compact_duration(state_duration)})")
|
|
colored_metrics.append(f"{BOLD}Connection:{RESET} {state_color}{self.connection_state} ({format_compact_duration(state_duration)}){RESET}")
|
|
|
|
# Failures
|
|
failure_color = RED if self.consecutive_failures > 0 else GREEN
|
|
failure_info = f"Failures: {self.consecutive_failures}"
|
|
last_failure_info = f" (Last: {GREEN}Never{RESET})"
|
|
if self.last_failure_time:
|
|
time_since_failure = current_time - self.last_failure_time
|
|
failure_recency_color = GREEN if time_since_failure > 300 else YELLOW # Green after 5min
|
|
last_failure_info = f" (Last: {failure_recency_color}{format_compact_duration(time_since_failure)} ago{RESET})"
|
|
metrics.append(f"{failure_info}{last_failure_info.replace(RESET, '').replace(GREEN, '').replace(YELLOW, '')}")
|
|
colored_metrics.append(f"{BOLD}Failures:{RESET} {failure_color}{failure_info}{RESET}{last_failure_info}")
|
|
|
|
# Remediation
|
|
remediation_color = YELLOW if self.remediation_state else GREEN
|
|
remediation_info = f"Remediation: {self.remediation_state or 'None'}"
|
|
if self.remediation_state:
|
|
remediation_duration = current_time - self.remediation_start_time
|
|
remediation_info += f" ({format_compact_duration(remediation_duration)})"
|
|
metrics.append(remediation_info)
|
|
colored_metrics.append(f"{BOLD}Remediation:{RESET} {remediation_color}{remediation_info}{RESET}")
|
|
|
|
# Stability timer
|
|
stability_info = "Stability: Not active"
|
|
if self.stability_start_time:
|
|
elapsed_stable = current_time - self.stability_start_time
|
|
stability_color = GREEN if elapsed_stable >= 1800 else YELLOW # Green after 30min
|
|
stability_info = f"Stability: {stability_color}{format_compact_duration(elapsed_stable)}{RESET}"
|
|
else:
|
|
stability_info = f"{CYAN}Stability: Not active{RESET}"
|
|
metrics.append(stability_info.replace(RESET, '').replace(GREEN, '').replace(YELLOW, '').replace(CYAN, ''))
|
|
colored_metrics.append(f"{BOLD}Stability:{RESET} {stability_info}")
|
|
|
|
# VPN status
|
|
vpn_color = GREEN if self.vpn_status == 'running' else RED
|
|
vpn_info = f"VPN Uptime: {format_compact_duration(current_time - self.last_vpn_status_change)}"
|
|
metrics.append(vpn_info)
|
|
colored_metrics.append(f"{BOLD}VPN:{RESET} {vpn_color}{vpn_info}{RESET}")
|
|
|
|
# Public IP
|
|
ip_color = GREEN if self.public_ip and self.public_ip != 'unknown' else YELLOW
|
|
ip_info = f"IP: {self.public_ip or 'unknown'}"
|
|
if self.public_ip and self.last_public_ip_change:
|
|
ip_duration = current_time - self.last_public_ip_change
|
|
ip_info += f" ({format_compact_duration(ip_duration)})"
|
|
metrics.append(ip_info)
|
|
colored_metrics.append(f"{BOLD}IP:{RESET} {ip_color}{ip_info}{RESET}")
|
|
|
|
return {
|
|
'multiline': "\n ".join([""] + colored_metrics), # For hierarchical format with colors
|
|
'compact': " | ".join(metrics) # For single-line format without colors
|
|
}
|
|
|
|
def _format_duration(self, seconds: float) -> str:
|
|
"""Format duration in human-readable format"""
|
|
from utils.time_utils import format_human_readable_time
|
|
return format_human_readable_time(seconds) |