<reponame>yatch/opentestbed-cli
import errno
import json
import logging
import queue
import socket
import sys
import time

import paho.mqtt.client as mqtt

from cli.utils import get_default_config

logger = logging.getLogger(__name__)

class MQTTClient(object):
    RESP_TIMEOUT_SECONDS = 5

    def __init__(self, config={}):
        # these attributes must be defined by child classes
        self.device_id = None
        self.device_type = None

        if config:
            self.config = config
        else:
            self.config = get_default_config()
        self.queue = queue.SimpleQueue()

        # shorthand
        mqtt_broker_config = self.config['opentestbed']['mqtt_broker']
        self.broker_address = mqtt_broker_config['address']
        self.broker_port = mqtt_broker_config['port']
        self.broker = '{0}:{1}'.format(self.broker_address, self.broker_port)

        if hasattr(sys, '_called_from_test'):
            # for offline testing, skip establishing a MQTT connection
            pass
        else:
            self.mqttc = mqtt.Client(userdata=self)
            self.mqttc.enable_logger(logging.getLogger('paho-mqtt'))
            self.mqttc.on_connect = self._on_connect
            self.mqttc.on_disconnect = self._on_disconnect
            self.mqttc.on_message = self._on_message

            try:
                logger.debug('Connecting to {0}'.format(self.broker))
                self.mqttc.connect(self.broker_address, self.broker_port)
                self.mqttc.loop_start()
            except OSError as err:
                if (isinstance(err, socket.gaierror)
                    and err.errno == errno.ENOEXEC):
                    err_str = ('Cannot resolve {0} '.format(self.broker_address)
                               + 'or {0}'.format(self.broker_port))
                elif (isinstance(err, socket.error)
                      and err.errno == errno.ETIMEDOUT):
                    err_str = 'Connection timeout with {0}'.format(self.broker)
                elif err.errno == errno.ENETUNREACH:
                    err_str = 'Cannot connect to {0}'.format(self.broker)
                else:
                    err_str = None
                if err_str:
                    logger.critical(err_str)
                    exit(1)
                else:
                    raise

    def __del__(self):
        logger.info('The instance for {0} '.format(self.device_id)
                    + 'is being destroyed')
        if hasattr(self, 'mqtcc') and self.mqttc:
            # disconnect from the broker and stop the loop
            self.mqttc.disconnect()
            self.mqttc.loop_stop()

        # if the parent class has __del__, call it ('object' doesn't,
        # but check it just in case for future enhancements)
        if hasattr(super(), '__del__'):
            super().__del__()

    @staticmethod
    def _on_connect(client, userdata, flags, rc):
        self = userdata
        logger.info('Connection for "{0}" '.format(self.device_id)
                    + 'established with {0}'.format(self.broker))

    @staticmethod
    def _on_disconnect(client, userdata, rc):
        self = userdata
        logger.info('Connection for "{0}" '.format(self.device_id)
                    + 'disconnected from {0}'.format(self.broker))

    @staticmethod
    def _on_message(client, userdata, msg):
        self = userdata
        message = {'topic': msg.topic,
                   'payload': json.loads(msg.payload.decode('utf-8'))}
        logger.debug('Receive {0} '.format(message['topic'])
                     + 'on {0}'.format(message['payload']))
        self.queue.put(message)

    def _cmd_topic(self, command):
        return self._topic('cmd', command)

    def _resp_topic(self, command):
        return self._topic('resp', command)

    def _topic(self, topic_type, command):
        if self.device_id == 'all' and topic_type == 'resp':
            device_id = '+'
        else:
            device_id = self.device_id

        return ('opentestbed/deviceType/{0}'.format(self.device_type)
                +'/deviceId/{0}/'.format(device_id)
                +'{0}/{1}'.format(topic_type, command))

    def _execute(self, command, payload={}, timeout=RESP_TIMEOUT_SECONDS):
        while True:
            try:
                self.queue.get_nowait()
            except queue.Empty:
                break
        assert self.queue.empty()

        if 'token' not in payload:
            payload['token'] = self.config['opentestbed']['token']

        self.mqttc.subscribe(self._resp_topic(command))
        logger.debug('Publishing {0} to {1}'.format(payload,
                                                    self._cmd_topic(command)))
        self.mqttc.publish(self._cmd_topic(command), json.dumps(payload))

        time_left = timeout
        end_time = time.time() + timeout

        logger.debug('Waiting on {0} '.format(self._resp_topic(command))
                     + 'for {0} seconds'.format(timeout))

        ret = [] if self.device_id == 'all' else None
        while time_left > 0:
            try:
                resp = self.queue.get(timeout=time_left)
            except queue.Empty:
                break

            if resp:
                # check the topic
                _resp_topic = resp['topic'].split('/')
                device_id_index = _resp_topic.index('deviceId') + 1
                device_id = _resp_topic[device_id_index]
                if self.device_id == 'all':
                    _resp_topic[device_id_index] = '+'
                assert '/'.join(_resp_topic) == self._resp_topic(command)
                payload = resp['payload']

                if payload['success']:
                    if (payload['returnVal']
                        and (payload['returnVal'] is not True)):
                        return_value = payload['returnVal']
                    else:
                        return_value = 'Success'
                    logger.debug('Receive from {0}, '.format(device_id)
                                 + '{0}'.format(return_value))
                else:
                    logger.error('Command {0} to '.format(command)
                                 + '{0} failed'.format(device_id))
                    return_value = None
            else:
                # timeout without ret
                logger.error('Reception timeout on '
                             + '{0}'.format(self._resp_topic(command)))
                return_value = None

            if self.device_id == 'all':
                time_left = end_time - time.time()
                if return_value:
                    ret.append((device_id, return_value))
                else:
                    # do nothing
                    pass
            else:
                # done
                ret = return_value
                break

        self.mqttc.unsubscribe(self._resp_topic(command))
        return ret

