storing capability byte in the HIDDevice object

Author: Venkat Venkataraju

.travis.yml

@@ -22,7 +22,7 @@ addons:
 
 install:
   - pip install "pip>=7.0.2" wheel
-  - pip install -r dev-requirements.txt
+  - $TRAVIS_BUILD_DIR/dev-requirements.py
   - pip install -e .
 
 script:

dev-requirements-2.7.txt

@@ -0,0 +1,5 @@
+coverage
+coveralls
+Cython # For building hidapi in Tox and on Travis
+cryptography>=1.0
+mock

dev-requirements.txt dev-requirements-3.x.txtdev-requirements.txt renamed to dev-requirements-3.x.txt

@@ -0,0 +1,11 @@
+#!/usr/bin/env python
+
+import os
+import sys
+
+if sys.version_info.major == 2:
+    os.system('pip install -r dev-requirements-2.7.txt')
+elif sys.version_info.major == 3:
+    os.system('pip install -r dev-requirements-3.x.txt')
+else:
+    raise Exception('Unsupported python version!')

dev-requirements.py

@@ -0,0 +1,85 @@
+import os
+import unittest
+from u2flib_host import hid_transport
+from u2flib_host import exc
+from u2flib_host.yubicommon.compat import byte2int, int2byte
+from mock import patch
+
+class TestHIDDevice():
+    def write(self, payload):
+        self.cid = payload[1:5]
+        self.cmd = payload[5] ^ hid_transport.TYPE_INIT
+        self.size = (payload[6] << 8) + payload[7]
+        self.data = list(map(int2byte, payload[8:(8 + self.size)]))
+
+    def read(self, size):
+        self.response += [0] * (hid_transport.HID_RPT_SIZE - len(self.response) + 1)
+        types = list(map(type, self.response))
+        return self.response
+
+    def close(self):
+        return None
+
+
+class HIDDeviceTest(unittest.TestCase):
+    @classmethod
+    def build_response(cls, cid, cmd, data):
+        size = len(data)
+        size_low = size & 0xff
+        size_high = (size >> 8) & 0xff
+        response = list(map(byte2int, cid)) + list(map(byte2int, cmd)) + [size_high, size_low]
+        response += list(map(byte2int, data))
+        return response
+
+
+    def test_init(self):
+        with patch.object(os, 'urandom', return_value=(b'\xab'*8)) as mock_method:
+            hid_device = TestHIDDevice()
+            hid_device.response = HIDDeviceTest.build_response(
+                b'\xff'*4,
+                b'\x86',
+                b'\xab'*8 + b'\x01\x02\x03\x04' + b'\x01\x02\x03\x04\x05'
+            )
+
+            dev = hid_transport.HIDDevice('/dev/null')
+            dev.handle = hid_device
+            dev.init()
+            self.assertEqual(dev.capabilities, 0x05)
+
+    def test_init_invalid_nounce(self):
+        with patch.object(os, 'urandom', return_value=(b'\xab'*8)) as mock_method:
+            hid_device = TestHIDDevice()
+            hid_device.response = HIDDeviceTest.build_response(
+                b'\xff'*4,
+                b'\x86',
+                b'\x00'*8 + b'\x01\x02\x03\x04' + b'\x01\x02\x03\x04\x05'
+            )
+
+            dev = hid_transport.HIDDevice('/dev/null')
+            dev.handle = hid_device
+            with self.assertRaises(exc.DeviceError) as context:
+                dev.init()
+                self.assertTrue('Wrong INIT response from device' in context.exception)
+
+    def test_init_invalid_length(self):
+        with patch.object(os, 'urandom', return_value=(b'\xab'*8)) as mock_method:
+            hid_device = TestHIDDevice()
+            hid_device.response = HIDDeviceTest.build_response(
+                b'\xff'*4,
+                b'\x86',
+                b'\xab'*8 + b'\x01\x02\x03\x04' + b'\x01\x02\x03\x04'
+            )
+
+            dev = hid_transport.HIDDevice('/dev/null')
+            dev.handle = hid_device
+
+            with self.assertRaises(exc.DeviceError) as context:
+                dev.init()
+                self.assertTrue('Wrong INIT response from device' in context.exception)
+
+    def test_ctap2_enabled(self):
+        dev = hid_transport.HIDDevice('/dev/null')
+        dev.capabilities = 0x01
+        self.assertFalse(dev.ctap2_enabled())
+        dev.capabilities = 0x04
+        self.assertTrue(dev.ctap2_enabled())

test/test_hid_transport.py

@@ -32,7 +32,7 @@
     import hidraw as hid  # Prefer hidraw
 except ImportError:
     import hid
-from time import time
+from time import time, sleep
 from u2flib_host.device import U2FDevice
 from u2flib_host.yubicommon.compat import byte2int, int2byte
 from u2flib_host import exc
@@ -110,6 +110,7 @@ class HIDDevice(U2FDevice):
     def __init__(self, path):
         self.path = path
         self.cid = b"\xff\xff\xff\xff"
+        self.capabilities = 0x00
 
     def open(self):
         self.handle = hid.device()
@@ -125,10 +126,19 @@ def close(self):
     def init(self):
         nonce = os.urandom(8)
         resp = self.call(CMD_INIT, nonce)
-        while resp[:8] != nonce:
-            print("Wrong nonce, read again...")
+
+        timeout = time() + 2.0
+        while (len(resp) != 17 or resp[:8] != nonce):
+            if timeout < time():
+                raise exc.DeviceError('Wrong INIT response from device')
+            sleep(0.1)
             resp = self._read_resp(self.cid, CMD_INIT)
+
         self.cid = resp[8:12]
+        self.capabilities = byte2int(resp[16])
+
+    def ctap2_enabled(self):
+        return (self.capabilities >> 2) & 0x01
 
     def set_mode(self, mode):
         data = mode + b"\x0f\x00\x00"