bumper.py

@@ -6,9 +6,10 @@ import sys, socket
 import time
 import platform
 
-bumperlog = logging.getLogger("bumper")
+
 def main():
     args = sys.argv
+    
     if len(args) > 0:
         if '--debug' in args:
             logging.basicConfig(level=logging.DEBUG,        
@@ -18,14 +19,7 @@ def main():
                 format="[%(asctime)s] :: %(levelname)s :: %(name)s :: %(message)s")
                 #format="[%(asctime)s] :: %(levelname)s :: %(name)s :: %(module)s :: %(funcName)s :: %(lineno)d :: %(message)s")
 
-    # A default bot could be set here to automatically add it as available
-    # dbot = bumper.VacBotDevice("did", "class", "resource", "name","nick" )
-    # bclient = bumper.bumper_bots_var
-    # bclienttemp = bclient.get()
-    # bclienttemp.append(dbot.asdict())
-    # bclient.set(bclienttemp)
-
-    if platform.system() == "Darwin":
+    if platform.system() == "Darwin": #If a Mac, use 0.0.0.0 for listening
         listen_host = "0.0.0.0"
     else:
         listen_host = socket.gethostbyname(socket.gethostname())
@@ -35,36 +29,46 @@ def main():
     conf_address_8007 = (listen_host, 8007)
     xmpp_address = (listen_host, 5223)
     mqtt_address = (listen_host, 8883)
+    
+    xmpp_server = bumper.XMPPServer(xmpp_address)
+    mqtt_server = bumper.MQTTServer(mqtt_address,bumper_bots=bumper.bumper_bots_var,bumper_clients=bumper.bumper_clients_var,remove_clients=bumper.bumper_removeclients_var)
+    mqtt_helperbot = bumper.MQTTHelperBot(mqtt_address, bumper_bots=bumper.bumper_bots_var,bumper_clients=bumper.bumper_clients_var)
+    conf_server = bumper.ConfServer(conf_address_443, usessl=True, bumper_bots=bumper.bumper_bots_var,bumper_clients=bumper.bumper_clients_var, remove_clients=bumper.bumper_removeclients_var,helperbot=mqtt_helperbot)
+    conf_server_2 = bumper.ConfServer(conf_address_8007, usessl=False, bumper_bots=bumper.bumper_bots_var,bumper_clients=bumper.bumper_clients_var, helperbot=mqtt_helperbot,remove_clients=bumper.bumper_removeclients_var)
+    
+    # start xmpp server on port 5223 (sync)
+    xmpp_server.run(run_async=True) #Start in new thread
 
-    # start mqtt server on port 8883 (async)
-    startmqttserver = "Starting MQTT Server at {}".format(mqtt_address)    
-    bumperlog.info("{}".format(startmqttserver))
-    print("{}".format(startmqttserver))
-    mqtt_server = bumper.MQTTServer(mqtt_address, run_async=True,bumper_bots=bumper.bumper_bots_var,bumper_clients=bumper.bumper_clients_var)
+    # start mqtt server on port 8883 (async)    
+    mqtt_server.run(run_async=True) #Start in new thread    
+    
     time.sleep(1.5) #Wait for broker startup
 
-    # start mqtt_helperbot (async)
-    bumperlog.info("Starting MQTT HelperBot")
-    print("Starting MQTT HelperBot")
-    mqtt_helperbot = bumper.MQTTHelperBot(mqtt_address, run_async=True,bumper_bots=bumper.bumper_bots_var,bumper_clients=bumper.bumper_clients_var)
+    # start mqtt_helperbot (async)       
+    mqtt_helperbot.run(run_async=True) #Start in new thread
 
-    # start conf server on port 443 (async) - Used for most https calls
-    startconf443 = "Starting Main ConfServer at {}".format(conf_address_443)
-    bumperlog.info("{}".format(startconf443))
-    print("{}".format(startconf443))
-    conf_server = bumper.ConfServer(conf_address_443, usessl=True, run_async=True,bumper_bots=bumper.bumper_bots_var,bumper_clients=bumper.bumper_clients_var, helperbot=mqtt_helperbot)
+    # start conf server on port 443 (async) - Used for most https calls        
+    conf_server.run(run_async=True) #Start in new thread
 
-    # start conf server on port 8007 (async) - Used for a load balancer request
-    startconf8007 = "Starting LoadBalancer ConfServer at {}".format(conf_address_8007)
-    bumperlog.info("{}".format(startconf8007))
-    print("{}".format(startconf8007))    
-    conf_server_2 = bumper.ConfServer(conf_address_8007, usessl=False, run_async=True,bumper_bots=bumper.bumper_bots_var,bumper_clients=bumper.bumper_clients_var, helperbot=mqtt_helperbot)
+    # start conf server on port 8007 (async) - Used for a load balancer request    
+    conf_server_2.run(run_async=True) #Start in new thread
 
-    # start xmpp server on port 5223 (sync)
-    startxmpp = "Starting XMPP Server at {}".format(xmpp_address)
-    bumperlog.info("{}".format(startxmpp))
-    print("{}".format(startxmpp))
-    xmpp_server = bumper.XMPPServer(xmpp_address)
+    while True:
+        try:
+            time.sleep(0.25)
 
+            # WIP: Remove clients that have disconnected
+            remove_clients = bumper.bumper_removeclients_var.get()
+            if len(remove_clients) > 0:
+                for uid in remove_clients:
+                    if uid != "":
+                        xmpp_server.remove_client_byuid(uid)  #Remove clients from xmpp server
+                        remove_clients.remove(uid)
+                
+                bumper.bumper_removeclients_var.set(remove_clients)
+        
+        except KeyboardInterrupt:                
+            print("Bumper Exiting")
+    
 if __name__ == "__main__":
     main()

bumper/__init__.py

@@ -7,13 +7,30 @@ from .xmppserver import XMPPServer
 import asyncio
 import contextvars
 import time