import errno
import json
import logging
import queue
import socket
import sys
import time

import paho.mqtt.client as mqtt

from cli.utils import get_default_config

logger = logging.getLogger(__name__)

class MQTTClient(object):
    RESP_TIMEOUT_SECONDS = 5

    def __init__(self, config={}):
        # these attributes must be defined by child classes
        self.device_id = None
        self.device_type = None

        if config:
            self.config = config
        else:
            self.config = get_default_config()
        self.queue = queue.SimpleQueue()

        # shorthand
        mqtt_broker_config = self.config['opentestbed']['mqtt_broker']
        self.broker_address = mqtt_broker_config['address']
        self.broker_port = mqtt_broker_config['port']
        self.broker = '{0}:{1}'.format(self.broker_address, self.broker_port)

        if hasattr(sys, '_called_from_test'):
            # for offline testing, skip establishing a MQTT connection
            pass
        else:
            self.mqttc = mqtt.Client(userdata=self)
            self.mqttc.enable_logger(logging.getLogger('paho-mqtt'))
            self.mqttc.on_connect = self._on_connect
            self.mqttc.on_disconnect = self._on_disconnect
            self.mqttc.on_message = self._on_message

            try:
                logger.debug('Connecting to {0}'.format(self.broker))
                self.mqttc.connect(self.broker_address, self.broker_port)
                self.mqttc.loop_start()
            except OSError as err:
                if (isinstance(err, socket.gaierror)
                    and err.errno == errno.ENOEXEC):
                    err_str = ('Cannot resolve {0} '.format(self.broker_address)
                               + 'or {0}'.format(self.broker_port))
                elif (isinstance(err, socket.error)
                      and err.errno == errno.ETIMEDOUT):
                    err_str = 'Connection timeout with {0}'.format(self.broker)
                elif err.errno == errno.ENETUNREACH:
                    err_str = 'Cannot connect to {0}'.format(self.broker)
                else:
                    err_str = None
                if err_str:
                    logger.critical(err_str)
                    exit(1)
                else:
                    raise

    def __del__(self):
        logger.info('The instance for {0} '.format(self.device_id)
                    + 'is being destroyed')
        if hasattr(self, 'mqtcc') and self.mqttc:
            # disconnect from the broker and stop the loop
            self.mqttc.disconnect()
            self.mqttc.loop_stop()

        # if the parent class has __del__, call it ('object' doesn't,
        # but check it just in case for future enhancements)
        if hasattr(super(), '__del__'):
            super().__del__()

    @staticmethod
    def _on_connect(client, userdata, flags, rc):
        self = userdata
        logger.info('Connection for "{0}" '.format(self.device_id)
                    + 'established with {0}'.format(self.broker))

    @staticmethod
    def _on_disconnect(client, userdata, rc):
        self = userdata
        logger.info('Connection for "{0}" '.format(self.device_id)
                    + 'disconnected from {0}'.format(self.broker))

    @staticmethod
    def _on_message(client, userdata, msg):
        self = userdata
        message = {'topic': msg.topic,
                   'payload': json.loads(msg.payload.decode('utf-8'))}
        logger.debug('Receive {0} '.format(message['topic'])
                     + 'on {0}'.format(message['payload']))
        self.queue.put(message)

    def _cmd_topic(self, command):
        return self._topic('cmd', command)

    def _resp_topic(self, command):
        return self._topic('resp', command)

    def _topic(self, topic_type, command):
        if self.device_id == 'all' and topic_type == 'resp':
            device_id = '+'
        else:
            device_id = self.device_id

        return ('opentestbed/deviceType/{0}'.format(self.device_type)
                +'/deviceId/{0}/'.format(device_id)
                +'{0}/{1}'.format(topic_type, command))

    def _execute(self, command, payload={}, timeout=RESP_TIMEOUT_SECONDS):
        while True:
            try:
                self.queue.get_nowait()
            except queue.Empty:
                break
        assert self.queue.empty()

        if 'token' not in payload:
            payload['token'] = self.config['opentestbed']['token']

        self.mqttc.subscribe(self._resp_topic(command))
        logger.debug('Publishing {0} to {1}'.format(payload,
                                                    self._cmd_topic(command)))
        self.mqttc.publish(self._cmd_topic(command), json.dumps(payload))

        time_left = timeout
        end_time = time.time() + timeout

        logger.debug('Waiting on {0} '.format(self._resp_topic(command))
                     + 'for {0} seconds'.format(timeout))

        ret = [] if self.device_id == 'all' else None
        while time_left > 0:
            try:
                resp = self.queue.get(timeout=time_left)
            except queue.Empty:
                break

            if resp:
                # check the topic
                _resp_topic = resp['topic'].split('/')
                device_id_index = _resp_topic.index('deviceId') + 1
                device_id = _resp_topic[device_id_index]
                if self.device_id == 'all':
                    _resp_topic[device_id_index] = '+'
                assert '/'.join(_resp_topic) == self._resp_topic(command)
                payload = resp['payload']

                if payload['success']:
                    if (payload['returnVal']
                        and (payload['returnVal'] is not True)):
                        return_value = payload['returnVal']
                    else:
                        return_value = 'Success'
                    logger.debug('Receive from {0}, '.format(device_id)
                                 + '{0}'.format(return_value))
                else:
                    logger.error('Command {0} to '.format(command)
                                 + '{0} failed'.format(device_id))
                    return_value = None
            else:
                # timeout without ret
                logger.error('Reception timeout on '
                             + '{0}'.format(self._resp_topic(command)))
                return_value = None

            if self.device_id == 'all':
                time_left = end_time - time.time()
                if return_value:
                    ret.append((device_id, return_value))
                else:
                    # do nothing
                    pass
            else:
                # done
                ret = return_value
                break

        self.mqttc.unsubscribe(self._resp_topic(command))
        return ret
# these attributes must be defined by child classes # shorthand # for offline testing, skip establishing a MQTT connection # disconnect from the broker and stop the loop # if the parent class has __del__, call it ('object' doesn't, # but check it just in case for future enhancements) # check the topic # timeout without ret # do nothing # done