Add ping, lock, fix error read, read capflags, fix device open error

u2flib_host/hid_transport.py

@@ -57,18 +57,24 @@
 U2F_VENDOR_FIRST = 0x40
 
 # USB Commands
+CMD_PING = 0x01 # https://github.com/Yubico/python-u2flib-host/pull/18 & 
+CMD_APDU = 0x03
+CMD_LOCK = 0x04 # https://github.com/peter-conalgo/python-u2flib-host/blob/47b965aec6c78a23ab04a38c746df13000d802b1/u2flib_host/hid_transport.py
 CMD_INIT = 0x06
 CMD_WINK = 0x08
-CMD_APDU = 0x03
+
 U2FHID_YUBIKEY_DEVICE_CONFIG = U2F_VENDOR_FIRST
 
+# Cap flags -- https://github.com/Yubico/libu2f-host/blob/master/u2f-host/inc/u2f_hid.h#L105
+CAPFLAG_WINK = 0x01 # Device supports WINK command
+CAPFLAG_LOCK = 0x02 # Device supports LOCK command
+
 STAT_ERR = 0xbf
 
 def list_devices():
     devices = []
     for d in hid.enumerate(0, 0):
-        usage_page = d['usage_page']
-        if usage_page == 0xf1d0 and d['usage'] == 1:
+        if d['usage_page'] == 0xf1d0 and d['usage'] == 1:
             devices.append(HIDDevice(d['path']))
         # Usage page doesn't work on Linux
         elif (d['vendor_id'], d['product_id']) in DEVICES:
@@ -79,6 +85,8 @@ def list_devices():
                 devices.append(HIDDevice(d['path']))
             except exc.DeviceError:
                 pass
+            except OSError: # OSError: open failed
+                pass
     return devices
 
 
@@ -124,7 +132,16 @@ def init(self):
         while resp[:8] != nonce:
             print("Wrong nonce, read again...")
             resp = self._read_resp(self.cid, CMD_INIT)
+        self.initRes = resp
         self.cid = resp[8:12]
+        self.versionInterface = byte2int(resp[12])
+        self.versionMajor = byte2int(resp[13])
+        self.versionMinor = byte2int(resp[14])
+        self.versionBuild = byte2int(resp[15])
+        self.version = '{}.{}.{}'.format(self.versionMajor, self.versionMinor, self.versionBuild)
+        self.capFlags = byte2int(resp[16])
+        self.capWink = CAPFLAG_WINK & self.capFlags == CAPFLAG_WINK
+        self.capLock = CAPFLAG_LOCK & self.capFlags == CAPFLAG_LOCK
 
     def set_mode(self, mode):
         data = mode + b"\x0f\x00\x00"
@@ -136,6 +153,15 @@ def _do_send_apdu(self, apdu_data):
     def wink(self):
         self.call(CMD_WINK)
 
+    def ping(self, msg=b'Hello U2F'):
+        resp = self.call(CMD_PING, msg)
+        if resp != msg:
+            raise exc.DeviceError("Incorrect PING readback")
+        return resp
+
+    def lock(self, lock_time=10):
+        self.call(CMD_LOCK, lock_time)
+
     def _send_req(self, cid, cmd, data):
         size = len(data)
         bc_l = int2byte(size & 0xff)
@@ -160,7 +186,7 @@ def _read_resp(self, cid, cmd):
             resp_vals = _read_timeout(self.handle, HID_RPT_SIZE)
             resp = b''.join(int2byte(v) for v in resp_vals)
             if resp[:5] == cid + int2byte(STAT_ERR):
-                raise U2FHIDError(byte2int(resp[6]))
+                raise U2FHIDError(byte2int(resp[7]))
 
         if not resp:
             raise exc.DeviceError("Invalid response from device!")
@@ -175,7 +201,7 @@ def _read_resp(self, cid, cmd):
             resp = b''.join(int2byte(v) for v in resp_vals)
             if resp[:4] != cid:
                 raise exc.DeviceError("Wrong CID from device!")
-            if byte2int(resp[4:5]) != seq & 0x7f:
+            if byte2int(resp[4]) != seq & 0x7f:
                 raise exc.DeviceError("Wrong SEQ from device!")
             seq += 1
             new_data = resp[5:min(5 + data_len, HID_RPT_SIZE)]