WIP: Bot and Client

- Using tinydb
- Bot and clients tracked in db

Pipfile

@@ -6,6 +6,6 @@ name = "pypi"
 [packages]
 hbmqtt = "*"
 aiohttp = "*"
-black = "*"
+tinydb = "*"
 
 [dev-packages]

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"

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":

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))

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()