+import logging
 
 bumper_clients_var = contextvars.ContextVar('bumper_clients', default=[])
 bumper_bots_var = contextvars.ContextVar('bumper_bots', default=[])
+bumper_removeclients_var = contextvars.ContextVar('bumper_removeclients', default=[])
 ca_cert = './certs/CA/cacert.pem'
 server_cert = './certs/cert.pem'
 server_key = './certs/key.pem'
 
+#Logs
+bumperlog = logging.getLogger("bumper")
+confserverlog = logging.getLogger("confserver")
+#Override the logging level
+#confserverlog.setLevel(logging.INFO) 
+mqttserverlog = logging.getLogger("mqttserver")
+#Override the logging level
+#mqttserverlog.setLevel(logging.INFO)
+helperbotlog = logging.getLogger("helperbot")
+#Override the logging level
+#helperbotlog.setLevel(logging.INFO)
+xmppserverlog = logging.getLogger("xmppserver")
+#Override the logging level
+#xmppserverlog.setLevel(logging.INFO)
+
 def get_milli_time(timetoconvert):
     return int(round(timetoconvert * 1000))
 

bumper/confserver.py

@@ -32,33 +32,37 @@ class ConfServer():
     bumper_clients = contextvars.ContextVar
     bumper_bots = contextvars.ContextVar
 
-    def __init__(self, address, usessl=False, run_async=True, bumper_bots=contextvars.ContextVar, bumper_clients=contextvars.ContextVar, helperbot=None):
+    def __init__(self, address, usessl=False, bumper_bots=contextvars.ContextVar, bumper_clients=contextvars.ContextVar, remove_clients=contextvars.ContextVar,helperbot=None):
         self.bumper_bots = bumper_bots
         self.bumper_clients = bumper_clients
+        self.remove_clients = remove_clients
         self.helperbot = helperbot
         self.usessl = usessl
-        self.run_async = run_async
         self.address = address  
-        
+        self.confthread = None
 
+        
+    def run(self, run_async=False):
         try:                        
             if run_async:
                 confserverlog.debug("Starting ConfServer Thread: 1")
-                confserver = Thread(name="ConfServer_Thread",target=self.run_server)
-                self.server = confserver
-                confserver.setDaemon(True)
-                confserver.start()
+                self.confthread = Thread(name="ConfServer_{}_Thread".format(self.address[1]),target=self.run_server)
+                self.confthread.setDaemon(True)
+                self.confthread.start()
                 
             else:
                 try:
                     self.run_server()
                 except KeyboardInterrupt:
                     self.disconnect()
+
         except Exception as e:
             confserverlog.exception('{}'.format(e))
 
 
     def run_server(self):
+        logging.info("Starting ConfServer at {}".format(self.address))
+        print("Starting ConfServer at {}".format(self.address))
         try:
             loop = asyncio.get_event_loop()
         except:
