Merge pull request #744 from andrewwhitehead/feat/custom-prefix

Allow custom message prefixes; simplify message type registration
openwallet-foundation · Oct 7, 2020 · c402a60 · c402a60
2 parents eae7a4e + 04ee93e
commit c402a60
Show file tree

Hide file tree

Showing 13 changed files with 124 additions and 173 deletions.
diff --git a/aries_cloudagent/protocols/actionmenu/v1_0/message_types.py b/aries_cloudagent/protocols/actionmenu/v1_0/message_types.py
@@ -9,24 +9,14 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.actionmenu.v1_0"
 
-MESSAGE_TYPES = {
-    **{
-        pfx.qualify(MENU): f"{PROTOCOL_PACKAGE}.messages.menu.Menu"
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(MENU_REQUEST): (
-            f"{PROTOCOL_PACKAGE}.messages.menu_request.MenuRequest"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(PERFORM): f"{PROTOCOL_PACKAGE}.messages.perform.Perform"
-        for pfx in DIDCommPrefix
-    },
-}
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {
+        MENU: f"{PROTOCOL_PACKAGE}.messages.menu.Menu",
+        MENU_REQUEST: f"{PROTOCOL_PACKAGE}.messages.menu_request.MenuRequest",
+        PERFORM: f"{PROTOCOL_PACKAGE}.messages.perform.Perform",
+    }
+)
 
-CONTROLLERS = {
-    pfx.qualify("action-menu/1.0"): f"{PROTOCOL_PACKAGE}.controller.Controller"
-    for pfx in DIDCommPrefix
-}
+CONTROLLERS = DIDCommPrefix.qualify_all(
+    {"action-menu/1.0": f"{PROTOCOL_PACKAGE}.controller.Controller"}
+)
diff --git a/aries_cloudagent/protocols/basicmessage/v1_0/message_types.py b/aries_cloudagent/protocols/basicmessage/v1_0/message_types.py
@@ -12,7 +12,6 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.basicmessage.v1_0"
 
-MESSAGE_TYPES = {
-    pfx.qualify(BASIC_MESSAGE): f"{PROTOCOL_PACKAGE}.messages.basicmessage.BasicMessage"
-    for pfx in DIDCommPrefix
-}
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {BASIC_MESSAGE: f"{PROTOCOL_PACKAGE}.messages.basicmessage.BasicMessage"}
+)
diff --git a/aries_cloudagent/protocols/connections/v1_0/message_types.py b/aries_cloudagent/protocols/connections/v1_0/message_types.py
@@ -15,29 +15,17 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.connections.v1_0"
 
