Add support for SET_CODE and VALIDATE

pyauth · May 1, 2019 · 91f3744 · 91f3744
1 parent fa8815d
commit 91f3744
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 5 deletions.
diff --git a/exile/ykoath/__init__.py b/exile/ykoath/__init__.py
@@ -1,4 +1,4 @@
-import base64, struct, typing
+import base64, struct, typing, hashlib, hmac
 from collections import namedtuple
 from datetime import datetime
 from urllib.parse import urlparse, parse_qs
@@ -12,16 +12,21 @@ class YKOATH(YKOATHConstants):
     """
     See https://developers.yubico.com/OATH/YKOATH_Protocol.html
     """
-    def __init__(self, device: SCardReader = None) -> None:
+    def __init__(self, device: SCardReader = None, password: str = None) -> None:
         if device is None:
             for reader in SCardManager():
                 if reader.name.startswith(self.device_prefix):
                     device = reader
                     break
             else:
-                raise YKOATHError("No Yubikey found")
+                raise YKOATHError("No YubiKey found")
         self.device = device
-        self.send_apdu(cla=0, ins=self.Instruction.SELECT, p1=0x04, p2=0, data=self.Application.OATH)
+        res = self.send_apdu(cla=0, ins=self.Instruction.SELECT, p1=0x04, p2=0, data=self.Application.OATH)
+        _, self._version, res = self.parse_tlv(res, self.Tag.VERSION)
+        _, self._id, res = self.parse_tlv(res, self.Tag.NAME)
+        _, self._challenge, res = self.parse_tlv(res)
+        if self._challenge and password is not None:
+            self.validate(password)
 
     def send_apdu(self, **kwargs):
         with self.device:
@@ -32,6 +37,14 @@ def send_apdu(self, **kwargs):
             raise YKOATHError(self.Response(res[-2:]))
         return res
 
+    def parse_tlv(self, data, expect_tag=None):
+        assert isinstance(data, bytes)
+        tag, length = data[0], data[1]
+        if expect_tag:
+            assert tag == expect_tag
+        value, data = data[2:2 + length], data[2 + length:]
+        return tag, value, data
+
     def put(self, credential_name: str, secret: bytes, require_touch=False,
             oath_type=YKOATHConstants.OATHType.TOTP, algorithm=YKOATHConstants.Algorithm.SHA1, digits=6):
         secret_header = i2b(oath_type.value | algorithm.value) + i2b(digits)
@@ -66,6 +79,22 @@ def calculate(self, credential_name: str, challenge: typing.Union[bytes, int], w
         else:
             return res[3:3 + res_len - 1]
 
+    def set_code(self, password):
+        key = hashlib.pbkdf2_hmac('sha256', password.encode(), self._id, 1000)
+        test_challenge = b'01234567'
+        test_response = hmac.new(key, test_challenge, 'sha256').digest()
+        data = i2b(self.Tag.KEY) + i2b(len(key) + 1) + i2b(self.Algorithm.SHA256.value) + key
+        data += i2b(self.Tag.CHALLENGE) + i2b(len(test_challenge)) + test_challenge
+        data += i2b(self.Tag.RESPONSE) + i2b(len(test_response)) + test_response
+        return self.send_apdu(cla=0, ins=self.Instruction.SET_CODE, p1=0, p2=0, data=data)
+
+    def validate(self, password):
+        key = hashlib.pbkdf2_hmac('sha256', password.encode(), self._id, 1000)
+        response = hmac.new(key, self._challenge, 'sha256').digest()
+        data = i2b(self.Tag.RESPONSE) + i2b(len(response)) + response
+        data += i2b(self.Tag.CHALLENGE) + i2b(len(self._challenge)) + self._challenge
+        return self.send_apdu(cla=0, ins=self.Instruction.VALIDATE, p1=0, p2=0, data=data)
+
     def __iter__(self):
         res = self.list()[:-2]
         while res:

diff --git a/exile/ykoath/const.py b/exile/ykoath/const.py
@@ -54,7 +54,7 @@ class Response(Enum):
         SUCCESS = b'\x90\x00'
         NO_SPACE = b'\x6a\x84'
         NOT_FOUND = b'\x69\x84'
-        AUTH_REQUIRED = b'x69\x82'
+        AUTH_REQUIRED = b'\x69\x82'
         WRONG_SYNTAX = b'\x6a\x80'
         GENERIC_ERROR = b'\x65\x81'
         MORE_DATA_AVAILABLE = b'\x61'