ironic-python-agent/teeth_agent/client.py

"""
Copyright 2013 Rackspace, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import time
import json
import random
import tempfile

from twisted.application.service import MultiService
from twisted.application.internet import TCPClient
from twisted.internet.protocol import ReconnectingClientFactory
from twisted.internet.defer import maybeDeferred
from twisted.python.failure import Failure

from teeth_agent import __version__ as AGENT_VERSION
from teeth_agent.protocol import TeethAgentProtocol
from teeth_agent.logging import get_logger
log = get_logger()


__all__ = ["TeethClientFactory", "TeethClient"]


class TeethClientFactory(ReconnectingClientFactory, object):
    """
    Protocol Factory for the Teeth Client.
    """
    protocol = TeethAgentProtocol
    initialDelay = 1.0
    maxDelay = 120

    def __init__(self, encoder, parent):
        super(TeethClientFactory, self).__init__()
        self._encoder = encoder
        self._parent = parent

    def buildProtocol(self, addr):
        """Create protocol for an address."""
        self.resetDelay()
        proto = self.protocol(self._encoder, addr, self._parent)
        self._parent.add_protocol_instance(proto)
        return proto

    def clientConnectionFailed(self, connector, reason):
        """clientConnectionFailed"""
        log.err('Failed to connect, re-trying', delay=self.delay)
        super(TeethClientFactory, self).clientConnectionFailed(connector, reason)

    def clientConnectionLost(self, connector, reason):
        """clientConnectionLost"""
        log.err('Lost connection, re-connecting', delay=self.delay)
        super(TeethClientFactory, self).clientConnectionLost(connector, reason)


class TeethClient(MultiService, object):
    """
    High level Teeth Client.
    """
    client_factory_cls = TeethClientFactory
    client_encoder_cls = json.JSONEncoder

    def __init__(self, addrs):
        super(TeethClient, self).__init__()
        self.setName('teeth-agent')
        self._client_encoder = self.client_encoder_cls()
        self._client_factory = self.client_factory_cls(self._client_encoder, self)
        self._log = get_logger()
        self._start_time = time.time()
        self._protocols = []
        self._outmsg = []
        self._connectaddrs = addrs
        self._handlers = {
            'v1': {
                'status': self._handle_status,
            }
        }

    @property
    def conf_image_cache_path(self):
        """Path to iamge cache."""
        # TODO: improve:
        return tempfile.gettempdir()

    def startService(self):
        """Start the Service."""
        super(TeethClient, self).startService()

        for host, port in self._connectaddrs:
            service = TCPClient(host, port, self._client_factory)
            service.setName("teeth-agent[%s:%d]".format(host, port))
            self.addService(service)
        self._connectaddrs = []

    def remove_endpoint(self, host, port):
        """Remove an Agent Endpoint from the active list."""

        def op(protocol):
            if protocol.address.host == host and protocol.address.port == port:
                protocol.loseConnectionSoon()
                return True
            return False
        self._protocols[:] = [protocol for protocol in self._protocols if not op(protocol)]

    def add_endpoint(self, host, port):
        """Add an agent endpoint to the """
        self._connectaddrs.append([host, port])
        self.start()

    def add_protocol_instance(self, protocol):
        """Add a running protocol to the parent."""
        protocol.on('command', self._on_command)
        self._protocols.append(protocol)

    def _on_command(self, topic, message):
        if message.version not in self._handlers:
            message.protocol.fatal_error('unknown message version')
            return

        if message.method not in self._handlers[message.version]:
            message.protocol.fatal_error('unknown message method')
            return

        handler = self._handlers[message.version][message.method]
        d = maybeDeferred(handler, message)
        d.addBoth(self._send_response, message)

    def _send_response(self, result, message):
        """Send a response to a message."""
        if isinstance(result, Failure):
            self._log.err(result)
            message.protocol.send_error_response(result.value.message, message)
        else:
            message.protocol.send_response(result, message)

    def _handle_status(self, command):
        return {
            'mode': self.AGENT_MODE,
            'uptime': time.time() - self._start_time,
            'version': AGENT_VERSION,
        }

    def _add_handler(self, version, command, func):
        self._handlers[version][command] = func

    def _send_command(self, method, params):
        protocol = random.choice(self._protocols)
        return protocol.send_command(method, params)

    def send_log(self, message, **kwargs):
        """
        Send a log message to the endpoint.
        """
        event = {}
        event.update(kwargs)
        event['message'] = message
        event['time'] = time.time()
        return self._send_command('log', event)