55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
from enum import Enum
|
|
|
|
from colony.core.storages.default import DefaultStorage
|
|
from colony.core.roles.replicator.role import Replicator
|
|
from colony.core.roles.role import Role, Command
|
|
|
|
|
|
class VaultWarden(Role):
|
|
REPLICATE_TIMEOUT = 3
|
|
|
|
class COMMANDS(Enum):
|
|
LS = 'VaultWarden.ls'
|
|
OPEN = 'VaultWarden.open'
|
|
OVERWRITE = 'VaultWarden.overwrite'
|
|
|
|
def __init__(self, context):
|
|
commands = {
|
|
self.COMMANDS.LS.value: Command(self.ls),
|
|
self.COMMANDS.OPEN.value: Command(self.open),
|
|
self.COMMANDS.OVERWRITE.value: Command(self.overwrite)
|
|
}
|
|
|
|
self.storage = DefaultStorage()
|
|
|
|
super().__init__(commands, context)
|
|
|
|
def ls(self, dir):
|
|
return self.storage.ls(dir)
|
|
|
|
def open(self, file):
|
|
return self.storage.open(file)
|
|
|
|
def overwrite(self, filename, content):
|
|
path = self.storage.save(filename, content)
|
|
if self.storage.SHARED:
|
|
self._push_to_replicators(filename, content)
|
|
return path
|
|
|
|
def _push_to_replicators(self, filename, content):
|
|
for record in self.context.membership().alive_records():
|
|
if Replicator.__name__ not in record.roles:
|
|
continue
|
|
target = f"{record.ip}:{record.port}"
|
|
try:
|
|
resp = self.context.execute_to_unregistered_node(
|
|
record.ip, record.port, Replicator.COMMANDS.REPLICATE, filename, content,
|
|
timeout=self.REPLICATE_TIMEOUT,
|
|
)
|
|
except Exception as e:
|
|
self.logger.warning(f"Could not replicate {filename} to {target}: {e}")
|
|
continue
|
|
if not isinstance(resp, dict) or resp.get("status") != "ok":
|
|
error = resp.get("error") if isinstance(resp, dict) else "bad reply"
|
|
self.logger.warning(f"{target} did not accept {filename}: {error}")
|