From 32f115ca517bcbe20769684baf8ae1a4608a95e0 Mon Sep 17 00:00:00 2001
From: aritroCoder <ogcodinggod@gmail.com>
Date: Mon, 12 Aug 2024 18:05:18 +0530
Subject: [PATCH] feat: added did:web support

---
 .../present_proof/dif/pres_exch_handler.py    |  5 +++-
 aries_cloudagent/vc/vc_ld/manager.py          |  2 +-
 .../default_verification_key_strategy.py      | 23 +++++++++++++++----
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/aries_cloudagent/protocols/present_proof/dif/pres_exch_handler.py b/aries_cloudagent/protocols/present_proof/dif/pres_exch_handler.py
index 0a32d9b155..59a2dbba48 100644
--- a/aries_cloudagent/protocols/present_proof/dif/pres_exch_handler.py
+++ b/aries_cloudagent/protocols/present_proof/dif/pres_exch_handler.py
@@ -120,7 +120,10 @@ async def _get_issue_suite(
         verkey_id_strategy = self.profile.context.inject(BaseVerificationKeyStrategy)
         verification_method = (
             await verkey_id_strategy.get_verification_method_id_for_did(
-                issuer_id, self.profile, proof_purpose="assertionMethod"
+                issuer_id,
+                self.profile,
+                proof_type=self.proof_type,
+                proof_purpose="assertionMethod",
             )
         )
 
diff --git a/aries_cloudagent/vc/vc_ld/manager.py b/aries_cloudagent/vc/vc_ld/manager.py
index de3f578158..361c99e141 100644
--- a/aries_cloudagent/vc/vc_ld/manager.py
+++ b/aries_cloudagent/vc/vc_ld/manager.py
@@ -346,7 +346,7 @@ async def _get_suite_for_document(
         verification_method = (
             options.verification_method
             or await verkey_id_strategy.get_verification_method_id_for_did(
-                issuer_id, self.profile, proof_purpose="assertionMethod"
+                issuer_id, self.profile, proof_type, proof_purpose="assertionMethod"
             )
         )
 
diff --git a/aries_cloudagent/wallet/default_verification_key_strategy.py b/aries_cloudagent/wallet/default_verification_key_strategy.py
index beb4047d89..84735e7bed 100644
--- a/aries_cloudagent/wallet/default_verification_key_strategy.py
+++ b/aries_cloudagent/wallet/default_verification_key_strategy.py
@@ -19,6 +19,7 @@ async def get_verification_method_id_for_did(
         self,
         did: str,
         profile: Optional[Profile],
+        proof_type: Optional[str] = None,
         allowed_verification_method_types: Optional[List[KeyType]] = None,
         proof_purpose: Optional[str] = None,
     ) -> Optional[str]:
@@ -40,11 +41,18 @@ class DefaultVerificationKeyStrategy(BaseVerificationKeyStrategy):
 
     Supports did:key: and did:sov only.
     """
+    def __init__(self):
+        """Initialize the key types mapping."""
+        self.key_types_mapping = {
+            "Ed25519Signature2018": ["Ed25519VerificationKey2018"],
+            "Ed25519Signature2020": ["Ed25519VerificationKey2020", "Multikey"],
+        }
 
     async def get_verification_method_id_for_did(
         self,
         did: str,
         profile: Optional[Profile],
+        proof_type: Optional[str] = None,
         allowed_verification_method_types: Optional[List[KeyType]] = None,
         proof_purpose: Optional[str] = None,
     ) -> Optional[str]:
@@ -66,9 +74,14 @@ async def get_verification_method_id_for_did(
         elif did.startswith("did:web:"):
             did_resolver = profile.inject(DIDResolver)
             did_document = await did_resolver.resolve(profile=profile, did=did)
-            # Get the verification method ID from the DID document
-            # take first verification method for now
-            verification_method_list = did_document.get("verificationMethod", {})
-            verification_method_id = verification_method_list[0].get("id")
-            return verification_method_id
+            if proof_type:
+                verification_method_types = self.key_types_mapping[proof_type]
+                verification_method_list = did_document.get("verificationMethod", None)
+                for method in verification_method_list:
+                    if method.get("type") in verification_method_types:
+                        return method.get("id")
+            else:
+                # taking the first verification method from did document
+                verification_method_id = verification_method_list[0].get("id")
+                return verification_method_id
         return None