@@ -76,6 +80,7 @@ class ConfServer():
             app = web.Application()        
             
             app.add_routes([
+                web.get('', self.handle_base),
                 web.get('/{apiversion}/private/{country}/{language}/{devid}/{apptype}/{appversion}/{devtype}/{aid}/user/login', self.handle_login),
             #    web.get('/{apiversion}/private/{country}/{language}/{devid}/{apptype}/{appversion}/{devtype}/{aid}/user/checkLogin', self.handle_checkLogin),            
                 web.get('/{apiversion}/private/{country}/{language}/{devid}/{apptype}/{appversion}/{devtype}/{aid}/user/logout', self.handle_logout),
@@ -115,6 +120,16 @@ class ConfServer():
             confserverlog.exception('{}'.format(e))            
             exit(1)
 
+    async def handle_base(self, request):              
+        try:
+            
+            text = "Bumper!"
+                                                
+            return web.json_response(text)
+        
+        except Exception as e:
+            confserverlog.exception('{}'.format(e))                          
+
     async def handle_login(self, request):              
         try:
             #Could implement basic auth if you wanted, or just accept anything    
@@ -162,8 +177,14 @@ class ConfServer():
 
     async def handle_logout(self, request):                      
         try:
+            uid = request.query['uid']
             body = {"code": "0000","data": None,"msg": "操作成功", "time": bumper.get_milli_time(time.time())}
-            #TODO - when logging out close out any other connections MQTT/XMPP
+            # QUERY String
+            # 'uid=fuid_CUOVIn&accessToken=tempaccesstoken&requestId=e584de79f9cca854df6fb3352c6893b6&authTimespan=1550168440575&authTimeZone=GMT-5&authAppkey=eJUWrzRv34qFSaYk&authSign=14e38f95c7316e111c5815cf15f0f972'     
+            if not uid == "":
+                remove_clients = self.remove_clients.get()
+                remove_clients.append(request.query['uid'])
+                self.remove_clients.set(remove_clients)                
             
             return web.json_response(body)        
 
@@ -341,9 +362,9 @@ class ConfServer():
         try:
             confserverlog.info('shutting down')        
             if(self.run_async):
-                self.server.join()
+                self.confthread.join()
             else:
-                self.server.disconnect()  
+                self.confthread.disconnect()  
     
         except Exception as e:
             confserverlog.exception('{}'.format(e))      

bumper/mqttserver.py

@@ -28,28 +28,28 @@ logging.getLogger("hbmqtt.client").setLevel(logging.CRITICAL + 1) #Ignore this l
 class MQTTHelperBot():
     
     Client = MQTTClient()
-    def __init__(self, address, run_async=False, bumper_bots=contextvars.ContextVar, bumper_clients=contextvars.ContextVar):
-        
+    def __init__(self, address,  bumper_bots=contextvars.ContextVar, bumper_clients=contextvars.ContextVar):        
         self.address = address
         self.client_id = "helper1@bumper/helper1"       
         self.command_responses = contextvars.ContextVar('command_responses', default=[])
-        
-        try:            
-            if run_async: 
+        self.helperthread = None        
+    
+
+    def run(self, run_async=False):
+        if run_async: 
                 hloop = asyncio.new_event_loop()                
                 helperbotlog.debug("Starting MQTT HelperBot Thread: 1")
-                helperbot = Thread(name="MQTTHelperBot_Thread",target=self.run_helperbot, args=(hloop,))
-                helperbot.setDaemon(True)
-                helperbot.start()                    
+                self.helperthread = Thread(name="MQTTHelperBot_Thread",target=self.run_helperbot, args=(hloop,))
+                self.helperthread.setDaemon(True)
+                self.helperthread.start()                    
                 
-            else:
-                self.run_helperbot()
+        else:
+            self.run_helperbot()
 
-        except Exception as e:
-            helperbotlog.exception('{}'.format(e))
-            pass    
 
-    def run_helperbot(self, loop):     
+    def run_helperbot(self, loop):    
+        logging.info("Starting MQTT HelperBot")
+        print("Starting MQTT HelperBot") 
         try:
             asyncio.set_event_loop(loop)                
             self.Client = MQTTClient(client_id=self.client_id, config={'check_hostname':False})      
@@ -77,7 +77,10 @@ class MQTTHelperBot():
                 message = await self.Client.deliver_message()          
             
                 #helperbotlog.debug("HelperBot MQTT Received Message on Topic: {} - Message: {}".format(message.topic, str(message.payload.decode("utf-8"))))
-                cresp = self.command_responses.get()        
+                cresp = self.command_responses.get()    
+                
+                if (str(message.topic).split("/")[6] == "helper1"):
+                    cresp.append({"time": time.time() ,"topic": message.topic,"payload":str(message.data.decode("utf-8"))})                          
                 
                 #Cleanup "expired messages" > 60 seconds from time
                 for msg in cresp:           
@@ -85,8 +88,7 @@ class MQTTHelperBot():
                     if time.time() > expire_time:
                         #helperbotlog.debug("Pruning Message Time: {}, MsgTime: {}, MsgTime+60: {}".format(time.time(), msg['time'], expire_time))
                         cresp.remove(msg)                                  
-
-                cresp.append({"time": time.time() ,"topic": message.topic,"payload":str(message.data.decode("utf-8"))})
+                                                           
                 self.command_responses.set(cresp)                          
                 #helperbotlog.debug("MQTT Command Response List Count: %s" %len(cresp))
         
@@ -95,7 +97,9 @@ class MQTTHelperBot():
 
     async def wait_for_resp(self, requestid):               
         try:
+            
             t_end = (datetime.now() + timedelta(seconds=10)).timestamp()
+            
             while time.time() < t_end:
                 await asyncio.sleep(0.1)
                 responses = self.command_responses.get()   
@@ -118,6 +122,7 @@ class MQTTHelperBot():
                             self.command_responses.set(cresp)                             
                             return resp
 
+            
             return { "id": requestid, "errno": "timeout", "ret": "fail" }
         except asyncio.CancelledError as e:
             helperbotlog.debug('wait_for_resp cancelled by asyncio')
@@ -169,15 +174,21 @@ class MQTTServer():
         except Exception as e:
             mqttserverlog.exception('{}'.format(e))                
 
-    def __init__(self, address, run_async=False, bumper_bots=contextvars.ContextVar, bumper_clients=contextvars.ContextVar):
+    def __init__(self, address, bumper_bots=contextvars.ContextVar, bumper_clients=contextvars.ContextVar,remove_clients=contextvars.ContextVar):
         try:
+           
+            self.bumper_bots = bumper_bots
+            self.bumper_clients = bumper_clients      
+            self.remove_clients = remove_clients      
+            self.mqttserverthread = None  
+            self.address = address
+
             #The below adds a plugin to the hbmqtt.broker.plugins without having to futz with setup.py
             distribution = pkg_resources.Distribution("hbmqtt.broker.plugins")
-            bumper_plugin = pkg_resources.EntryPoint.parse('bumper = bumper.mqttserver:BumperMQTTServer_Plugin', dist=distribution)        
+            bumper_plugin = pkg_resources.EntryPoint.parse('bumper = bumper.mqttserver:BumperMQTTServer_Plugin', dist=distribution)                            
             distribution._ep_map = {"hbmqtt.broker.plugins": {"bumper": bumper_plugin}}
             pkg_resources.working_set.add(distribution)
-            self.bumper_bots = bumper_bots
-            self.bumper_clients = bumper_clients        
+
             # Initialize bot server
             self.default_config = {
                 'listeners': {
@@ -203,26 +214,32 @@ class MQTTServer():
                     'enabled': False
                 },
                 'clients':{
-                    'connected_bots': bumper_bots,
-                    'connected_clients': bumper_clients
-
-                }
-            }
-            if run_async:
-                sloop = asyncio.new_event_loop()                
-                mqttserverlog.debug("Starting MQTTServer Thread: 1")
-                mqttserver = Thread(name="MQTTServer_Thread",target=self.run_server, args=(sloop,))
-                mqttserver.setDaemon(True)
-                mqttserver.start()                   
-                
-            else:
-                self.run_server()
+                    'connected_bots': self.bumper_bots,
+                    'connected_clients': self.bumper_clients,
+                    'remove_clients': self.remove_clients
+                },                               
+            }            
         
         except Exception as e:
             mqttserverlog.exception('{}'.format(e))    
-    
 
