Refactor agent controller

* Use dataclass to make it more comprehensive
* Improve header and token handling
** tennant_jwt and api_key are now class attributes
** updates now attr and header
** can be removed
* remove duplicate session code
* Add docstrings

Author: morrieinmaas

libs/aries-basic-controller/aries_basic_controller/aries_controller.py

@@ -3,6 +3,7 @@
     ClientSession,
     ClientRequest,
 )
+from dataclasses import dataclass
 from pubsub import pub
 
 from .controllers.connections import ConnectionsController
@@ -25,51 +26,84 @@
 
 logger = logging.getLogger("aries_controller")
 
-class AriesAgentController:
 
+@dataclass
+class AriesAgentController:
+    """The Aries Agent Controller class
+    
+    This class allows you to interact with Aries by exposing the aca-py API. 
+    
+    Attributes
+    ----------
+    webhook_host : str
+        The url of the webhook host 
+    webhook_port : int 
+        The exposed port for webhooks on the host
+    admin_url : str
+        The URL for the Admin API
+    webhook_base : str
+        The base url for webhooks (default is "")
+    connections : bool
+        Specify whether to create connecitons (default is True)
+    messaging : bool
+        Initialise the messaging interface (default is True)
+    multitenant : bool
+        Initialise the multitenant interface (default is False)
+    mediation : bool
+        Initialise the mediation interface (default is False)
+    issuer : bool
+        Initialise the issuer interface (defautl is True)
+    action_menu : bool
+        Initialise the action menu interface (default is True)
+    revocations : bool
+        Initialise revocation interface for credentials (default is True)
+    api_key : str 
+        The API key (default is None)
+    tennant_jwt: str
+        The tennant JW token (default is None)
+    """
+    
     ## TODO rethink how to initialise. Too many args?
     ## is it important to let users config connections/issuer etc
-    def __init__(
-        self,
-        webhook_host: str,
-        webhook_port: int,
-        admin_url: str,
-        webhook_base: str = "",
-        connections: bool = True,
-        messaging: bool = True,
-        multitenant: bool = False,
-        mediation: bool = False,
-        issuer: bool = True,
-        action_menu: bool = True,
-        revocations: bool = True,
-        api_key: str = None,
-        tennant_jwt: str = None,
-    ):
-
+    webhook_host: str
+    webhook_port: int
+    admin_url: str
+    webhook_base: str = ""
+    connections: bool = True
+    messaging: bool = True
+    multitenant: bool = False
+    mediation: bool = False
+    issuer: bool = True
+    action_menu: bool = True
+    revocations: bool = True
+    api_key: str = None
+    tennant_jwt: str = None
+
+
+    def __post_init__(self):
+        """Constructs additional attributes, 
+        and logic defined by attributes set during initial instantiation
+        """
+        
         self.webhook_site = None
-        self.admin_url = admin_url
-        if webhook_base:
-            self.webhook_base = webhook_base
-        else:
-            self.webhook_base = ""
-        self.webhook_host = webhook_host
-        self.webhook_port = webhook_port
         self.connections_controller = None
-
-        headers = {}
-
-        if api_key:
-            headers.update({"X-API-Key": api_key})
-
-        if tennant_jwt:
-            headers.update({'Authorization': 'Bearer ' + tennant_jwt, 'content-type': "application/json"})
-
-        self.client_session: ClientSession = ClientSession(headers=headers)
 
+        # Construct headers for Client Session and the session itself
+        self.headers = {}
+        
+        if self.api_key:
+            self.headers.update({"X-API-Key": self.api_key})
 
-        if connections:
+        if self.tennant_jwt:
+            self.headers.update({'Authorization': 'Bearer ' + self.tennant_jwt, 'content-type': "application/json"})
+
+        self.client_session: ClientSession = ClientSession(headers=self.headers)
+    
+        # Instantiate controllers based on the provided attributes
+        if self.connections:
             self.connections = ConnectionsController(self.admin_url, self.client_session)
