Draft: Resolve "Add websockets"

Closes #11

Pipfile

 pyrabbit = "*"
 twisted = "*"
 websockets = "*"
+requests = "*"
 
 [dev-packages]
 

api_wrap_socket.py

 # Source: https://python.plainenglish.io/identify-websocket-clients-with-autobahn-twisted-and-python-3f90b4c135d4
 # Source: https://stackoverflow.com/questions/29951718/autobahn-sending-user-specific-and-broadcast-messages-from-external-application
 
+import concurrent.futures
 import uuid
-from twisted.internet import reactor
+import time
+from twisted.internet import reactor, threads
 from twisted.web.server import Site
 from twisted.web.resource import Resource
 from twisted.web.wsgi import WSGIResource
 from endpoints.clients import ClientsFile
 from models.token_model import TokenModel
 from api_rest import app
+from rmq import write_to_queue, listen, stop_listening
 
+SUBSCRIPTIONS = {}
+PORT = 8087
 
 class SubscriptionServerProtocol(WebSocketServerProtocol):
     def __init__(self):
         WebSocketServerProtocol.__init__(self)
         self.token = TokenModel()
-        self.apiClients = ClientsFile()
+        self.apiClientsFile = ClientsFile()
+        self.apiClients = self.apiClientsFile.get()
         self.client = None
 
     def auth(self, request): 
             print(str(parsed))
             if valid:
                 try:
+                    client_id = parsed.get('client_id')
+                    print(f"client_id: {client_id}")
                     client = self.apiClients.get(parsed.get('client_id'))
+                    print(str(client))                    
                     if client: 
                         self.client = client
                         allow = True
                 except Exception: 
                     allow = False
-            # allow = valid
         return allow
         
     def onOpen(self):
         print("Client connecting: {}".format(request.peer))
         allow = self.auth(request)
         print(f"Allowed: {allow}")
-        # if allow:
-        #     setup queue consumer 
-
+        if allow:
+            # setup queue consumer
+            print("call createListener")
+            SubscriptionListener.createListener(self)
+             
     def onMessage(self, payload, isBinary):
-        if not isBinary:
-            if (msg := payload.decode("utf-8")).startswith("@a"):
-                self.factory.broadcast("Got message: " + msg)
-                # identify client 
-                # deliver to client-outbox rmq
-
+        
+        # if not isBinary:
+            # if (msg := payload.decode("utf-8")).startswith("@a"):
+            #     self.factory.broadcast("Got message: " + msg)
+            # identify client 
+            # deliver to client-outbox rmq
+        if self.client:
+            message = payload.decode('utf-8')
+            print(f"ws:onMessage message: {message}")
+            client_outbox = f"{self.client['client_id']}-outbox"
+            write_to_queue(client_outbox, message)
+            
     def connectionLost(self, reason):
         WebSocketServerProtocol.connectionLost(self, reason)
         self.factory.unregister(self)
+        # stop listener
+        if self.client: 
+            SubscriptionListener.stopListening(self.client['client_id'])
 
 
+class SubscriptionListener():
+    def __init__(self, socket):
+        self.socket = socket
+        self.running = False
+        self.thread = None
+        self.tag = None
+    
+    def listenToInbox(self):
+        client_id = self.socket.client['client_id']
+        client_inbox = f"{client_id}-inbox"
+        self.tag = f"{client_id}-socket-consumer"
+        listen(client_inbox, self.receive, self.tag)
+    
+    def receive(self, ch, method, properties, body):
+        try: 
+            self.socket.sendMessage(body)
+            ch.basic_ack(delivery_tag=method.delivery_tag)
+        except Disconnected: 
+            ch.basic_reject(delivery_tag=method.delivery_tag)
+
+    @staticmethod
+    def createListener(socket):
+        client_id = socket.client['client_id']
+        print(f"createListener: {client_id}")
+        listener = SubscriptionListener(socket)
+        SUBSCRIPTIONS[client_id] = listener
+        return listener
+
+    @staticmethod
+    def monitorSubscriptions():
+        while True:
+            for client_id, subscription in SUBSCRIPTIONS.items():
+                if not subscription.running: 
+                    subscription.running = True
+                    reactor.callInThread(subscription.listenToInbox)
+                    
+            time.sleep(5)
+    
+    @staticmethod
+    def stopListening(client_id):
+        sub = SUBSCRIPTIONS.get(client_id)
+        if sub:    
+            del SUBSCRIPTIONS[client_id]
+            stop_listening(sub.tag)
+            
+ 
 class SubscriptionServerFactory(WebSocketServerFactory):
     def __init__(self, url):
         WebSocketServerFactory.__init__(self, url)
         self.clients = {}
-
+    
     def register(self, client):
         registered = [self.clients[i] for i in list(self.clients.keys())]
         ids = list(self.clients.keys())
         
         for cid in disconnected:  
             del self.clients[cid]
-    
-    def deliver(self, cid, msg):
-        print("delivering message '{}' to {} client ...".format(msg, cid))
-        if cid in self.clients:
-            self.clients[cid].sendMessage(msg.encode('utf-8'))
 
 
 def run():
     ServerFactory = SubscriptionServerFactory
-    factory = ServerFactory("ws://localhost:8087")
+    factory = ServerFactory(f"ws://localhost:{PORT}")
 
     factory.protocol = SubscriptionServerProtocol
-    # listenWS(factory)
     socket_resource = WebSocketResource(factory)
 
     flask_resource = WSGIResource(reactor, reactor.getThreadPool(), app)
     root.putChild(b'api', flask_resource)
     root.putChild(b'ws', socket_resource)
 
-    reactor.listenTCP(8087, Site(root))
+    reactor.listenTCP(PORT, Site(root))
 
+    reactor.callInThread(SubscriptionListener.monitorSubscriptions)
+    
     reactor.run()
 
 
 if __name__ == "__main__":
-    run()
\ No newline at end of file
+    run()
+        

get_token.py

+# export SOAR_TOKEN=\
+# $(bash -c "cd to/repo/root && pipenv run python ./get_token.py http://localhost:8087/api [client-id]")
+import requests
+import sys
+from endpoints.clients import ClientsFile
+
+api_root = sys.argv[1]
+client_id = sys.argv[2]
+
+clients_file = ClientsFile()
+client = clients_file.find(client_id)
+
+credentials = {
+  "client_id": client['client_id'],
+  "secret": client['secret']
+}
+
+response = requests.get(f"{api_root}/token", params=credentials).json()
+token = response['token']
+
+print(token)
\ No newline at end of file

rmq.py

             delivery_mode=pika.spec.PERSISTENT_DELIVERY_MODE
         ),
     )
-
+    
     connection.close()