def test_key_exchange(self):
"""
Tests the auth function before the client has paired. This triggers
key exchanging.
Verifies:
* remote command "exchange_keys" is sent
* Avatar response is received and decoded
* response triggers client save_key to be called
* server key is received
* response triggers server save_key to be called
* client key is received
"""
client = RSAClient(self.priv_key, self.pub_key_values)
remote = RemoteProxy()
save_key_server = CallProxy(None, False)
save_key_client = CallProxy(None, False)
client.auth(remote, save_key=save_key_client)
args, kwargs, deferred = remote.assertCalled(self, 'exchange_keys')
avatar = RSAAvatar(self.priv_key, self.pub_key_values, self.pub_key, save_key=save_key_server, key_size=KEY_SIZE)
deferred.callback(avatar.perspective_exchange_keys(*args[1:]))
args, kwargs = save_key_server.assertCalled(self)
key = simplejson.loads(''.join(args[0]))
self.assert_(key==self.pub_key_values, 'keys do not match')
args, kwargs = save_key_client.assertCalled(self)
key = simplejson.loads(''.join(args[0]))
self.assert_(key==self.pub_key_values, 'keys do not match')
def test_auth(self):
"""
Tests the auth function
"""
client = RSAClient(self.priv_key)
remote = RemoteProxy()
client.auth(remote, server_key=self.pub_key)
self.assertEqual(remote.func, 'auth_challenge', 'Calling auth should trigger auth_challenge call on server')
def auth_result_response_before_challenge(self):
"""
Tests responding to a failure due to response before challenge error
"""
client = RSAClient(self.priv_key)
remote = RemoteProxy()
client.auth_result(0, remote)
#verify that auth_response got called
self.assertEqual(remote.func, 'auth_challenge', 'Calling auth_response before auth_challegne should trigger auth_response call on server')
def test_auth(self):
"""
Tests the auth function
Verifies:
* auth_challenge is sent
"""
client = RSAClient(self.priv_key)
remote = RemoteProxy()
client.auth(remote, server_key=self.pub_key)
remote.assertCalled(self, 'auth_challenge')
def test_auth_challenge_no_challenge(self):
"""
Tests auth_challenge when the challenge received is None
"""
client = RSAClient(self.priv_key)
remote = RemoteProxy()
challenge = None
client.auth_challenge(challenge, remote, self.pub_key)
#verify that auth_response got called
self.assertEqual(remote.func, 'auth_response', 'Calling auth_challenge should trigger auth_response call on server')
#verify the correct response was sent
self.assertFalse(remote.kwargs['response'], 'Response did not match the expected response')
def __init__(self):
#locks
self._lock = Lock() #general lock, use when multiple shared resources are touched
self._listeners = {
'MANAGER_INIT':self.connect,
'NODE_CREATED':self.connect_node,
'NODE_UPDATED':self.connect_node,
}
self._interfaces = [
self.connect
]
#load rsa crypto
self.pub_key, self.priv_key = load_crypto('%s/master.key' \
% pydra_settings.RUNTIME_FILES_DIR)
self.rsa_client = RSAClient(self.priv_key, self.pub_key, \
callback=self.init_node)
#connection management
self.connecting = True
self.reconnect_count = 0
self.attempts = None
self.reconnect_call_ID = None
self.host = 'localhost'
def test_auth_challenge_no_server_key(self):
"""
Tests auth_challenge when server key is None.
"""
client = RSAClient(self.priv_key)
avatar = RSAAvatar(self.priv_key, None, self.pub_key, key_size=KEY_SIZE)
remote = RemoteProxy()
challenge = avatar.perspective_auth_challenge()
client.auth_challenge(challenge, remote, None)
#verify that auth_response got called
self.assertEqual(remote.func, 'auth_response', 'Calling auth_challenge should trigger auth_response call on server')
#verify the correct response was sent
self.assertFalse(remote.kwargs['response'], 'Response did not match the expected response')
def test_auth_challenge(self):
"""
Tests a normal challenge string
Verifies:
* remote call is sent
* response equals challenge
"""
client = RSAClient(self.priv_key)
avatar = RSAAvatar(self.priv_key, None, self.pub_key, key_size=KEY_SIZE)
remote = RemoteProxy()
challenge = avatar.perspective_auth_challenge()
client.auth_challenge(challenge, remote, self.pub_key)
#verify that auth_response got called
args, kwargs, deferred = remote.assertCalled(self, 'auth_response')
self.assertEqual(kwargs['response'], avatar.challenge, 'Response did not match the expected response')
def test_auth_challenge_no_challenge(self):
"""
Tests auth_challenge when the challenge received is None
Verifies:
* remote call is sent
* auth is denied
"""
client = RSAClient(self.priv_key)
remote = RemoteProxy()
challenge = None
client.auth_challenge(challenge, remote, self.pub_key)
#verify that auth_response got called
args, kwargs, deferred = remote.assertCalled(self, 'auth_response')
#verify the correct response was sent
self.assertFalse(kwargs['response'], 'Response did not match the expected response')
def test_auth_challenge_no_server_key(self):
"""
Tests auth_challenge when server key is None.
Verifies:
* remote call is sent
* auth is denied
"""
client = RSAClient(self.priv_key)
avatar = RSAAvatar(self.priv_key, None, self.pub_key, key_size=KEY_SIZE)
remote = RemoteProxy()
challenge = avatar.perspective_auth_challenge()
client.auth_challenge(challenge, remote, None)
#verify that auth_response got called
args, kwargs, deferred = remote.assertCalled(self, 'auth_response')
#verify the correct response was sent
self.assertFalse(kwargs['response'], 'Response did not match the expected response')
def __init__(self):
self._listeners = {"MANAGER_INIT": self.connect, "WORKER_FINISHED": self.disconnect}
self._remotes = [("MASTER", os.getpid)]
self.reconnect_count = 0
self.factory = MasterClientFactory(self.reconnect)
# load crypto for authentication
# workers use the same keys as their parent Node
self.pub_key, self.priv_key = load_crypto("%s/node.key" % pydra_settings.RUNTIME_FILES_DIR)
# self.master_pub_key = load_crypto('./node.key', False, both=False)
self.rsa_client = RSAClient(self.priv_key)
class NodeConnectionManager(Module):
_signals = [
'WORKER_CONNECTED',
'WORKER_DISCONNECTED',
'NODE_CONNECTED',
'NODE_DISCONNECTED',
'NODE_AUTHENTICATED'
]
_shared = [
'nodes',
'known_nodes',
'workers'
]
def __init__(self):
#locks
self._lock = Lock() #general lock, use when multiple shared resources are touched
self._listeners = {
'MANAGER_INIT':self.connect,
'NODE_CREATED':self.connect_node,
'NODE_UPDATED':self.connect_node,
}
self._interfaces = [
self.connect
]
#load rsa crypto
self.pub_key, self.priv_key = load_crypto('%s/master.key' \
% pydra_settings.RUNTIME_FILES_DIR)
self.rsa_client = RSAClient(self.priv_key, self.pub_key, \
callback=self.init_node)
#connection management
self.connecting = True
self.reconnect_count = 0
self.attempts = None
self.reconnect_call_ID = None
self.host = 'localhost'
def _register(self, manager):
Module._register(self, manager)
#cluster management
self.nodes = self.load_nodes()
self.workers = {}
def load_nodes(self):
"""
Load node configuration from the database
"""
logger.info('loading nodes')
nodes = Node.objects.exclude(deleted=True)
node_dict = {}
for node in nodes:
node_dict[node.id] = node
logger.info('%i nodes loaded' % len(nodes))
return node_dict
def connect_node(self, node):
"""
Callback for node creation and updates. Connect() cannot be called
directly because those signals include a parameter
"""
self.connect()
def connect(self):
"""
Make connections to all Nodes that are not connected. This method is a single control
for connecting to nodes. individual nodes cannot be connected to. This is to ensure that
only one attempt at a time is ever made to connect to a node.
"""
#lock for two reasons:
# 1) connect() cannot be called more than once at a time
# 2) if a node fails while connecting the reconnect call will block till
# connections are finished
with self._lock:
self.connecting=True
# make sure the various states are in sync
for i in Node.objects.exclude(deleted=True):
if i.id not in self.nodes:
self.nodes[i.id] = i
if (i.host, i.port) in self.known_nodes:
self.known_nodes.discard((i.host, i.port))
logger.info("Connecting to nodes")
connections = []
self.attempts = []
for id, node in self.nodes.items():
#only connect to nodes that aren't connected yet
if not node.ref:
factory = NodeClientFactory(node, self)
reactor.connectTCP(node.host, node.port, factory)
# SSH authentication is not currently supported with perspectiveBroker.
# For now we'll perform a key handshake within the info/init handshake that already
# occurs. Prior to the handshake completing access will be limited to info which
# is informational only.
#
# Note: The first time connecting to the node will accept and register whatever
# key is passed to it. This is a small trade of in temporary insecurity to simplify
# this can be avoided by manually generating and setting the keys
#
#credential = credentials.SSHPrivateKey('master', 'RSA', node.pub_key, '', '')
credential = credentials.UsernamePassword('master', '1234')
# construct referenceable with remotes for NODE
client = ModuleReferenceable(self.manager._remotes['NODE'])
deferred = factory.login(credential, client=client)
connections.append(deferred)
self.attempts.append(node)
defer.DeferredList(connections, consumeErrors=True).addCallbacks(
self.nodes_connected, errbackArgs=("Failed to Connect"))
# Release the connection flag.
self.connecting=False
def nodes_connected(self, results):
"""
Called with the results of all connection attempts. Store connections and retrieve info from node.
The node will respond with info including how many workers it has.
"""
# process each connected node
failures = False
for result, node in zip(results, self.attempts):
#successes
if result[0]:
# save reference for remote calls
node.ref = result[1]
d = node.ref.callRemote('get_key')
d.addCallback(self.check_node, node)
#failures
else:
logger.error('node:%s:%s - failed to connect' % (node.host, node.port))
node.ref = None
failures = True
#single call to reconnect for all failures
if failures:
self.reconnect_nodes()
else:
self.reconnect_count = 0
def check_node(self, key, node):
# node.pub_key is set only for paired nodes, make sure we don't attempt
# to pair with a known pub key
duplicate = ''.join(key) in [i.pub_key for i in self.nodes.values()]
if duplicate and not node.pub_key:
logger.info('deleting %s:%s - duplicate' % (node.host, node.port))
node.delete()
return
# Authenticate with the node
pub_key = node.load_pub_key()
self.rsa_client.auth(node.ref, self.receive_key_node, server_key=pub_key, node=node)
logger.info('node:%s:%s - connected' % (node.host, node.port))
def reconnect_nodes(self, reset_counter=False):
"""
Called to signal that a reconnection attempt is needed for one or more nodes. This is the single control
for requested reconnection. This single control is used to ensure at most
one request for reconnection is pending.
"""
#lock - Blocking here ensures that connect() cannot happen while requesting
# a reconnect.
with self._lock:
#reconnecting flag ensures that connect is only called a single time
#it's possible that multiple nodes can have problems at the same time
#reset_counter overrides this
if not self.connecting or reset_counter:
self.connecting = True
#reset the counter, useful when a new failure occurs
if reset_counter:
#cancel existing call if any
if self.reconnect_call_ID:
try:
self.reconnect_call_ID.cancel()
# There is a slight chance that this method can be called
# and receive the lock, after connect() has been called.
# in that case reconnect_call_ID will point to an already called
# item. The error can just be ignored as the locking will ensure
# the call we are about to make does not start
# until the first one does.
except AlreadyCalled:
pass
self.reconnect_count = 0
reconnect_delay = 5*pow(2, self.reconnect_count)
#let increment grow exponentially to 5 minutes
if self.reconnect_count < 6:
self.reconnect_count += 1
logger.debug('reconnecting in %i seconds' % reconnect_delay)
self.reconnect_call_ID = reactor.callLater(reconnect_delay, self.connect)
def receive_key_node(self, key, node=None, **kwargs):
"""
Receives the public key from the node
"""
logger.debug("saving public key from node: %s" % node)
node.pub_key = key
node.save()
def init_node(self, node):
"""
Start the initialization sequence with the node. The first
step is to query it for its information.
"""
d = node.ref.callRemote('info')
d.addCallback(self.add_node, node=node)
def add_node(self, info, node):
"""
Process Node information. Most will just be stored for later use. Info will include
a list of workers. The master will then connect to all Workers.
"""
# save node's information in the database
node.total_memory = info['total_memory']
node.avail_memory = info['avail_memory']
node.cores = info['cores']
node.stones = info['stones']
node.save()
#node key to be used by node and its workers
node_key_str = '%s:%s' % (node.host, node.port)
# create worker avatars
with self._lock:
for i in range(node.cores):
worker_key = '%s:%i' % (node_key_str, i)
avatar = WorkerAvatarWrapper(worker_key, node.ref)
self.workers[worker_key] = avatar
self.emit('WORKER_CONNECTED', avatar)
# we have allowed access for all the workers, tell the node to init
d = node.ref.callRemote('init_node', self.host, pydra_settings.PORT, node_key_str)
d.addCallback(self.node_ready, node)
def node_ready(self, result, node):
"""
Called when a call to initialize a Node is successful
"""
logger.info('node:%s - ready' % node)
class WorkerConnectionManager(Module):
"""
Module that manages connection with the Master for a Worker
"""
_signals = ["MASTER_CONNECTED", "MASTER_DISCONNECTED"]
_shared = ["worker_key", "master", "master_port", "_lock_connection"]
def __init__(self):
self._listeners = {"MANAGER_INIT": self.connect, "WORKER_FINISHED": self.disconnect}
self._remotes = [("MASTER", os.getpid)]
self.reconnect_count = 0
self.factory = MasterClientFactory(self.reconnect)
# load crypto for authentication
# workers use the same keys as their parent Node
self.pub_key, self.priv_key = load_crypto("%s/node.key" % pydra_settings.RUNTIME_FILES_DIR)
# self.master_pub_key = load_crypto('./node.key', False, both=False)
self.rsa_client = RSAClient(self.priv_key)
def _register(self, manager):
Module._register(self, manager)
self._lock_connection = Lock()
def connect(self):
"""
Make initial connections to all Nodes
"""
logger.info("worker:%s - connecting to Node @ %s:%s" % (self.worker_key, "localhost", self.master_port))
reactor.connectTCP("localhost", self.master_port, self.factory)
# construct referenceable with remotes for MASTER
client = ModuleReferenceable(self.manager._remotes["MASTER"])
deferred = self.factory.login(credentials.UsernamePassword(self.worker_key, "1234"), client=client)
deferred.addCallbacks(self.connected, self.reconnect, errbackArgs=("Failed to Connect"))
def disconnect(self):
if self.factory:
self.factory.disconnecting = True
self.factory.disconnect()
def reconnect(self, *arg, **kw):
with self._lock_connection:
self.master = None
reconnect_delay = 5 * pow(2, self.reconnect_count)
# let increment grow exponentially to 5 minutes
if self.reconnect_count < 6:
self.reconnect_count += 1
logger.debug("worker:%s - reconnecting in %i seconds" % (self.worker_key, reconnect_delay))
self.reconnect_call_ID = reactor.callLater(reconnect_delay, self.connect)
def connected(self, result):
"""
Callback called when connection to master is made
"""
with self._lock_connection:
self.master = result
self.reconnect_count = 0
logger.info("worker:%s - connected to master @ %s:%s" % (self.worker_key, "localhost", self.master_port))
# Authenticate with the master
self.rsa_client.auth(result, None, self.priv_key)
def connect_failed(self, result):
"""
Callback called when conenction to master fails
"""
self.reconnect()