-MESSAGE_TYPES = {
-    **{
-        pfx.qualify(CONNECTION_INVITATION): (
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {
+        CONNECTION_INVITATION: (
             f"{PROTOCOL_PACKAGE}.messages.connection_invitation.ConnectionInvitation"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(CONNECTION_REQUEST): (
+        ),
+        CONNECTION_REQUEST: (
             f"{PROTOCOL_PACKAGE}.messages.connection_request.ConnectionRequest"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(CONNECTION_RESPONSE): (
+        ),
+        CONNECTION_RESPONSE: (
             f"{PROTOCOL_PACKAGE}.messages.connection_response.ConnectionResponse"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(PROBLEM_REPORT): (
-            f"{PROTOCOL_PACKAGE}.messages.problem_report.ProblemReport"
-        )
-        for pfx in DIDCommPrefix
-    },
-}
+        ),
+        PROBLEM_REPORT: f"{PROTOCOL_PACKAGE}.messages.problem_report.ProblemReport",
+    }
+)
diff --git a/aries_cloudagent/protocols/didcomm_prefix.py b/aries_cloudagent/protocols/didcomm_prefix.py
@@ -1,9 +1,18 @@
 """DIDComm prefix management."""
+import re
 
 from enum import Enum
 from os import environ
 from typing import Mapping
 
+QUALIFIED = re.compile(r"^[a-zA-Z\-\+]+:.+")
+
+
+def qualify(msg_type: str, prefix: str):
+    """Qualify a message type with a prefix, if unqualified."""
+
+    return msg_type if QUALIFIED.match(msg_type or "") else f"{prefix}/{msg_type}"
+
 
 class DIDCommPrefix(Enum):
     """Enum for DIDComm Prefix, old or new style, per Aries RFC 384."""
@@ -24,13 +33,23 @@ def set(settings: Mapping):
     def qualify(self, msg_type: str = None) -> str:
         """Qualify input message type with prefix and separator."""
 
-        return f"{self.value}/{msg_type or ''}"
+        return qualify(msg_type, self.value)
+
+    @classmethod
+    def qualify_all(cls, messages: dict) -> dict:
+        """Apply all known prefixes to a dictionary of message types."""
+
+        result = {}
+        for pfx in cls:
+            for k, v in messages.items():
+                result[qualify(k, pfx.value)] = v
+        return result
 
     @staticmethod
     def qualify_current(slug: str = None) -> str:
         """Qualify input slug with prefix currently in effect and separator."""
 
-        return f"{environ.get('DIDCOMM_PREFIX', DIDCommPrefix.NEW.value)}/{slug or ''}"
+        return qualify(slug, environ.get("DIDCOMM_PREFIX", DIDCommPrefix.NEW.value))
 
     @staticmethod
     def unqualify(qual: str) -> str:

diff --git a/aries_cloudagent/protocols/discovery/v1_0/message_types.py b/aries_cloudagent/protocols/discovery/v1_0/message_types.py
@@ -13,13 +13,9 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.discovery.v1_0"
 
-MESSAGE_TYPES = {
-    **{
-        pfx.qualify(DISCLOSE): (f"{PROTOCOL_PACKAGE}.messages.disclose.Disclose")
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(QUERY): (f"{PROTOCOL_PACKAGE}.messages.query.Query")
-        for pfx in DIDCommPrefix
-    },
-}
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {
+        DISCLOSE: f"{PROTOCOL_PACKAGE}.messages.disclose.Disclose",
+        QUERY: f"{PROTOCOL_PACKAGE}.messages.query.Query",
+    }
+)
diff --git a/aries_cloudagent/protocols/introduction/v0_1/message_types.py b/aries_cloudagent/protocols/introduction/v0_1/message_types.py
@@ -8,21 +8,14 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.introduction.v0_1"
 
-MESSAGE_TYPES = {
-    **{
-        pfx.qualify(INVITATION_REQUEST): (
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {
+        INVITATION_REQUEST: (
             f"{PROTOCOL_PACKAGE}.messages.invitation_request.InvitationRequest"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(INVITATION): (f"{PROTOCOL_PACKAGE}.messages.invitation.Invitation")
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(FORWARD_INVITATION): (
+        ),
+        INVITATION: f"{PROTOCOL_PACKAGE}.messages.invitation.Invitation",
+        FORWARD_INVITATION: (
             f"{PROTOCOL_PACKAGE}.messages.forward_invitation.ForwardInvitation"
-        )
-        for pfx in DIDCommPrefix
-    },
-}
+        ),
+    }
+)
diff --git a/aries_cloudagent/protocols/issue_credential/v1_0/message_types.py b/aries_cloudagent/protocols/issue_credential/v1_0/message_types.py
@@ -16,38 +16,23 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.issue_credential.v1_0"
 
-MESSAGE_TYPES = {
-    **{
-        pfx.qualify(CREDENTIAL_PROPOSAL): (
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {
+        CREDENTIAL_PROPOSAL: (
             f"{PROTOCOL_PACKAGE}.messages.credential_proposal.CredentialProposal"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(CREDENTIAL_OFFER): (
+        ),
+        CREDENTIAL_OFFER: (
             f"{PROTOCOL_PACKAGE}.messages.credential_offer.CredentialOffer"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(CREDENTIAL_REQUEST): (
+        ),
+        CREDENTIAL_REQUEST: (
             f"{PROTOCOL_PACKAGE}.messages.credential_request.CredentialRequest"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(CREDENTIAL_ISSUE): (
+        ),
+        CREDENTIAL_ISSUE: (
             f"{PROTOCOL_PACKAGE}.messages.credential_issue.CredentialIssue"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(CREDENTIAL_ACK): (
-            f"{PROTOCOL_PACKAGE}.messages.credential_ack.CredentialAck"
-        )
-        for pfx in DIDCommPrefix
-    },
-}
+        ),
+        CREDENTIAL_ACK: f"{PROTOCOL_PACKAGE}.messages.credential_ack.CredentialAck",
+    }
+)
 
 # Inner object types
 CREDENTIAL_PREVIEW = f"issue-credential/1.0/credential-preview"

diff --git a/aries_cloudagent/protocols/out_of_band/v1_0/message_types.py b/aries_cloudagent/protocols/out_of_band/v1_0/message_types.py
@@ -12,7 +12,6 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.out_of_band.v1_0"
 
-MESSAGE_TYPES = {
-    pfx.qualify(INVITATION): f"{PROTOCOL_PACKAGE}.messages.invitation.Invitation"
-    for pfx in DIDCommPrefix
-}
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {INVITATION: f"{PROTOCOL_PACKAGE}.messages.invitation.Invitation"}
+)
diff --git a/aries_cloudagent/protocols/present_proof/v1_0/message_types.py b/aries_cloudagent/protocols/present_proof/v1_0/message_types.py
@@ -15,32 +15,20 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.present_proof.v1_0"
 
-MESSAGE_TYPES = {
-    **{
-        pfx.qualify(PRESENTATION_PROPOSAL): (
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {
+        PRESENTATION_PROPOSAL: (
             f"{PROTOCOL_PACKAGE}.messages.presentation_proposal.PresentationProposal"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(PRESENTATION_REQUEST): (
+        ),
+        PRESENTATION_REQUEST: (
             f"{PROTOCOL_PACKAGE}.messages.presentation_request.PresentationRequest"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(PRESENTATION): (
-            f"{PROTOCOL_PACKAGE}.messages.presentation.Presentation"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(PRESENTATION_ACK): (
+        ),
+        PRESENTATION: f"{PROTOCOL_PACKAGE}.messages.presentation.Presentation",
+        PRESENTATION_ACK: (
             f"{PROTOCOL_PACKAGE}.messages.presentation_ack.PresentationAck"
-        )
-        for pfx in DIDCommPrefix
-    },
-}
+        ),
+    }
+)
 
 # Inner object types
 PRESENTATION_PREVIEW = f"present-proof/1.0/presentation-preview"

diff --git a/aries_cloudagent/protocols/problem_report/v1_0/message_types.py b/aries_cloudagent/protocols/problem_report/v1_0/message_types.py
@@ -12,7 +12,6 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.problem_report.v1_0"
 
-MESSAGE_TYPES = {
-    pfx.qualify(PROBLEM_REPORT): f"{PROTOCOL_PACKAGE}.message.ProblemReport"
-    for pfx in DIDCommPrefix
-}
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {PROBLEM_REPORT: f"{PROTOCOL_PACKAGE}.message.ProblemReport"}
+)
diff --git a/aries_cloudagent/protocols/routing/v1_0/message_types.py b/aries_cloudagent/protocols/routing/v1_0/message_types.py
@@ -11,33 +11,20 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.routing.v1_0"
 
-MESSAGE_TYPES = {
-    **{
-        pfx.qualify(FORWARD): f"{PROTOCOL_PACKAGE}.messages.forward.Forward"
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(ROUTE_QUERY_REQUEST): (
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {
+        FORWARD: f"{PROTOCOL_PACKAGE}.messages.forward.Forward",
+        ROUTE_QUERY_REQUEST: (
             f"{PROTOCOL_PACKAGE}.messages.route_query_request.RouteQueryRequest"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(ROUTE_QUERY_RESPONSE): (
+        ),
+        ROUTE_QUERY_RESPONSE: (
             f"{PROTOCOL_PACKAGE}.messages.route_query_response.RouteQueryResponse"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(ROUTE_UPDATE_REQUEST): (
+        ),
+        ROUTE_UPDATE_REQUEST: (
             f"{PROTOCOL_PACKAGE}.messages.route_update_request.RouteUpdateRequest"
-        )
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(ROUTE_UPDATE_RESPONSE): (
+        ),
+        ROUTE_UPDATE_RESPONSE: (
             f"{PROTOCOL_PACKAGE}.messages.route_update_response.RouteUpdateResponse"
-        )
-        for pfx in DIDCommPrefix
-    },
-}
+        ),
+    }
+)
diff --git a/aries_cloudagent/protocols/tests/test_didcomm_prefix.py b/aries_cloudagent/protocols/tests/test_didcomm_prefix.py
@@ -6,7 +6,7 @@
 
 
 class TestDIDCommPrefix(AsyncTestCase):
-    async def test_didcomm_prefix(self):
+    def test_didcomm_prefix(self):
         DIDCommPrefix.set({})
         assert environ.get("DIDCOMM_PREFIX") == DIDCommPrefix.OLD.value
 
@@ -31,3 +31,17 @@ async def test_didcomm_prefix(self):
         assert DIDCommPrefix.unqualify(new_q_hello) == "hello"
         assert DIDCommPrefix.unqualify("already unqualified") == "already unqualified"
         assert DIDCommPrefix.unqualify(None) is None
+
+        qualified = "http://custom-prefix/message/type/1.0"
+        assert DIDCommPrefix.qualify_current(qualified) == qualified
+        assert DIDCommPrefix.NEW.qualify(qualified) == qualified
+        assert DIDCommPrefix.unqualify(qualified) == qualified
+
+    def test_didcomm_qualify_all(self):
+        mtype = "a/message/1.0"
+        mcls = "protocol.Cls"
+        messages = {mtype: mcls}
+        assert DIDCommPrefix.qualify_all(messages) == {
+            DIDCommPrefix.NEW.qualify(mtype): mcls,
+            DIDCommPrefix.OLD.qualify(mtype): mcls,
+        }
diff --git a/aries_cloudagent/protocols/trustping/v1_0/message_types.py b/aries_cloudagent/protocols/trustping/v1_0/message_types.py
@@ -13,15 +13,9 @@
 
 PROTOCOL_PACKAGE = "aries_cloudagent.protocols.trustping.v1_0"
 
-MESSAGE_TYPES = {
-    **{
-        pfx.qualify(PING): f"{PROTOCOL_PACKAGE}.messages.ping.Ping"
-        for pfx in DIDCommPrefix
-    },
-    **{
-        pfx.qualify(PING_RESPONSE): (
-            f"{PROTOCOL_PACKAGE}.messages.ping_response.PingResponse"
-        )
-        for pfx in DIDCommPrefix
-    },
-}
+MESSAGE_TYPES = DIDCommPrefix.qualify_all(
+    {
+        PING: f"{PROTOCOL_PACKAGE}.messages.ping.Ping",
+        PING_RESPONSE: f"{PROTOCOL_PACKAGE}.messages.ping_response.PingResponse",
+    }
+)