-    def run_server(self, loop):      
+    def run(self, run_async=False,):
+
+        if run_async:
+                sloop = asyncio.new_event_loop()                
+                mqttserverlog.debug("Starting MQTTServer Thread: 1")
+                self.mqttserverthread = Thread(name="MQTTServer_Thread",target=self.run_server, args=(sloop,))
+                self.mqttserverthread.setDaemon(True)
+                self.mqttserverthread.start()                   
+                
+        else:
+            self.run_server()            
+
+
+    def run_server(self, loop):     
+        
+        logging.info("Starting MQTT Server at {}".format(self.address))
+        print("Starting MQTT Server at {}".format(self.address))
         try:     
             asyncio.set_event_loop(loop)
             loop.run_until_complete(self.broker_coro())
@@ -237,6 +254,7 @@ class BumperMQTTServer_Plugin:
         self.context = context        
         try:            
             self.clients = self.context.config['clients']
+
         except KeyError:
              self.context.logger.warning("'clients' section not found in context configuration")
         except Exception as e:
@@ -277,7 +295,7 @@ class BumperMQTTServer_Plugin:
                     if client['userid'] == newuser.userid:
                         clientactive = True
                 
-                if clientactive == False:
+                if clientactive == False and newuser.userid != 'helper1':
                     connected_clients.append(newuser.asdict())
                     mqttserverlog.info("new client {}".format(newuser.userid))
                 
@@ -295,20 +313,25 @@ class BumperMQTTServer_Plugin:
         try:
             #mqttserverlog.debug('%s disconnected' % client_id)
             connected_bots = self.clients['connected_bots'].get()       
-            connected_clients = self.clients['connected_clients'].get()               
+            connected_clients = self.clients['connected_clients'].get()     
+            remove_clients = self.clients['remove_clients'].get()     
             didsplit = str(client_id).split("@")
             #If the did is in the list, remove it
             for bot in connected_bots:
                 if didsplit[0] == bot['did']:
                     mqttserverlog.info("bot disconnected {}".format(bot['did']))
                     connected_bots.remove(bot)
-                    self.clients['connected_bots'].set(connected_bots)           
+                    #remove_clients.append(bot['did'])                    
+                    self.clients['connected_bots'].set(connected_bots)                          
 
             for client in connected_clients:
-                if didsplit[0] == client['userid']:
-                    mqttserverlog.info("client disconnected {}".format(client['userid']))
+                if didsplit[0] == client['userid'] and client['userid'] != 'helper1':
+                    mqttserverlog.info("client disconnected {}".format(client['userid']))                    
                     connected_clients.remove(client)
+                    #remove_clients.append(client['userid'])                                    
                     self.clients['connected_clients'].set(connected_clients)
+                    
+            #self.clients['remove_clients'].set(remove_clients)        
 
             #mqttserverlog.debug('Connected Bots: %s' %self.clients['connected_bots'].get())
             #mqttserverlog.debug('Connected Clients: %s' %self.clients['connected_clients'].get())

bumper/xmppserver.py

@@ -1,6 +1,10 @@
 #!/usr/bin/env python3
 
+from threading import Thread
 import sys, socket, threading, re, time, logging, uuid, xml.etree.ElementTree as ET
+import base64
+import ssl
+import bumper
 
 xmppserverlog = logging.getLogger("xmppserver")
 
@@ -12,51 +16,115 @@ class XMPPServer():
     exit_flag = False
 
     def __init__(self, address):
+        # Initialize bot server
+        self.address = address
+        
+
+    def run(self, run_async=False):    
+        if run_async:
+                xmppserverlog.debug("Starting XMPPServer Thread: 1")
+                self.xmppthread = Thread(name="XMPPServer_Thread",target=self.run_server)                
+                self.xmppthread.setDaemon(True)
+                self.xmppthread.start()
+                
+        else:
+            try:
+                self.run_server()
+            except KeyboardInterrupt:
+                self.disconnect()
+
+    def run_server(self):        
+        logging.info("Starting XMPP Server at {}".format(self.address))
+        print("Starting XMPP Server at {}".format(self.address))
+        
+        #xmppserverlog.setLevel(logging.DEBUG)       
+
+        #Set SSL Context
+        self.ssl_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)                        
+        self.ssl_ctx.load_cert_chain(certfile=bumper.server_cert,keyfile=bumper.server_key)
+        
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)               
+        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)   
+
         try:
-            # Initialize bot server
-            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-            self.socket.bind(address)
-            self.socket.listen(1)
-            xmppserverlog.debug('listening on {}:{}'.format(address[0], address[1]))
-            while not self.exit_flag:
+            self.socket.bind(self.address)
+            self.socket.listen(5)              
+
+            xmppserverlog.debug('listening on {}:{}'.format(self.address[0], self.address[1]))
+            while not self.exit_flag:                                         
                 connection, client_address = self.socket.accept()
+                
                 # disconnect any clients with this ip
                 for client in self.clients:
                     if client.address == client_address[0]:                        
-                        client.disconnect()
-                        client.join()
-                thread_id = uuid.uuid4()
-                client = Client(thread_id, connection, client_address)
+                        xmppserverlog.debug('disconnecting existing client {} with resource {}'.format(client.address, client.clientresource))
+                        client._disconnect()
+                        self.remove_client_byip(client.address)                
                 
+                xmppserverlog.debug('starting new client with ip {}'.format(client_address[0]))
+                thread_id = uuid.uuid4()                
+                client = Client(thread_id, connection, client_address)       
+                client.setDaemon(True)         
                 client.start()
-                self.clients.append(client)
-            self.socket.close()
+                self.clients.append(client)       
+            
         except PermissionError as e:
             if "bind" in e.strerror:
                 xmppserverlog.exception("Error binding XMPPServer, exiting. Try using a different hostname or IP - {}".format(e))        
-            exit(1)            
+                exit(1)            
+                            
         except Exception as e:
             xmppserverlog.exception('{}'.format(e))
             exit(1)
+        
         except KeyboardInterrupt as e:
             xmppserverlog.exception('{}'.format(e))
+        
         finally:
+            connection.shutdown(socket.SHUT_RDWR)
+            connection.close()
             self.disconnect()
