Source code for ykman.scripting

# Copyright (c) 2021 Yubico AB
# All rights reserved.
#
#   Redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following
#   conditions are met:
#
#    1. Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#    2. Redistributions in binary form must reproduce the above
#       copyright notice, this list of conditions and the following
#       disclaimer in the documentation and/or other materials provided
#       with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.


from time import sleep
from typing import Generator, Optional

from smartcard.Exceptions import CardConnectionException, NoCardException

from yubikit.core import TRANSPORT
from yubikit.core.fido import FidoConnection
from yubikit.core.otp import OtpConnection
from yubikit.core.smartcard import SmartCardConnection
from yubikit.management import DeviceInfo
from yubikit.support import get_name, read_info

from .base import YkmanDevice
from .device import list_all_devices, scan_devices
from .pcsc import list_devices as list_ccid

"""
Various helpers intended to simplify scripting.

Add an import to your script:

  from ykman import scripting as s

Example usage:

  yubikey = s.single()
  print("Here is a YubiKey:", yubikey)


  print("Insert multiple YubiKeys")
  for yubikey in s.multi():
      print("You inserted {yubikey}")
  print("You pressed Ctrl+C, end of script")

"""


[docs] class ScriptingDevice: """Scripting-friendly proxy for YkmanDevice. This wrapper adds some helpful utility methods useful for scripting. """ def __init__(self, wrapped, info): self._wrapped = wrapped self._info = info self._name = get_name(info, self.pid.yubikey_type if self.pid else None)
[docs] def __getattr__(self, attr): return getattr(self._wrapped, attr)
[docs] def __str__(self): serial = self._info.serial return f"{self._name} ({serial})" if serial else self._name
@property def info(self) -> DeviceInfo: return self._info @property def name(self) -> str: return self._name
[docs] def otp(self) -> OtpConnection: """Establish a OTP connection.""" return self.open_connection(OtpConnection)
[docs] def smart_card(self) -> SmartCardConnection: """Establish a Smart Card connection.""" return self.open_connection(SmartCardConnection)
[docs] def fido(self) -> FidoConnection: """Establish a FIDO connection.""" return self.open_connection(FidoConnection)
YkmanDevice.register(ScriptingDevice)
[docs] def single(*, prompt=True) -> ScriptingDevice: """Connect to a YubiKey. :param prompt: When set, you will be prompted to insert a YubiKey. """ pids, state = scan_devices() n_devs = sum(pids.values()) if prompt and n_devs == 0: print("Insert YubiKey...") while n_devs == 0: sleep(1.0) pids, new_state = scan_devices() n_devs = sum(pids.values()) devs = list_all_devices() if len(devs) == 1: return ScriptingDevice(*devs[0]) raise ValueError("Failed to get single YubiKey")
[docs] def multi( *, ignore_duplicates: bool = True, allow_initial: bool = False, prompt: bool = True ) -> Generator[ScriptingDevice, None, None]: """Connect to multiple YubiKeys. :param ignore_duplicates: When set, duplicates are ignored. :param allow_initial: When set, YubiKeys can be connected at the start of the function call. :param prompt: When set, you will be prompted to insert a YubiKey. """ state = None handled_serials: set[Optional[int]] = set() pids, _ = scan_devices() n_devs = sum(pids.values()) if n_devs == 0: if prompt: print("Insert YubiKeys, one at a time...") elif not allow_initial: raise ValueError("YubiKeys must not be present initially.") while True: # Run this until we stop the script with Ctrl+C pids, new_state = scan_devices() if new_state != state: state = new_state # State has changed serials = set() if len(pids) == 0 and None in handled_serials: handled_serials.remove(None) # Allow one key without serial at a time for device, info in list_all_devices(): serials.add(info.serial) if info.serial not in handled_serials: handled_serials.add(info.serial) yield ScriptingDevice(device, info) if not ignore_duplicates: # Reset handled serials to currently connected handled_serials = serials else: try: sleep(1.0) # No change, sleep for 1 second. except KeyboardInterrupt: return # Stop waiting
def _get_reader(reader) -> YkmanDevice: readers = [d for d in list_ccid(reader) if d.transport == TRANSPORT.NFC] if not readers: raise ValueError(f"No NFC reader found matching filter: '{reader}'") elif len(readers) > 1: names = [r.fingerprint for r in readers] raise ValueError(f"Multiple NFC readers matching filter: '{reader}' {names}") return readers[0]
[docs] def single_nfc(reader="", *, prompt=True) -> ScriptingDevice: """Connect to a YubiKey over NFC. :param reader: The name of the NFC reader. :param prompt: When set, you will prompted to place a YubiKey on NFC reader. """ device = _get_reader(reader) while True: try: with device.open_connection(SmartCardConnection) as connection: info = read_info(connection) return ScriptingDevice(device, info) except NoCardException: if prompt: print("Place YubiKey on NFC reader...") prompt = False sleep(1.0)
[docs] def multi_nfc( reader="", *, ignore_duplicates=True, allow_initial=False, prompt=True ) -> Generator[ScriptingDevice, None, None]: """Connect to multiple YubiKeys over NFC. :param reader: The name of the NFC reader. :param ignore_duplicates: When set, duplicates are ignored. :param allow_initial: When set, YubiKeys can be connected at the start of the function call. :param prompt: When set, you will be prompted to place YubiKeys on the NFC reader. """ device = _get_reader(reader) prompted = False try: with device.open_connection(SmartCardConnection) as connection: if not allow_initial: raise ValueError("YubiKey must not be present initially.") except NoCardException: if prompt: print("Place YubiKey on NFC reader...") prompted = True sleep(1.0) handled_serials: set[Optional[int]] = set() current: Optional[int] = -1 while True: # Run this until we stop the script with Ctrl+C try: with device.open_connection(SmartCardConnection) as connection: info = read_info(connection) if info.serial in handled_serials or current == info.serial: if prompt and not prompted: print("Remove YubiKey from NFC reader.") prompted = True else: current = info.serial if ignore_duplicates: handled_serials.add(current) yield ScriptingDevice(device, info) prompted = False except NoCardException: if None in handled_serials: handled_serials.remove(None) # Allow one key without serial at a time current = -1 if prompt and not prompted: print("Place YubiKey on NFC reader...") prompted = True except CardConnectionException: pass try: sleep(1.0) # No change, sleep for 1 second. except KeyboardInterrupt: return # Stop waiting