113 lines
3.9 KiB
Python
113 lines
3.9 KiB
Python
import json
|
|
from enum import Enum
|
|
from typing import Optional
|
|
|
|
from colony.core.network.nodes import LocalNode, Node
|
|
from colony.core.network.requests import CommandRequest
|
|
from colony.core.logger.logging import get_logger
|
|
from colony.settings import configuration
|
|
from colony.transport.whisperer import Whisperer
|
|
from colony.utils.network import get_local_ip, get_hostname
|
|
|
|
|
|
class Network:
|
|
def __init__(self):
|
|
self.logger = get_logger(self.__class__.__name__)
|
|
|
|
self.self_node = LocalNode(
|
|
hostname=get_hostname(),
|
|
ip=get_local_ip(),
|
|
roles=[],
|
|
port=configuration.listener.port
|
|
)
|
|
self.logger.info(f"Self node created: {self.self_node.id}")
|
|
|
|
self.nodes = [self.self_node] + [
|
|
node for node in self.load_network_from_snapshot()
|
|
if (node.ip, node.port) != (self.self_node.ip, self.self_node.port)
|
|
]
|
|
|
|
def get_self(self):
|
|
return self.self_node
|
|
|
|
def get_all_nodes(self):
|
|
return self.nodes
|
|
|
|
def get_node_of_role(self, role_name: str) -> Optional[Node]:
|
|
for node in self.nodes:
|
|
for role in node.roles:
|
|
if role == role_name:
|
|
return node
|
|
|
|
self.logger.warning(f"No node found with role: {role_name}")
|
|
return None
|
|
|
|
def execute(self, command: Enum, *args, **kwargs):
|
|
request = CommandRequest.form(command, *args, **kwargs)
|
|
self.logger.info(f"Executing command {request.id} on role {request.role}")
|
|
|
|
node = self.get_node_of_role(request.role)
|
|
if node is None:
|
|
raise ValueError(f"No node found with role {request.role}")
|
|
|
|
result = node.execute(command, *args, **kwargs)
|
|
self.logger.debug(f"Command {request.id} executed")
|
|
return result
|
|
|
|
|
|
@staticmethod
|
|
def execute_to_unregistered_node(ip, port, command: Enum, *args, **kwargs):
|
|
request = CommandRequest.form(command, *args, **kwargs)
|
|
|
|
return Whisperer(ip, port).send(request.collapse())
|
|
|
|
|
|
def load_network_from_snapshot(self, file_path: str = "storage/colony_snapshot.json"):
|
|
self.logger.info(f"Loading network snapshot from {file_path}")
|
|
nodes = []
|
|
try:
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
snapshot = json.load(f)
|
|
except FileNotFoundError:
|
|
self.logger.warning(f"Snapshot file {file_path} not found. Starting with empty network")
|
|
return nodes
|
|
except (json.JSONDecodeError, OSError) as e:
|
|
self.logger.warning(f"Could not read snapshot {file_path}: {e}")
|
|
return nodes
|
|
|
|
if not isinstance(snapshot, dict):
|
|
self.logger.warning(f"Snapshot {file_path} is not an object. Ignoring it")
|
|
return nodes
|
|
|
|
for node_ip, node_data in snapshot.items():
|
|
if not isinstance(node_data, dict):
|
|
self.logger.warning(f"Skipping node {node_ip}: entry is not a dict")
|
|
continue
|
|
|
|
if node_data.get("status") == "unreachable":
|
|
self.logger.info(f"Skipping node {node_ip}: last seen unreachable")
|
|
continue
|
|
|
|
result = node_data.get("result")
|
|
|
|
if not isinstance(result, dict):
|
|
self.logger.warning(f"Skipping node {node_ip}: result is not a dict")
|
|
continue
|
|
|
|
roles = result.get("roles", [])
|
|
try:
|
|
node = Node(
|
|
hostname=result["hostname"],
|
|
ip=result["ip"],
|
|
port=result["port"],
|
|
roles=roles if isinstance(roles, list) else []
|
|
)
|
|
except KeyError as e:
|
|
self.logger.warning(f"Skipping node {node_ip}: result has no {e}")
|
|
continue
|
|
|
|
self.logger.info(f"Loaded node from snapshot: {node.id}")
|
|
nodes.append(node)
|
|
|
|
return nodes
|