From 4bcabf1101a56e341096493370c791a426e771a4 Mon Sep 17 00:00:00 2001
From: hyugogirubato <65763543+hyugogirubato@users.noreply.github.com>
Date: Sun, 7 Jul 2024 16:27:47 +0200
Subject: [PATCH] Improve library vendor detection

---
 keydive/__main__.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/keydive/__main__.py b/keydive/__main__.py
index d2b80b2..8e9c57c 100644
--- a/keydive/__main__.py
+++ b/keydive/__main__.py
@@ -107,6 +107,17 @@ def main() -> None:
         # Initialize Core instance for interacting with the device
         core = Core(cdm=cdm, device=args.device, functions=args.functions)
 
+        # Map process keys to their compatible CDM vendors
+        cdm_vendor = {}
+        for key, vendors in CDM_VENDOR_API.items():
+            for vendor in vendors:
+                # Check if vendor's SDK matches the core's SDK or the previous version
+                if vendor.sdk in (core.sdk_api, core.sdk_api - 1):
+                    cdm_vendor.setdefault(key, []).append(vendor)
+
+        if not cdm_vendor:
+            raise NotImplementedError('SDK version is not supported.')
+
         # Process watcher loop
         logger.info('Watcher delay: %ss' % args.delay)
         current = None
@@ -115,7 +126,7 @@ def main() -> None:
             processes = {
                 key: (name, pid)
                 for name, pid in core.enumerate_processes().items()
-                for key in CDM_VENDOR_API.keys() if key in name
+                for key in cdm_vendor.keys() if key in name
             }
 
             if not processes:
@@ -136,7 +147,7 @@ def main() -> None:
                 for key, (name, pid) in processes.items():
                     if current:
                         break
-                    for vendor in CDM_VENDOR_API[key]:
+                    for vendor in cdm_vendor[key]:
                         if core.hook_process(pid=pid, vendor=vendor):
                             logger.info('Process: %s (%s)', pid, name)
                             current = pid