Use with conn so Yubikey connections close sooner

This commit is contained in:
Jay Lee
2021-07-26 14:46:58 -04:00
parent 375e36ff96
commit ed20fe252e

View File

@@ -29,6 +29,7 @@ class YubiKey():
def get_certificate(self): def get_certificate(self):
try: try:
conn, _, _ = connect_to_device(self.serial_number) conn, _, _ = connect_to_device(self.serial_number)
with conn:
session = PivSession(conn) session = PivSession(conn)
if self.pin: if self.pin:
try: try:
@@ -37,14 +38,14 @@ class YubiKey():
controlflow.system_error_exit(7, f'YubiKey - {err}') controlflow.system_error_exit(7, f'YubiKey - {err}')
try: try:
cert = session.get_certificate(self.slot) cert = session.get_certificate(self.slot)
except ApduError as err:
controlflow.system_error_exit(9, f'Yubikey = {err}')
cert_pem = cert.public_bytes( cert_pem = cert.public_bytes(
serialization.Encoding.PEM).decode() serialization.Encoding.PEM).decode()
publicKeyData = b64encode(cert_pem.encode()) publicKeyData = b64encode(cert_pem.encode())
if isinstance(publicKeyData, bytes): if isinstance(publicKeyData, bytes):
publicKeyData = publicKeyData.decode() publicKeyData = publicKeyData.decode()
return publicKeyData return publicKeyData
except ApduError as err:
controlflow.system_error_exit(8, f'YubiKey - {err}')
except ValueError as err: except ValueError as err:
controlflow.system_error_exit(9, f'YubiKey - {err}') controlflow.system_error_exit(9, f'YubiKey - {err}')
@@ -53,6 +54,7 @@ class YubiKey():
mplock.acquire() mplock.acquire()
try: try:
conn, _, _ = connect_to_device(self.serial_number) conn, _, _ = connect_to_device(self.serial_number)
with conn:
session = PivSession(conn) session = PivSession(conn)
if self.pin: if self.pin:
try: try: