From a4fec85717b9385f585e009462b98523ae963d66 Mon Sep 17 00:00:00 2001 From: Brian Martin Date: Sat, 2 Mar 2019 16:33:56 -0500 Subject: [PATCH] WIP: Bot and Client - Using tinydb - Bot and clients tracked in db --- Pipfile | 2 +- bumper/__init__.py | 104 ++++++++++++++++++++++++++++++++++--------- bumper/confserver.py | 66 +++++++++++---------------- bumper/mqttserver.py | 48 +++++++++----------- bumper/xmppserver.py | 49 ++++++++------------ 5 files changed, 147 insertions(+), 122 deletions(-) diff --git a/Pipfile b/Pipfile index 012b6da..89d3377 100644 --- a/Pipfile +++ b/Pipfile @@ -6,6 +6,6 @@ name = "pypi" [packages] hbmqtt = "*" aiohttp = "*" -black = "*" +tinydb = "*" [dev-packages] diff --git a/bumper/__init__.py b/bumper/__init__.py index 3ae2598..9d07871 100644 --- a/bumper/__init__.py +++ b/bumper/__init__.py @@ -7,8 +7,11 @@ from .xmppserver import XMPPServer import asyncio import contextvars import time +import platform +import os import logging from base64 import b64decode, b64encode +from tinydb import TinyDB, Query bumper_users_var = contextvars.ContextVar("bumper_users", default=[]) bumper_clients_var = contextvars.ContextVar("bumper_clients", default=[]) @@ -39,6 +42,22 @@ xmppserverlog = logging.getLogger("xmppserver") def get_milli_time(timetoconvert): return int(round(timetoconvert * 1000)) +def db_file(): + if platform.system() == 'Windows': + return os.path.join(os.getenv('APPDATA'), 'bumper.db') + else: + return os.path.expanduser('~/.config/bumper.db') + +def db_get(): + #Will create the database if it doesn't exist + db = TinyDB(db_file()) + + #Will create the tables if they don't exist + users_table = db.table('users') + clients_table = db.table('clients') + bots_table = db.table('bots') + + return db class BumperUser(object): def __init__(self, userid=""): @@ -108,6 +127,8 @@ class VacBotDevice(object): "name": self.name, "nick": self.nick, "resource": self.resource, + "mqtt_connection": self.mqtt_connection, + "xmpp_connection": self.xmpp_connection } @@ -120,7 +141,13 @@ class VacBotClient(object): self.xmpp_connection = False def asdict(self): - return {"userid": self.userid, "realm": self.realm, "resource": self.resource} + return { + "userid": self.userid, + "realm": self.realm, + "resource": self.resource, + "mqtt_connection": self.mqtt_connection, + "xmpp_connection": self.xmpp_connection + } def check_authcode(uid, authcode): @@ -132,46 +159,79 @@ def check_authcode(uid, authcode): return False -def add_bot(sn, did, devclass, resource, company): - +def bot_add(sn, did, devclass, resource, company): newbot = VacBotDevice() newbot.did = did newbot.name = sn newbot.vac_bot_device_class = devclass newbot.resource = resource newbot.company = company + + bot = bot_get(did) + if not bot: + bumperlog.info("Adding new bot with SN: {} DID: {}".format(newbot.name, newbot.did)) + bot_full_upsert(newbot.asdict()) - bots = bumper_bots_var.get() - existingbot = False - for bot in bots: - if bot.did == newbot.did: - existingbot = True +def bot_remove(did): + bots = db_get().table('bots') + bot = bot_get(did) + bots.remove(doc_ids=[bot.doc_id]) + +def bot_get(did): + bots = db_get().table('bots') + Bot = Query() + return bots.get(Bot.did == did) - if existingbot == False: - bots.append(newbot) - bumperlog.info("new bot added SN: {} DID: {}".format(newbot.name, newbot.did)) - bumper_bots_var.set(bots) +def bot_full_upsert(vacbot): + bots = db_get().table('bots') + Bot = Query() + bots.upsert(vacbot, Bot.did == vacbot['did']) +def bot_set_nick(did, nick): + bots = db_get().table('bots') + Bot = Query() + bots.upsert({'nick': nick}, Bot.did == did) -def add_client(userid, realm, resource): +def bot_set_mqtt(did, mqtt): + bots = db_get().table('bots') + Bot = Query() + bots.upsert({'mqtt_connection': mqtt}, Bot.did == did) +def bot_set_xmpp(did, xmpp): + bots = db_get().table('bots') + Bot = Query() + bots.upsert({'xmpp_connection': xmpp}, Bot.did == did) + +def client_add(userid, realm, resource): newclient = VacBotClient() newclient.userid = userid newclient.realm = realm newclient.resource = resource - clients = bumper_clients_var.get() + client = client_get(resource) + if not client: + bumperlog.info("Adding new client with resource {}".format(newclient.resource)) + client_full_upsert(newclient.asdict()) - existingclient = False - for client in clients: - if client.userid == newclient.userid: - existingclient = True +def client_get(resource): + clients = db_get().table('clients') + Client = Query() + return clients.get(Client.resource == resource) - if existingclient == False: - clients.append(newclient) - bumperlog.info("new client added {}".format(newclient.userid)) - bumper_clients_var.set(clients) +def client_full_upsert(client): + clients = db_get().table('clients') + Client = Query() + clients.upsert(client, Client.resource == client['resource']) +def client_set_mqtt(resource, mqtt): + clients = db_get().table('clients') + Client = Query() + clients.upsert({'mqtt_connection': mqtt}, Client.resource == resource) + +def client_set_xmpp(resource, xmpp): + clients = db_get().table('clients') + Client = Query() + clients.upsert({'xmpp_connection': xmpp}, Client.resource == resource) RETURN_API_SUCCESS = "0000" ERR_ACTIVATE_TOKEN_TIMEOUT = "1006" diff --git a/bumper/confserver.py b/bumper/confserver.py index 6cb58e4..042f101 100644 --- a/bumper/confserver.py +++ b/bumper/confserver.py @@ -567,36 +567,20 @@ class ConfServer: } elif todo == "GetDeviceList": - active_bots = self.bumper_bots.get() - bot_list = [] - for bot in active_bots: - bot_list.append(bot.asdict()) - body = {"devices": bot_list, "result": "ok", "todo": "result"} + body = {"devices": bumper.db_get().table('bots').all(), "result": "ok", "todo": "result"} elif todo == "SetDeviceNick": - bots = self.bumper_bots.get() - for bot in bots: - if postbody["did"] == bot.did: - bot.nick = postbody["nick"] - self.bumper_bots.set(bots) - body = {"result": "ok", "todo": "result"} + bumper.bot_set_nick(postbody["did"], postbody["nick"]) + body = {"result": "ok", "todo": "result"} elif todo == "AddOneDevice": - bots = self.bumper_bots.get() - for bot in bots: - if postbody["did"] == bot.did: - bot.nick = postbody["nick"] - self.bumper_bots.set(bots) - body = {"result": "ok", "todo": "result"} + bumper.bot_set_nick(postbody["did"], postbody["nick"]) + body = {"result": "ok", "todo": "result"} elif todo == "DeleteOneDevice": - bots = self.bumper_bots.get() - for bot in bots: - if postbody["did"] == bot.did: - #bots.remove(bot) - #self.bumper_bots.set(bots) - body = {"result": "ok", "todo": "result"} + bumper.bot_remove(postbody["did"]) + body = {"result": "ok", "todo": "result"} confserverlog.debug( "\r\n POST: {} \r\n Response: {}".format(postbody, body) @@ -607,6 +591,7 @@ class ConfServer: except Exception as e: confserverlog.exception("{}".format(e)) + async def handle_lookup(self, request): try: @@ -651,26 +636,25 @@ class ConfServer: try: json_body = json.loads(await request.text()) randomid = "".join(random.sample(string.ascii_letters, 6)) - bots = self.bumper_bots.get() + if "toId" in json_body: #Its a command - for bot in bots: - if bot.company == 'eco-ng': - if bot.did == json_body["toId"] and bot.mqtt_connection == True: - retcmd = await self.helperbot.send_command(json_body, randomid) - body = retcmd - confserverlog.debug( - "\r\n POST: {} \r\n Response: {}".format(json_body, body) - ) - return web.json_response(body) - - #No response, send error back - confserverlog.error( - "No bots with DID: {} connected to MQTT".format( - json_body["toId"] + bot = bumper.bot_get(json_body["toId"]) + if bot['company'] == 'eco-ng' and bot['mqtt_connection'] == True: + retcmd = await self.helperbot.send_command(json_body, randomid) + body = retcmd + confserverlog.debug( + "\r\n POST: {} \r\n Response: {}".format(json_body, body) ) - ) - body = {"id": randomid, "errno": bumper.ERR_COMMON, "ret": "fail"} - return web.json_response(body) + return web.json_response(body) + else: + #No response, send error back + confserverlog.error( + "No bots with DID: {} connected to MQTT".format( + json_body["toId"] + ) + ) + body = {"id": randomid, "errno": bumper.ERR_COMMON, "ret": "fail"} + return web.json_response(body) else: if "td" in json_body: #Seen when doing initial wifi config if json_body["td"] == "PollSCResult": diff --git a/bumper/mqttserver.py b/bumper/mqttserver.py index 6dc3f10..2c4f344 100644 --- a/bumper/mqttserver.py +++ b/bumper/mqttserver.py @@ -332,7 +332,7 @@ class BumperMQTTServer_Plugin: or str(didsplit[0]).startswith("helper") ): tmpbotdetail = str(didsplit[1]).split("/") - bumper.add_bot( + bumper.bot_add( username, didsplit[0], tmpbotdetail[0], tmpbotdetail[1], "eco-ng" ) mqttserverlog.debug( @@ -358,7 +358,7 @@ class BumperMQTTServer_Plugin: auth = True if auth: - bumper.add_client(userid, realm, resource) + bumper.client_add(userid, realm, resource) mqttserverlog.debug( "client authenticated {}".format(userid) ) @@ -375,44 +375,36 @@ class BumperMQTTServer_Plugin: async def on_broker_client_connected(self, client_id): try: - bumper_users = self.bumper_config["bumper_users"].get() - bumper_bots = self.bumper_config["bumper_bots"].get() - bumper_clients = self.bumper_config["bumper_clients"].get() didsplit = str(client_id).split("@") - for bot in bumper_bots: - if didsplit[0] == bot.did: - bot.mqtt_connection = True - mqttserverlog.debug("bot connected {}".format(bot.did)) - self.bumper_config["bumper_bots"].set(bumper_bots) + bot = bumper.bot_get(didsplit[0]) + if bot: + bumper.bot_set_mqtt(bot['did'], True) + return - for client in bumper_clients: - if didsplit[0] == client.userid and client.userid != "helper1": - client.mqtt_connection = True - # mqttserverlog.info("client connected {}".format(client.userid)) - self.bumper_config["bumper_clients"].set(bumper_clients) + clientuserid = didsplit[0] + clientresource = didsplit[1].split("/")[1] + client = bumper.client_get(clientresource) + if client: + bumper.client_set_mqtt(client['resource'], True) + return except Exception as e: mqttserverlog.exception("{}".format(e)) async def on_broker_client_disconnected(self, client_id): try: - bumper_users = self.bumper_config["bumper_users"].get() - bumper_bots = self.bumper_config["bumper_bots"].get() - bumper_clients = self.bumper_config["bumper_clients"].get() didsplit = str(client_id).split("@") - for bot in bumper_bots: - if didsplit[0] == bot.did: - bot.mqtt_connection = False - mqttserverlog.debug("bot disconnected {}".format(bot.did)) - self.bumper_config["bumper_bots"].set(bumper_bots) + bot = bumper.bot_get(didsplit[0]) + if bot: + bumper.bot_set_mqtt(bot['did'], False) - for client in bumper_clients: - if didsplit[0] == client.userid and client.userid != "helper1": - client.mqtt_connection = False - # mqttserverlog.info("client disconnected {}".format(client.userid)) - self.bumper_config["bumper_clients"].set(bumper_clients) + clientuserid = didsplit[0] + clientresource = didsplit[1].split("/")[1] + client = bumper.client_get(clientresource) + if client: + bumper.client_set_mqtt(client['userid'], False) except Exception as e: mqttserverlog.exception("{}".format(e)) diff --git a/bumper/xmppserver.py b/bumper/xmppserver.py index 78cf78b..3261093 100644 --- a/bumper/xmppserver.py +++ b/bumper/xmppserver.py @@ -228,22 +228,15 @@ class Client(threading.Thread): def _disconnect(self): try: - bumper_bots = self.bumper_bots.get() - bumper_clients = self.bumper_clients.get() - for bot in bumper_bots: - if self.uid == bot.did: - bot.xmpp_connection = False - # xmppserverlog.info("bot disconnected {}".format(bot.did)) + + bot = bumper.bot_get(self.uid) + if bot: + bumper.bot_set_xmpp(bot['did'], False) - self.bumper_bots.set(bumper_bots) - - for client in bumper_clients: - if self.uid == client.userid and client.userid != "helper1": - client.xmpp_connection = False - # xmppserverlog.info("client disconnected {}".format(client.userid)) - - self.bumper_clients.set(bumper_clients) - # xmppserverlog.debug('client {} with resource {} disconnecting'.format(self.address, self.clientresource)) + client = bumper.client_get(self.clientresource) + if client: + bumper.client_set_xmpp(client['resource'], False) + self.connection.close() except Exception as e: @@ -545,7 +538,7 @@ class Client(threading.Thread): if not self.uid.startswith("fuid"): # Need sample data to see details here - bumper.add_bot("", self.uid, "", resource) + bumper.bot_add("", self.uid, "", resource) xmppserverlog.info("bot authenticated {}".format(self.uid)) # Client authenticated, move to next state @@ -562,7 +555,7 @@ class Client(threading.Thread): auth = True if auth: - bumper.add_client(self.uid, "bumper", self.clientresource) + bumper.client_add(self.uid, "bumper", self.clientresource) xmppserverlog.debug("client authenticated {}".format(self.uid)) # Client authenticated, move to next state @@ -617,7 +610,7 @@ class Client(threading.Thread): if not self.uid.startswith("fuid"): # Need sample data to see details here - bumper.add_bot(self.uid, self.uid, self.devclass, "atom","eco-legacy") + bumper.bot_add(self.uid, self.uid, self.devclass, "atom","eco-legacy") self.type = self.BOT xmppserverlog.info("bot authenticated {}".format(self.uid)) # Send response @@ -637,7 +630,7 @@ class Client(threading.Thread): if auth: self.type = self.CONTROLLER - bumper.add_client(self.uid, "bumper", self.clientresource) + bumper.client_add(self.uid, "bumper", self.clientresource) xmppserverlog.debug("client authenticated {}".format(self.uid)) # Client authenticated, move to next state @@ -661,18 +654,14 @@ class Client(threading.Thread): try: bumper_bots = self.bumper_bots.get() bumper_clients = self.bumper_clients.get() + + bot = bumper.bot_get(self.uid) + if bot: + bumper.bot_set_xmpp(bot['did'], True) - for bot in bumper_bots: - if self.uid == bot.did: - bot.xmpp_connection = True - # xmppserverlog.info("bot connected {}".format(bot.did)) - self.bumper_bots.set(bumper_bots) - - for client in bumper_clients: - if self.uid == client.userid: - client.xmpp_connection = True - # xmppserverlog.info("client connected {}".format(client.userid)) - self.bumper_clients.set(bumper_clients) + client = bumper.client_get(self.clientresource) + if client: + bumper.client_set_xmpp(client['resource'], True) clientbindxml = xml.getchildren() clientresourcexml = clientbindxml[0].getchildren()