rocketchat.py

import time
import datetime
import requests

from DDPClient import DDPClient

from pyee import EventEmitter
import hashlib

class MeteorClientException(Exception):
    """Custom Exception"""
    pass

class CollectionData(object):
    def __init__(self):
        self.data = {}

    def add_data(self, collection, id, fields):
        if collection not in self.data:
            self.data[collection] = {}
        if not id in self.data[collection]:
            self.data[collection][id] = {}
        for key, value in fields.items():
            self.data[collection][id][key] = value

    def change_data(self, collection, id, fields):
        if collection not in self.data:
            self.data[collection] = {}
        if not id in self.data[collection]:
            self.data[collection][id] = {}
        for key, value in fields.items():
            self.data[collection][id][key] = value

    def remove_data(self, collection, id):
        del self.data[collection][id]

class RocketChatClient(EventEmitter):
    def __init__(self, url, auto_reconnect=True, auto_reconnect_timeout=0.5, debug=False):
        EventEmitter.__init__(self)
        self.collection_data = CollectionData()
        self.ddp_client = DDPClient(url, auto_reconnect=auto_reconnect,
                                    auto_reconnect_timeout=auto_reconnect_timeout, debug=debug)
        self.ddp_client.on('connected', self.connected)
        self.ddp_client.on('socket_closed', self.closed)
        self.ddp_client.on('failed', self.failed)
        self.ddp_client.on('added', self.added)
        self.ddp_client.on('changed', self.changed)
        self.ddp_client.on('removed', self.removed)
        self.ddp_client.on('reconnected', self._reconnected)
        self.connected = False
        self.subscriptions = {}
        self._login_data = None
        self._login_token = None

    def connect(self):
        """Connect to the meteor server"""
        self.ddp_client.connect()

    def close(self):
        """Close connection with meteor server"""
        self.ddp_client.close()

    def _reconnected(self):
        """Reconnect
        Currently we get a new session every time so we have
        to clear all the data an resubscribe"""

        if self._login_data or self._login_token:

            def reconnect_login_callback(error, result):
                if error:
                    if self._login_token:
                        self._login_token = None
                        self._login(self._login_data,
                                    callback=reconnect_login_callback)
                        return
                    else:
                        raise MeteorClientException(
                            'Failed to re-authenticate during reconnect')

                self.connected = True
                self._resubscribe()

            if self._login_token:
                self._resume(self._login_token,
                             callback=reconnect_login_callback)
            else:
                self._login(self._login_data,
                            callback=reconnect_login_callback)
        else:
            self._resubscribe()

    def _resubscribe(self):
        self.collection_data.data = {}
        cur_subs = self.subscriptions.items()
        self.subscriptions = {}
        for name, value in cur_subs:
            self.subscribe(name, value['params'])
        self.emit('reconnected')
    #
    # Account Management
    #

    def login(self, user, password, token=None, callback=None):
        """Login with a username and password
        Arguments:
        user - username or email address
        password - the password for the account
        Keyword Arguments:
        token - meteor resume token
        callback - callback function containing error as first argument and login data"""
        # TODO: keep the tokenExpires around so we know the next time
        #       we need to authenticate

        # hash the password
        hashed = hashlib.sha256(password).hexdigest()
        # handle username or email address
        if '@' in user:
            user_object = {
                'email': user
            }
        else:
            user_object = {
                'username': user
            }
        password_object = {
            'algorithm': 'sha-256',
            'digest': hashed
        }

        self._login_token = token
        self._login_data = {'user': user_object, 'password': password_object}

        if token:
            self._resume(token, callback=callback)
        else:
            self._login(self._login_data, callback=callback)

    def _resume(self, token, callback=None):
        login_data = {'resume': token}
        self._login(login_data, callback=callback)

    def _login(self, login_data, callback=None):
        self.emit('logging_in')

        def logged_in(error, data):
            if error:
                if self._login_token:
                    self._login_token = None
                    self._login(self._login_data, callback=callback)
                    return
                if callback:
                    callback(error, None)
                return

            self._login_token = data['token']

            if callback:
                callback(None, data)
            self.emit('logged_in', data)

        self.ddp_client.call('login', [login_data], callback=logged_in)

    def logout(self, callback=None):
        """Logout a user
        Keyword Arguments:
        callback - callback function called when the user has been logged out"""
        self.ddp_client.call('logout', [], callback=callback)
        self.emit('logged_out')

    #
    # Meteor Method Call
    #

    def call(self, method, params, callback=None):
        """Call a remote method
        Arguments:
        method - remote method name
        params - remote method parameters
        Keyword Arguments:
        callback - callback function containing return data"""
        self._wait_for_connect()
        self.ddp_client.call(method, params, callback=callback)

    #
    # Subscription Management
    #

    def subscribe(self, name, params=[], callback=None):
        """Subscribe to a collection
        Arguments:
        name - the name of the publication
        params - the subscription parameters
        Keyword Arguments:
        callback - a function callback that returns an error (if exists)"""
        self._wait_for_connect()

        def subscribed(error, sub_id):
            if error:
                self._remove_sub_by_id(sub_id)
                if callback:
                    callback(error.get('reason'))
                return
            if callback:
                callback(None)
            self.emit('subscribed', name)

        if name in self.subscriptions:
            raise MeteorClientException('Already subcribed to {}'.format(name))

        sub_id = self.ddp_client.subscribe(name, params, subscribed)
        self.subscriptions[name] = {
            'id': sub_id,
            'params': params
        }

    def unsubscribe(self, name):
        """Unsubscribe from a collection
        Arguments:
        name - the name of the publication"""
        self._wait_for_connect()
        if name not in self.subscriptions:
            raise MeteorClientException('No subscription for {}'.format(name))
        self.ddp_client.unsubscribe(self.subscriptions[name]['id'])
        del self.subscriptions[name]
        self.emit('unsubscribed', name)

    #
    # Collection Management
    #

    def find(self, collection, selector={}):
        """Find data in a collection
        Arguments:
        collection - collection to search
        Keyword Arguments:
        selector - the query (default returns all items in a collection)"""
        results = []
        for _id, doc in self.collection_data.data.get(collection, {}).items():
            doc.update({'_id': _id})
            if selector == {}:
                results.append(doc)
            for key, value in selector.items():
                if key in doc and doc[key] == value:
                    results.append(doc)
        return results

    def find_one(self, collection, selector={}):
        """Return one item from a collection
        Arguments:
        collection - collection to search
        Keyword Arguments:
        selector - the query (default returns first item found)"""
        for _id, doc in self.collection_data.data.get(collection, {}).items():
            doc.update({'_id': _id})
            if selector == {}:
                return doc
            for key, value in selector.items():
                if key in doc and doc[key] == value:
                    return doc
        return None

    def insert(self, collection, doc, callback=None):
        """Insert an item into a collection
        Arguments:
        collection - the collection to be modified
        doc - The document to insert. May not yet have an _id attribute,
        in which case Meteor will generate one for you.
        Keyword Arguments:
        callback - Optional. If present, called with an error object as the first argument and,
        if no error, the _id as the second."""
        self.call("/" + collection + "/insert", [doc], callback=callback)

    def update(self, collection, selector, modifier, callback=None):
        """Insert an item into a collection
        Arguments:
        collection - the collection to be modified
        selector - specifies which documents to modify
        modifier - Specifies how to modify the documents
        Keyword Arguments:
        callback - Optional. If present, called with an error object as the first argument and,
        if no error, the number of affected documents as the second."""
        self.call("/" + collection + "/update", [selector, modifier], callback=callback)

    def remove(self, collection, selector, callback=None):
        """Remove an item from a collection
        Arguments:
        collection - the collection to be modified
        selector - Specifies which documents to remove
        Keyword Arguments:
        callback - Optional. If present, called with an error object as its argument."""
        self.call("/" + collection + "/remove", [selector], callback=callback)

    #
    # Event Handlers
    #

    def connected(self):
        self.connected = True
        self.emit('connected')

    def closed(self, code, reason):
        self.connected = False
        self.emit('closed', code, reason)

    def failed(self, data):
        self.emit('failed', str(data))

    def added(self, collection, id, fields):
        self.collection_data.add_data(collection, id, fields)
        self.emit('added', collection, id, fields)

    def changed(self, collection, id, fields, cleared):
        self.collection_data.change_data(collection, id, fields)
        self.emit('changed', collection, id, fields)
        #self.collection_data.change_data(collection, id, fields, cleared)
        #self.emit('changed', collection, id, fields, cleared)

    def removed(self, collection, id):
        self.collection_data.remove_data(collection, id)
        self.emit('removed', collection, id)

    #
    # Helper functions
    #

    def _time_from_start(self, start):
        now = datetime.datetime.now()
        return now - start

    def _remove_sub_by_id(self, sub_id):
        for name, cur_sub_id in self.subscriptions.items():
            if cur_sub_id == sub_id:
                del self.subscriptions[name]

    def _wait_for_connect(self):
        start = datetime.datetime.now()
        while not self.connected and self._time_from_start(start).seconds < 5:
            time.sleep(0.1)

        if not self.connected:
            raise MeteorClientException('Could not subscribe because a connection has not been established')

class RocketChatBot():
    def __init__(self, user, password, server='open.rocket.chat', ssl=True):
        self.username = user
        self.password = password
        self.server = server
        self.ssl = ssl
        self.token = None
        self.userId = None
        self.debug = True

        self._prefixs = []
        if ssl:
            protocol = 'wss://'
        else:
            protocol = 'ws://'
        self.client = RocketChatClient(protocol + server + '/websocket',debug=self.debug)

        # registering internal handlers
        self.client.on('connected', self._connected)
        self.client.on('closed', self._closed)
        self.client.on('logged_in', self._logged_in)
        self.client.on('failed', self._failed)
        self.client.on('added', self._added)
        self.client.on('changed', self._changed)
        self.client.on('unsubscribed', self._unsubscribed)
        self.client.on('subscribed', self._subscribed)

    """
    Internal events handlers
    """
    def _connected(self):
        print("[+] rocketchat: connected")
        self.client.subscribe('stream-room-messages', ['__my_messages__', False], self.cb1)

    def _closed(self, code, reason):
        print('[-] rocketchat: connection closed: %s (%d)' % (reason, code))

    def _logged_in(self, data):
        print('[+] rocketchat: logged in')
        self.token = data['token']
        self.userId = data['id']
        print(data)

    def _failed(self, collection, data):
        print('[-] %s' % str(data))

    def _added(self, collection, id, fields):
        print('[+] added %s: %s' % (collection, id))

    def _changed(self, collection, mid, fields):
        print('[+] changed: %s %s' % (collection, mid))

        if not fields.get('args'):
            return

        args = fields['args']

        if args[0] == "GENERAL":
            print("[+] message: general, skipping")
            return

        if args[0].get('msg'):
            return self._incoming(args[0])

        if args[0].get('attachments'):
            return self._downloading(args[0])

    def _subscribed(self, subscription):
        print('[+] subscribed: %s' % subscription)

    def _unsubscribed(self, subscription):
        print('[+] unsubscribed: %s' % subscription)

    """
    Internal callback handlers
    """
    def cb(self, error, data):
        if not error:
            if self.debug:
                print(data)
            return

        print('[-] callback error:')
        print(error)

    def cb1(self, data):
        # if not self.debug:
        #     return
        if(len(data)>0):
            print(data)
            self._incoming(data)
        else:
            print("[+] callback success")

    """
    Internal dispatcher
    """
    def _incoming(self, data):
        print("[+] Message from %s: %s" % (data['u']['username'], data['msg']))
        print("[+] Incoming Message")
        #self.sendMessage(data['rid'],  "I hear you")
        # print(data)
        #Check if message was sent by another user

        for prefix in self._prefixs:
            if data['msg'].startswith(prefix['prefix']):
                prefix['handler'](self, data)

    def _downloading(self, data):
        print("[+] attachement from %s: %d files" % (data['u']['username'], len(data['attachments'])))

    """
    Public initializers
    """
    def start(self):
        self.client.connect()
        self.client.login(self.username, self.password.encode('utf-8'), callback=self.cb)

        # let's yeld to background task
        while True:
            time.sleep(3600)

    def addPrefixHandler(self, prefix, handler):
        self._prefixs.append({'prefix': prefix, 'handler': handler})

    def sendMessage(self, id, message):
        self.client.call('sendMessage', [{'msg': message, 'rid': id}], self.cb)

    def post(self, method, params=None):
        headers = {
            'X-User-Id': self.userId,
            'X-Auth-Token': self.token,
            'Content-type': 'application/json',
        }
        if self.ssl:
            protocol = 'https://'
        else:
            protocol = 'http://'

        res = requests.post(
            protocol + self.server + "/api/v1/" + method, 
            json = params,
            headers = headers
        )
        data = {"response": res, "data": None}
        if res and res.status_code == 200 and 'json' in res.headers.get('Content-Type'):
            data["data"] = res.json()
        return data

    def get(self, method, params=None):
        headers = {
            'X-User-Id': self.userId,
            'X-Auth-Token': self.token,
        }
        if self.ssl:
            protocol = 'https://'
        else:
            protocol = 'http://'

        res = requests.get(
            protocol + self.server + "/api/v1/" + method, 
            params,
            headers = headers
        )
        data = {"response": res, "data": None}
        if res and res.status_code == 200 and 'json' in res.headers.get('Content-Type'):
            data["data"] = res.json()
        return data