-            xmppserverlog.info('disconnecting')
+            xmppserverlog.info('disconnecting')                    
+
+        self.socket.close()
+        
+        
 
     def disconnect(self):
         try:
             xmppserverlog.debug('waiting for all client threads to exit')
             for client in self.clients:
-                client.disconnect()
-                client.join()
+                client._disconnect()
+            
             self.exit_flag = True
             xmppserverlog.debug('shutting down')
+        
         except Exception as e:
             xmppserverlog.exception("{}".format(e))  
 
 
+    def remove_client_byip(self, ip):
+        for client in self.clients:
+            if client.address == ip:
+                xmppserverlog.debug('removing client from client list with ip {} and resource {}'.format(client.address, client.clientresource))
+                client._disconnect()
+                self.clients.remove(client)
+
+    def remove_client_byresource(self, resource):
+        for client in self.clients:
+            if str(client.clientresource).lower() == str(resource).lower():
+                xmppserverlog.debug('removing client from client list with ip {} and resource {}'.format(client.address, client.clientresource))
+                client._disconnect()
+                self.clients.remove(client)   
+
+    def remove_client_byuid(self, uid):
+        for client in self.clients:
+            if str(client.uid).lower() == str(uid).lower():
+                xmppserverlog.debug('removing client from client list with ip {} and resource {}'.format(client.address, client.clientresource))
+                client._disconnect()
+                self.clients.remove(client)  
+
+
 class Client(threading.Thread):
     IDLE = 0
     CONNECT = 1
@@ -68,32 +136,47 @@ class Client(threading.Thread):
     BOT = 1
     CONTROLLER = 2
 
-    def __init__(self, thread_id, connection, client_address):
+    def __init__(self, thread_id, connection, client_address):        
         threading.Thread.__init__(self)
         self.id = thread_id
-        self.name = "XMPP Thread {}".format(client_address[0])
+        self.name = "XMPP_Client_{}".format(client_address[0])
         self.type = self.UNKNOWN
         self.state = self.IDLE
         self.connection = connection
         self.address = client_address[0]
         self.clientresource = ""
+        self.uid = ""        
+        self.log_sent_message = False #Set to true to log sends
+        self.log_incoming_data = True #Set to true to log sends
+
+        xmppserverlog.debug('new client thread init for client with ip {}'.format(self.address))
 
     def send(self, command):
         try:
-            xmppserverlog.debug('send {} - {}'.format(self.address, command))
-            self.connection.send(command.encode())
+            if not self.connection._closed:
+                if self.log_sent_message:
+                    xmppserverlog.debug('send {} - {}'.format(self.address, command))
+                self.connection.send(command.encode())       
+        
         except OSError as e:
-            xmppserverlog.error('{}'.format(e))            
+            xmppserverlog.error('{}'.format(e))           
+
+        except BrokenPipeError as e:
+            xmppserverlog.error('{}'.format(e))
+            #self._set_state('DISCONNECT')     
+
+        except ConnectionResetError as e:
+            xmppserverlog.error('{}'.format(e))
+            #self._set_state('DISCONNECT')             
+        
         except Exception as e:
             xmppserverlog.exception("{}".format(e))
+            
 
-
-
-    def disconnect(self):
+    def _disconnect(self):
         try:
             xmppserverlog.debug('client {} with resource {} disconnecting'.format(self.address, self.clientresource))
             self.connection.close()
-            self._set_state('DISCONNECT')
  
         except Exception as e:
             xmppserverlog.exception("{}".format(e))        
@@ -103,6 +186,7 @@ class Client(threading.Thread):
             if tag[0] == '{':
                 uri, ignore, tag = tag[1:].partition('}')
             return tag
+        
         except Exception as e:
             xmppserverlog.exception("{}".format(e))           
 
@@ -111,21 +195,40 @@ class Client(threading.Thread):
             new_state = getattr(Client, state)
             if self.state > new_state:
                 raise Exception('{} illegal state change {}->{}'.format(self.address, self.state, new_state))
+            
             xmppserverlog.debug('{} state: {}'.format(self.address, state))
+            
             self.state = new_state
-            if new_state == '5':
-                self.join()
+            
+            if new_state == 5:
+                self._disconnect()
+        
         except Exception as e:
             xmppserverlog.exception("{}".format(e))            
 
-    def _handle_ctl(self, xml, data):
+    def _handle_ctl(self, xml, data):                                            
         try:
-            ctl = xml[0][0]
-            if ctl.get('admin') and self.type == self.BOT:
-                xmppserverlog.debug('admin username received from bot: {}'.format(ctl.get('admin')))
-                XMPPServer.client_id = ctl.get('admin')
+
+            if data.decode('utf-8').find('roster') > -1:                                    
+                #Return not-implemented for roster                              
+                self.send('<iq type="error" id="{}"><error type="cancel" code="501"><feature-not-implemented xmlns="urn:ietf:params:xml:ns:xmpp-stanzas"/></error></iq>'.format(xml.get('id')))
                 return
-            # forward
+           
+            if xml.get('type') == 'set': 
+                if data.decode('utf-8').find('com:sf') > -1 and xml.get('to') == 'rl.ecorobot.net': #Android bind? Not sure what this does yet. 
+                    self.send('<iq id="{}" to="{}@{}/{}" from="rl.ecorobot.net" type="result"/>'.format(xml.get('id'), self.uid, XMPPServer.bot_id, self.clientresource))
+                
+                else:
+                    xmppserverlog.debug('Unknown set type: {}'.format(data.decode('utf-8')))
+
+            if xml[0][0]:
+                ctl = xml[0][0]
+                if ctl.get('admin') and self.type == self.BOT:
+                    xmppserverlog.debug('admin username received from bot: {}'.format(ctl.get('admin')))
+                    XMPPServer.client_id = ctl.get('admin')
+                    return                    
+
+            #forward
             for client in XMPPServer.clients:
                 if client.address != self.address and client.state == client.READY:
                     if client.type == self.BOT:
@@ -135,106 +238,298 @@ class Client(threading.Thread):
                             data = data[:id_index] + 'from="' + XMPPServer.client_id + '" ' + data[id_index:]
                             data = data.encode()
                     client.send(data.decode('utf-8'))
+        
         except Exception as e:
             xmppserverlog.exception("{}".format(e))
 
 
     def _handle_ping(self, xml, data):
-        if(xml.get('to').find('@') == -1):
-            # Ping to server - respond
-            self.send('<iq type="result" id="{}" from="{}" />'.format(xml.get('id'), xml.get('to')))
-        else:
-            for client in XMPPServer.clients:
-                if client.address != self.address and client.state == client.READY:
-                    client.send(data.decode('utf-8'))
+        try:
+            if(xml.get('to').find('@') == -1):
+                # Ping to server - respond
+                self.send('<iq type="result" id="{}" from="{}" />'.format(xml.get('id'), xml.get('to')))
 
-    def _handle_result(self, data):
+            else:
+                for client in XMPPServer.clients:
+                    if client.address != self.address and client.state == client.READY:
+                        client.send(data.decode('utf-8'))
+        
+        except Exception as e:
+            xmppserverlog.exception('{}'.format(e))                        
+
+    def _handle_result(self, data):                                        
         # forward
         try:
             for client in XMPPServer.clients:
                 if client.address != self.address and client.state == client.READY:
                     client.send(data.decode('utf-8'))
+        
         except Exception as e:
-            xmppserverlog.exception("{}".format(e))            
+            xmppserverlog.exception("{}".format(e))         
+                         
 
-    def run(self):
+    def _handle_connect(self, data):                                          
         try:
-            xmppserverlog.info('client connected - {}'.format(self.address))
-            self._set_state('CONNECT')
-            data = ""
-            while True:
-                time.sleep(0.2)
-                if not self.connection._closed:       
-                    try:        
-                        data = self.connection.recv(4096)
-                    except ConnectionResetError as e:
-                        xmppserverlog.error('{}'.format(e))
-                    except Exception as e:
-                        xmppserverlog.exception('{}'.format(e))
-                if data:
-                    xmppserverlog.debug('from {} - {}'.format(self.address, data.decode('utf-8')))
-                    try:
-                        if self.state == self.CONNECT:
-                            if data.decode('utf-8').find('jabber:client') > -1:
-                                self._set_state('INIT')
-                                # ack jabbr:client
-                                self.send('<stream:stream xmlns:stream="http://etherx.jabber.org/streams" xmlns="jabber:client" version="1.0" id="1" from="{}">'.format(XMPPServer.server_id))
-                                time.sleep(0.5)
-                                # session
-                                self.send('<stream:features><bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"/><session xmlns="urn:ietf:params:xml:ns:xmpp-session"/></stream:features>')
-                            continue
-                        xml = ET.fromstring(data)
-                        if len(xml):
-                            child = self._tag_strip_uri(xml[0].tag)
-                        else:
-                            child = None
-                        if xml.tag == 'iq':
-                            res = None
-                            if child == 'bind':
-                                clientbindxml = xml.getchildren()
-                                clientresourcexml = clientbindxml[0].getchildren()
-                                self.clientresource = clientresourcexml[0].text
-                                self.name = "XMPP Thread {}".format(self.clientresource)
-                                xmppserverlog.debug("new client {} using resource {}".format(self.address, self.clientresource))
-                                res = '<iq type="result" id="{}"><bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"><jid>{}</jid></bind></iq>'.format(xml.get('id'), XMPPServer.bot_id)
-                                self._set_state('BIND')
-                            elif child == 'session':
-                                res = '<iq type="result" id="{}" />'.format(xml.get('id'))
-                                self._set_state('READY')
-                            elif child == 'query':
-                                self._handle_ctl(xml, data)
-                            elif child == 'ping':
-                                self._handle_ping(xml, data)
-                            elif xml.get('type') == 'result':
-                                self._handle_result(data)
-                            if res:
-                                self.send(res)
-                        elif xml.tag == 'presence':
-                            if len(xml) and xml[0].tag == 'status':
-                                # bot announcing arrival
-                                self.type = self.BOT
-                                xmppserverlog.debug('{} type set to BOT (based on presence tag)'.format(self.address))
-                                # send a command from an unknown user  - the response will contain the correct admin username
-                                self.send('<iq type="set" id="{}" from="{}" to="{}"><query xmlns="com:ctl"><ctl td="GetCleanState" /></query></iq>'.format(uuid.uuid4(), 'unknown@ecouser.net', XMPPServer.bot_id))
-                            elif xml.get('type') == 'available':
-                                self.type = self.CONTROLLER
-                                xmppserverlog.debug('{} type set to CONTROLLER (based on presence tag)'.format(self.address))
-                    except ET.ParseError as e:
-                        xmppserverlog.debug('parse error - {}'.format(e))
-                    except Exception as e:
-                        xmppserverlog.exception('{}'.format(e))
-                        self._set_state('DISCONNECT')
-        except OSError as e:
-            xmppserverlog.error('{}'.format(e))
-            self._set_state('DISCONNECT')
+            
+            if self.state == self.CONNECT:
+                #Client first connecting, send our features
 
