From 4aed54f34a8356ab51d33080540da283a7f5a160 Mon Sep 17 00:00:00 2001
From: oll-bot <developers@openlawlib.org>
Date: Sun, 9 Jun 2024 01:53:10 -0400
Subject: [PATCH] - Updated code to support Yubikey Manager 5.1.x - Dropped
 support for Yubikey Manager 4.x

---
 setup.py       |  2 +-
 taf/yubikey.py | 51 ++++++++++++++++++++++++++------------------------
 2 files changed, 28 insertions(+), 25 deletions(-)

diff --git a/setup.py b/setup.py
index d73c9786b..e7c998c47 100644
--- a/setup.py
+++ b/setup.py
@@ -45,7 +45,7 @@ def finalize_options(self):
     "jsonschema==3.2.0",
 ]
 
-yubikey_require = ["yubikey-manager==4.0.*"]
+yubikey_require = ["yubikey-manager==5.1.*"]
 
 # Determine the appropriate version of pygit2 based on the Python version
 if sys.version_info > (3, 10):
diff --git a/taf/yubikey.py b/taf/yubikey.py
index 250dbe109..6bfa76a55 100644
--- a/taf/yubikey.py
+++ b/taf/yubikey.py
@@ -14,7 +14,7 @@
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.asymmetric import rsa, padding
 from tuf.repository_tool import import_rsakey_from_pem
-from ykman.device import list_all_devices, connect_to_device
+from ykman.device import list_all_devices
 from yubikit.core.smartcard import SmartCardConnection
 from ykman.piv import (
     KEY_TYPE,
@@ -110,31 +110,34 @@ def _yk_piv_ctrl(serial=None, pub_key_pem=None):
     # If pub_key_pem is given, iterate all devices, read x509 certs and try to match
     # public keys.
     if pub_key_pem is not None:
-        for _, info in list_all_devices():
-            connection, _, device = connect_to_device(
-                info.serial, [SmartCardConnection]
-            )
-            session = PivSession(connection)
-            device_pub_key_pem = (
-                session.get_certificate(SLOT.SIGNATURE)
-                .public_key()
-                .public_bytes(
-                    encoding=serialization.Encoding.PEM,
-                    format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        for dev, info in list_all_devices():
+            # Connect to a YubiKey over a SmartCardConnection, which is needed for PIV.
+            with dev.open_connection(SmartCardConnection) as connection:
+                session = PivSession(connection)
+                device_pub_key_pem = (
+                    session.get_certificate(SLOT.SIGNATURE)
+                    .public_key()
+                    .public_bytes(
+                        encoding=serialization.Encoding.PEM,
+                        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+                    )
+                    .decode("utf-8")
                 )
-                .decode("utf-8")
-            )
-            # Tries to match without last newline char
-            if (
-                device_pub_key_pem == pub_key_pem
-                or device_pub_key_pem[:-1] == pub_key_pem
-            ):
-                break
-            yield session, device.serial
+                # Tries to match without last newline char
+                if (
+                    device_pub_key_pem == pub_key_pem
+                    or device_pub_key_pem[:-1] == pub_key_pem
+                ):
+                    break
+                yield session, info.serial
     else:
-        connection, _, device = connect_to_device(serial, [SmartCardConnection])
-        session = PivSession(connection)
-        yield session, device.serial
+        for dev, info in list_all_devices():
+            if info.serial == serial:
+                with dev.open_connection(SmartCardConnection) as connection:
+                    session = PivSession(connection)
+                    yield session, info.serial
+            else:
+                pass
 
 
 def is_inserted():