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,22 +29,23 @@ 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)
session = PivSession(conn) with conn:
if self.pin: session = PivSession(conn)
if self.pin:
try:
session.verify_pin(self.pin)
except InvalidPinError as err:
controlflow.system_error_exit(7, f'YubiKey - {err}')
try: try:
session.verify_pin(self.pin) cert = session.get_certificate(self.slot)
except InvalidPinError as err: except ApduError as err:
controlflow.system_error_exit(7, f'YubiKey - {err}') controlflow.system_error_exit(9, f'Yubikey = {err}')
try: cert_pem = cert.public_bytes(
cert = session.get_certificate(self.slot) serialization.Encoding.PEM).decode()
cert_pem = cert.public_bytes( publicKeyData = b64encode(cert_pem.encode())
serialization.Encoding.PEM).decode() if isinstance(publicKeyData, bytes):
publicKeyData = b64encode(cert_pem.encode()) publicKeyData = publicKeyData.decode()
if isinstance(publicKeyData, bytes): return publicKeyData
publicKeyData = publicKeyData.decode()
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,20 +54,21 @@ class YubiKey():
mplock.acquire() mplock.acquire()
try: try:
conn, _, _ = connect_to_device(self.serial_number) conn, _, _ = connect_to_device(self.serial_number)
session = PivSession(conn) with conn:
if self.pin: session = PivSession(conn)
if self.pin:
try:
session.verify_pin(self.pin)
except InvalidPinError as err:
controlflow.system_error_exit(7, f'YubiKey - {err}')
try: try:
session.verify_pin(self.pin) signed = session.sign(slot=self.slot,
except InvalidPinError as err:
controlflow.system_error_exit(7, f'YubiKey - {err}')
try:
signed = session.sign(slot=self.slot,
key_type=self.key_type, key_type=self.key_type,
message=message, message=message,
hash_algorithm=hashes.SHA256(), hash_algorithm=hashes.SHA256(),
padding=padding.PKCS1v15()) padding=padding.PKCS1v15())
except ApduError as err: except ApduError as err:
controlflow.system_error_exit(8, f'YubiKey = {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}')
if 'mplock' in globals(): if 'mplock' in globals():