-        except ConnectionResetError as e:
-            xmppserverlog.error('{}'.format(e))
-            self._set_state('DISCONNECT')
-                        
+                if data.decode('utf-8').find('jabber:client') > -1:
+                    # ack jabbr:client
+                    #no STARTTLS
+                    self.send('<stream:stream xmlns:stream="http://etherx.jabber.org/streams" xmlns="jabber:client" version="1.0" id="1" from="{}">'.format(XMPPServer.server_id))
+                    #with STARTTLS
+                    #self.send('<stream:stream xmlns:stream="http://etherx.jabber.org/streams" xmlns:tls="http://www.ietf.org/rfc/rfc2595.txt" xmlns="jabber:client" version="1.0" id="1" from="{}">'.format(XMPPServer.server_id))
+                    time.sleep(0.25)
+                    # send authentication support for iq-auth (fallback) and SASL
+                    self.send('<stream:features><auth xmlns="http://jabber.org/features/iq-auth"/><mechanisms xmlns="urn:ietf:params:xml:ns:xmpp-sasl"><mechanism>PLAIN</mechanism></mechanisms></stream:features>')               
+                    #self.send('<stream:features><auth xmlns="http://jabber.org/features/iq-auth"/></stream:features>')               
+                                    
+                elif data.decode('utf-8').find('jabber:iq:auth') > -1: #Handle iq-auth
+                    self._handle_iq_auth(data)
+
+                elif data.decode('utf-8').find('urn:ietf:params:xml:ns:xmpp-sasl') > -1: #Handle SASL auth
+                    self._handle_sasl_auth(data)
+            
+            elif self.state == self.INIT:
+                #Client getting session after authentication
+                if data.decode('utf-8').find('jabber:client') > -1:
+                    # ack jabbr:client
+                    self.send('<stream:stream xmlns:stream="http://etherx.jabber.org/streams" xmlns="jabber:client" version="1.0" id="1" from="{}">'.format(XMPPServer.server_id))
+                    time.sleep(0.25)
+                    # session
+                    self.send('<stream:features><bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"/><session xmlns="urn:ietf:params:xml:ns:xmpp-session"/></stream:features>')                 
+
+                else: #Handle init bind
+                    xml = ET.fromstring(data.decode('utf-8'))                                
+                    if len(xml):
+                        child = self._tag_strip_uri(xml[0].tag)
+                    else:
+                        child = None
+                    
+                    if xml.tag == 'iq':
+                        if child == 'bind':
+                            self._handle_bind(xml)
+                           
         except Exception as e:
             xmppserverlog.exception('{}'.format(e))
-            self._set_state('DISCONNECT')
+
+
+    def _handle_iq_auth(self, data):                                         
+        try:
+            xml = ET.fromstring(data.decode('utf-8'))       
+            ctl = xml[0][0]                 
+            xmppserverlog.info("IQ AUTH XML: {}".format(xml))
+            #Received username and auth tag, send username/password requirement
+            if xml.get('type') == 'get' and "auth}username" in ctl.tag and self.type == self.UNKNOWN:                
+                self.send('<iq type="result" id="{}"><query xmlns="jabber:iq:auth"><username/><password/></query></iq>'.format(xml.get('id')))                               
+
+            #Received username, password, resource - Handle auth here and return pass or fail
+            if xml.get('type') == 'set' and "auth}username" in ctl.tag and self.type == self.UNKNOWN:                
+                xmlauth = xml[0].getchildren()
+                uid = ''
+                password = ''
+                resource = ''
+                for aitem in xmlauth:
+                    if 'username' in aitem.tag:
+                        self.uid = aitem.text
                         
-        finally:
-            self.disconnect()
+                    elif 'password' in aitem.tag:
+                        password = aitem.text.split("/")[2]
+
+                    elif 'resource' in aitem.tag:
+                        self.clientresource = aitem.text
+                             
+                if True:
+                    #Client authenticated, move to next state                
+                    self._set_state('INIT')    
+                    
+                    #Successful auth
+                    self.send('<iq type="result" id="{}"/>'.format(xml.get('id')))                        
+
+                else:
+                    #Failed auth
+                    self.send('<iq type="error" id="{}"><error code="401" type="auth"><not-authorized xmlns="urn:ietf:params:xml:ns:xmpp-stanzas"/></error></iq>'.format(xml.get('id')))                                                       
+              
+
+                                                    
+        except ET.ParseError as e:
+            if "no element found" in e.msg:
+                xmppserverlog.debug('xml parse error - {} - {} - this is common with ecovac protocol'.format(data.decode('utf-8'), e))                                                   
+            elif "not well-formed (invalid token)" in e.msg:
+                xmppserverlog.debug('xml parse error - {} - {}'.format(data.decode('utf-8'), e))                                                       
+            else:
+                xmppserverlog.debug('xml parse error - {} - {}'.format(data.decode('utf-8'), e))            
+
+        except Exception as e:
+            xmppserverlog.exception('{}'.format(e))  
+
+
+    def _handle_sasl_auth(self, data):                                                      
+        try:
+            xml = ET.fromstring(data.decode('utf-8'))                  
+            saslauth = base64.b64decode(xml.text).decode('utf-8').split("/")
+            username = saslauth[0]
+            username = saslauth[0].split('\x00')[1]     
+            self.uid = username       
+            resource = saslauth[1]
+            self.clientresource = resource
+            authcode = saslauth[2]
+            
+            if authcode == "us_tempauthcode": #Handle auth                
+                #Send response
+                self.send('<success xmlns="urn:ietf:params:xml:ns:xmpp-sasl"/>') #Success
+
+                #Client authenticated, move to next state                
+                self._set_state('INIT')    
+                
+            else:
+                #Failed to authenticate
+                self.send('<response xmlns="urn:ietf:params:xml:ns:xmpp-sasl"/>') #Fail
+            
+        except ET.ParseError as e:
+            if "no element found" in e.msg:
+                xmppserverlog.debug('xml parse error - {} - {} - this is common with ecovac protocol'.format(data.decode('utf-8'), e))                                                   
+            elif "not well-formed (invalid token)" in e.msg:
+                xmppserverlog.debug('xml parse error - {} - {}'.format(data.decode('utf-8'), e))                                                       
+            else:
+                xmppserverlog.debug('xml parse error - {} - {}'.format(data.decode('utf-8'), e))            
+
+        except Exception as e:
+            xmppserverlog.exception('{}'.format(e))                 
+
+    def _handle_bind(self, xml):                                         
+        try:            
+            clientbindxml = xml.getchildren()
+            clientresourcexml = clientbindxml[0].getchildren()
+            if len(clientresourcexml) > 0:
+                self.clientresource = clientresourcexml[0].text
+                self.name = "XMPP_Client_{}".format(self.clientresource)
+                xmppserverlog.debug("new client {} using resource {}".format(self.address, self.clientresource))
+                res = '<iq type="result" id="{}"><bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"><jid>{}@{}/{}</jid></bind></iq>'.format(xml.get('id'), self.uid, XMPPServer.bot_id, self.clientresource)
+            else:
+                xmppserverlog.debug("new client {}".format(self.address))
+                res = '<iq type="result" id="{}"><bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"><jid>{}</jid></bind></iq>'.format(xml.get('id'), XMPPServer.bot_id)
+            
+            self._set_state('BIND')
+            self.send(res)
+        
+        except Exception as e:
+            xmppserverlog.exception('{}'.format(e))            
+
+    def _handle_session(self, xml):
+        try:                                            
+            res = '<iq type="result" id="{}" />'.format(xml.get('id'))
+            self._set_state('READY')
+            self.send(res)
+        
+        except Exception as e:
+            xmppserverlog.exception('{}'.format(e))        
+
+    def _handle_presence(self, xml):                                                    
+        try:
+            if len(xml) and xml[0].tag == 'status':
+                # bot announcing arrival
+                self.type = self.BOT
+                xmppserverlog.debug('{} type set to BOT (based on presence tag)'.format(self.address))
+                # send a command from an unknown user  - the response will contain the correct admin username
+                
+                self.send('<iq type="set" id="{}" from="{}" to="{}"><query xmlns="com:ctl"><ctl td="GetCleanState" /></query></iq>'.format(uuid.uuid4(), 'unknown@ecouser.net', XMPPServer.bot_id))
+            
+            else:
+                self.type = self.CONTROLLER
+                xmppserverlog.debug('{} type set to CONTROLLER (based on presence tag)'.format(self.address))     
+                self.send('<presence to="{}@{}/{}"> dummy </presence>'.format(self.uid, XMPPServer.bot_id, self.clientresource))
+            
+        except Exception as e:
+            xmppserverlog.exception('{}'.format(e))                                
+
+    def _parse_data(self, data):
+        if self.log_incoming_data:
+            xmppserverlog.debug('from {} - {}'.format(self.address, data.decode('utf-8')))
+        
+        try:            
+            xml = ET.fromstring(data.decode('utf-8'))
+            self._handle_xml(xml, data)
+
+        except ET.ParseError as e:
+            if "no element found" in e.msg: #Element not closed or not all bytes received                 
+                #Happens wth connect stream often        
+                if '<stream:stream ' in data.decode('utf-8'):
+                    if self.state == self.CONNECT or self.state == self.INIT:
+                        self._handle_connect(data)                                                        
+                else:
+                    xmppserverlog.error('xml parse error - {} - {}'.format(data.decode('utf-8'), e))   
+                                                                     
+            elif "not well-formed (invalid token)" in e.msg:
+                #If a lone </stream:stream> - client is signalling end of session/disconnect
+                if not '</stream:stream>' in data.decode('utf-8'):
+                    xmppserverlog.error('xml parse error - {} - {}'.format(data.decode('utf-8'), e))  
+                else:
+                    self.send("</stream:stream>") #Close stream                    
+
+            elif "junk after document element" in e.msg: #More than one xml doc in data
+                #try to split it
+                data0 = data.decode('utf-8')
+                data1 = data0[e.position[1]:]                
+                data0 = data0[:e.position[1]]
+                #xmppserverlog.debug('xml parse error - {} - {} - split0: {} - split1: {}'.format(data.decode('utf-8'), e, data0, data1))
+                self._parse_data(data0.encode('utf-8'))
+                self._parse_data(data1.encode('utf-8'))
+                
+            else:
+                xmppserverlog.debug('xml parse error - {} - {}'.format(data.decode('utf-8'), e))                                                    
+
+        except Exception as e:
+            xmppserverlog.exception('{}'.format(e))  
+
+    def _handle_xml(self, xml, data):
+        try:
+            if self.state == self.CONNECT or self.state == self.INIT:
+                self._handle_connect(data)                
+
+            if len(xml):
+                child = self._tag_strip_uri(xml[0].tag)
+            else:
+                child = None
+            
+            if xml.tag == 'iq':
+                if child == 'bind':                    
+                    self._handle_bind(xml)
+                elif child == 'session':                    
+                    self._handle_session(xml)
+                elif child == 'query':                                        
+                    self._handle_ctl(xml, data)
+                elif child == 'ping':
+                    self._handle_ping(xml, data)
+                elif xml.get('type') == 'result':                    
+                    self._handle_result(data)                                 
+            elif xml.tag == 'presence':                
+                self._handle_presence(xml)
+
+        except Exception as e:
+            xmppserverlog.exception('{}'.format(e))              
+          
+
+    def run(self):        
+        xmppserverlog.info('client connected - {}'.format(self.address))
+        self._set_state('CONNECT')
+        
+        while not self.state == self.DISCONNECT and not self.connection._closed:
+            data = b''
+
+            time.sleep(0.2)
+            if not self.connection._closed:       
+                try:        
+                    data = self.connection.recv(4096)  
+   
+                except ConnectionResetError as e:
+                    xmppserverlog.error('{}'.format(e))                
+                except OSError as e:
+                    xmppserverlog.error('{}'.format(e))                                        
+                except Exception as e:
+                    xmppserverlog.exception('{}'.format(e))
+                            
+            if data != b'':                                                             
+                self._parse_data(data)
+