-        if messaging:
+            
+        if self.messaging:
             self.messaging = MessagingController(self.admin_url, self.client_session)
 
         self.proofs = ProofController(self.admin_url, self.client_session)
@@ -78,50 +112,9 @@ def __init__(
         self.server = ServerController(self.admin_url, self.client_session)
         self.oob = OOBController(self.admin_url, self.client_session)
 
-        if multitenant:
-            self.multitenant = MultitenancyController(self.admin_url, self.client_session)
-
-        if mediation:
-            self.mediation = MediationController(self.admin_url, self.client_session)
-
-        if issuer:
-            self.schema = SchemaController(self.admin_url, self.client_session)
-            self.wallet = WalletController(self.admin_url, self.client_session)
-            self.definitions = DefinitionsController(self.admin_url, self.client_session)
-            self.issuer = IssuerController(self.admin_url, self.client_session, self.connections,
-                                           self.wallet, self.definitions)
-
-        if action_menu:
-            self.action_menu = ActionMenuController(self.admin_url, self.client_session)
-
-        if revocations:
-            self.revocations = RevocationController(
-                self.admin_url,
-                self.client_session
-            )
-
-    # TODO: Determine whether we really want to essentially create a whole new ClientSession object as done below
-    # Ideally we'd just update the existing session along the lines of self.client_session(headers) 
-    # That does not work, though because it is not callable and updating cannot be achieved reliably 
-    # because headers can be of different type
-    # from https://docs.aiohttp.org/en/stable/client_reference.html :
-    # "May be either iterable of key-value pairs or Mapping (e.g. dict, CIMultiDict)."
-    # So for now let's create a new ClientSession and use all the instances current attributes 
-    # to update every attr using ClientSession
-    def update_tennant_jwt(self, tennant_jwt): 
-        self.tennant_jwt = tennant_jwt
-        headers = {'Authorization': 'Bearer ' + tennant_jwt, 'content-type': "application/json"}
-        self.client_session: ClientSession = ClientSession(headers=headers)
-
-        if self.connections:
-            self.connections = ConnectionsController(self.admin_url, self.client_session)
-
-        if self.messaging:
-            self.messaging = MessagingController(self.admin_url, self.client_session)
-
         if self.multitenant:
             self.multitenant = MultitenancyController(self.admin_url, self.client_session)
-            
+
         if self.mediation:
             self.mediation = MediationController(self.admin_url, self.client_session)
 
@@ -130,7 +123,7 @@ def update_tennant_jwt(self, tennant_jwt):
             self.wallet = WalletController(self.admin_url, self.client_session)
             self.definitions = DefinitionsController(self.admin_url, self.client_session)
             self.issuer = IssuerController(self.admin_url, self.client_session, self.connections,
-                                           self.wallet, self.definitions)
+                                            self.wallet, self.definitions)
 
         if self.action_menu:
             self.action_menu = ActionMenuController(self.admin_url, self.client_session)
@@ -141,8 +134,63 @@ def update_tennant_jwt(self, tennant_jwt):
                 self.client_session
             )
 
+    
+    def update_tennant_jwt(self, tennant_jwt: str): 
+        """Update the tenannt JW token attribute and the header
+        
+        Args:
+        ----
+        tennant_jwt : str 
+            The tennant's JW token
+        """
+        self.tennant_jwt = tennant_jwt
+        self.headers.update({'Authorization': 'Bearer ' + tennant_jwt, 'content-type': "application/json"})
+        self.client_session.headers.update(self.headers)
+        
+        
+    def update_api_key(self, api_key: str):
+        """Update the API Key attribute and the header
+        
+        Args:
+        ----
+        api_key : str
+            The API Key
+        """
+        self.api_key = api_key
+        self.headers.update({"X-API-Key": api_key})
+        self.client_session.headers.update(self.headers)
+        
+        
+    def remove_api_key(self):
+        """Removes the API key attribute and corresponding headers from the Client Session"""
+        self.api_key = None
+        if 'X-API-Key' in self.client_session.headers:
+            del self.client_session.headers['X-API-Key']
+            del self.headers['X-API-Key']
+      
+      
+    def remove_tennant_jwt(self):
+        """Removes the tennant's JW Token attribute and corresponding headers from the Client Session"""
+        self.tennant_jwt = None
+        if 'Authorization' in self.client_session.headers:
+            del self.client_session.headers['Authorization']
+            del self.headers['Authorization']
+        if 'content-type' in self.client_session.headers:
+            del self.client_session.headers['content-type']
+            del self.headers['content-type']
+
 
     def register_listeners(self, listeners, defaults=True):
+        """Registers the webhook listners
+        
+        Args:
+        ----
+        listeners : [dict]
+            A collection of dictionaries comprised of a "handler": handler (fct) and a "topic":"topicname" key-value pairs
+        defaults : bool
+            Whether to connect to the default handlers for connections, basicmessage and present_proof 
+            (default is True)
+        """
         if defaults:
             if self.connections:
                 pub.subscribe(self.connections.default_handler, "connections")
@@ -154,40 +202,91 @@ def register_listeners(self, listeners, defaults=True):
         for listener in listeners:
             self.add_listener(listener)
 
+
     def add_listener(self, listener):
+        """Subscribe to a listeners for a topic
+        
+        Args:
+        ----
+        listener : dict
+            A dictionary comprised of a "handler": handler (fct) and a "topic":"topicname" key-value pairs
+        """
         pub.subscribe(listener["handler"], listener["topic"])
 
+
     def remove_listener(self, listener):
+        """Remove a listener for a topic
+        
+        Args:
+        ----
+        listener : dict
+            A dictionary comprised of a "handler": handler (fct) and a "topic":"topicname" key-value pairs
+        """
         if pub.isSubscribed(listener["handler"], listener["topic"]):
             pub.unsubscribe(listener["handler"], listener["topic"])
         else:
             logger.debug("Listener not subscribed", listener)
 
+
     def remove_all_listeners(self, topic: str = None):
+        """Remove all listeners for one or all topics
+        
+        Args:
+        ----
+        topic : str
+            The topic to stop listening for (default is None). Default will cause unsubscribing from all topics.
+        """
         # Note advanced use of function can include both listenerFilter and topicFilter for this
         # Add when needed
         pub.unsubAll(topicName=topic)
 
+
     async def listen_webhooks(self):
+        """Create a server to listen to webhooks"""
         app = web.Application()
         app.add_routes([web.post(self.webhook_base + "/topic/{topic}/", self._receive_webhook)])
         runner = web.AppRunner(app)
         await runner.setup()
         self.webhook_site = web.TCPSite(runner, self.webhook_host, self.webhook_port)
         await self.webhook_site.start()
 
+
     async def _receive_webhook(self, request: ClientRequest):
+        """Helper to receive webhooks by requesting it
+        
+        Args:
+        ----
+        request : ClientRequest
+            The client request to which the corresponding webhooks shall be received
+            
+        Returns:
+        -------
+        Response:
+            A response with status 200
+        """
         topic = request.match_info["topic"]
         payload = await request.json()
         await self._handle_webhook(topic, payload)
         return web.Response(status=200)
 
+
     async def _handle_webhook(self, topic, payload):
+        """Helper handling a webhook
+        
+        Args:
+        ----
+        topic : str
+            The topic to handle webhooks for
+        payload : dict
+            A JSON-like dictionary representation of the payload
+        """
         logging.debug(f"Handle Webhook - {topic}", payload)
         pub.sendMessage(topic, payload=payload)
         # return web.Response(status=200)
 
+
     async def terminate(self):
+        """Terminate the controller client session and webhook listeners"""
         await self.client_session.close()
         if self.webhook_site:
             await self.webhook_site.stop()