2 files changed, 198 insertions, 22 deletions

diff --git a/freestyle_hid/_freestyle_encryption.py b/freestyle_hid/_freestyle_encryption.py
new file mode 100644
index 0000000..eb04303
--- /dev/null
+++ b/freestyle_hid/_freestyle_encryption.py
@@ -0,0 +1,115 @@
+class SpeckEncrypt():
+
+    def __init__(self, key):
+        # Perform key expansion and store the round keys
+        self.key = key & ((2 ** 128) - 1)
+        self.key_schedule = [self.key & 0xFFFFFFFF]
+        key_buf = [(self.key >> (x * 32)) & 0xFFFFFFFF for x in range(1, 4)]
+        for x in range(26):
+            k = self.encryption_round(key_buf[x], self.key_schedule[x], x)
+            key_buf.append(k[0])
+            self.key_schedule.append(k[1])
+
+    def encryption_round(self, x, y, k):
+        # Perform one encryption round of the speck cipher
+        x_shift = ((x << 24) + (x >> 8)) & 0xFFFFFFFF
+        x_enc = k ^ ((x_shift + y) & 0xFFFFFFFF)
+        y_shift = ((y >> 29) + (y << 3)) & 0xFFFFFFFF
+        y_enc = x_enc ^ y_shift
+
+        return x_enc, y_enc
+
+    def decryption_round(self, x, y, k):
+        # Perform one decryption round of the speck cipher
+        new_y = (((x ^ y) << 29) + ((x ^ y) >> 3)) & 0xFFFFFFFF
+        msub = (((x ^ k) - new_y) + 0x100000000) % 0x100000000
+        new_x = ((msub >> 24) + (msub << 8)) & 0xFFFFFFFF
+
+        return new_x, new_y
+
+    def encrypt_block(self, plain):
+        # Encrypt one 64 bit block
+        x = (plain >> 32) & 0xFFFFFFFF
+        y = plain & 0xFFFFFFFF
+
+        for k in self.key_schedule:
+            x, y = self.encryption_round(x, y, k)
+
+        encrypted = (x << 32) + y
+
+        return encrypted
+
+    def decrypt_block(self, encrypted):
+        # Decrypt one 64 bit block
+        x = (encrypted >> 32) & 0xFFFFFFFF
+        y = encrypted & 0xFFFFFFFF
+
+        for k in reversed(self.key_schedule):
+            x, y = self.decryption_round(x, y, k)
+
+        plain = (x << 32) + y
+
+        return plain
+
+    def encrypt(self, iv, plain):
+        plain = bytearray(plain)
+        input_length = len(plain)
+        plain.extend(bytes(b'\x00' * (8 - (input_length % 8))))
+        iv = int.from_bytes(iv.to_bytes(8, byteorder='big'), byteorder='little', signed=False)
+        output = bytearray()
+        for i in range(len(plain) // 8):
+            k = self.encrypt_block(iv)
+            res = k ^ int.from_bytes(plain[i*8:i*8+8], byteorder='little', signed=False)
+            output.extend(int.to_bytes(res, 8, byteorder='little', signed=False))
+            iv += 1
+        encrypted = output[:input_length]
+        return bytes(encrypted)
+
+    def decrypt(self, iv, encrypted):
+        return self.encrypt(iv, encrypted)
+
+class SpeckCMAC:
+
+    def __init__(self, key):
+        self.cipher = SpeckEncrypt(key)
+
+        k0 = self.cipher.encrypt_block(0)
+        k0 = int.from_bytes(k0.to_bytes(8, byteorder='big'), byteorder='little', signed=False)
+
+        k1 = (k0 << 1) & 0XFFFFFFFFFFFFFFFF
+        if (k0 >> 63 != 0):
+            k1 ^= 0x1B
+
+        k2 = (k1 << 1) & 0XFFFFFFFFFFFFFFFF
+        if (k1 >> 63 != 0):
+            k2 ^= 0x1B
+
+        k1 = int.from_bytes(k1.to_bytes(8, byteorder='big'), byteorder='little', signed=False)
+        k2 = int.from_bytes(k2.to_bytes(8, byteorder='big'), byteorder='little', signed=False)
+        self.k1 = k1
+        self.k2 = k2
+
+    def sign(self, data):
+        c = 0
+        i = 0
+        data_len = len(data)
+
+        while (i < data_len):
+            data_left = data_len - i
+            if (data_left == 8):
+                block = int.from_bytes(data[i:i+8], 'little') ^ self.k1
+            elif (data_left < 8):
+                block = int.from_bytes(data[i:i+data_left] + b'\x80' + b'\x00'*(7-data_left), 'little') ^ self.k2
+            else:
+                block = int.from_bytes(data[i:i+8], 'little')
+            c = self.cipher.encrypt_block(c ^ block)
+            i += 8
+
+        return c
+
+    def derive(self, label, context):
+        data = label + b'\x00' + context + b'\x80\x00'
+        d1 = self.sign(b'\x01' + data)
+        d2 = self.sign(b'\x02' + data) << 64
+
+        return d1 | d2
diff --git a/freestyle_hid/_session.py b/freestyle_hid/_session.py
index 527a5fc..581f3fe 100644
--- a/freestyle_hid/_session.py
+++ b/freestyle_hid/_session.py
@@ -4,6 +4,7 @@
 import csv
 import logging
 import pathlib
+import random
 import re
 from typing import AnyStr, Callable, Iterator, Optional, Sequence, Tuple
 
@@ -11,9 +12,15 @@ import construct
 
 from ._exceptions import ChecksumError, CommandError
 from ._hidwrapper import HidWrapper
+from ._freestyle_encryption import SpeckEncrypt, SpeckCMAC
 
 ABBOTT_VENDOR_ID = 0x1A61
 
+_AUTH_ENC_MASTER_KEY = 0xdeadbeef
+_AUTH_MAC_MASTER_KEY = 0xdeadbeef
+_SESS_ENC_MASTER_KEY = 0xdeadbeef
+_SESS_MAC_MASTER_KEY = 0xdeadbeef
+
 _INIT_COMMAND = 0x01
 _INIT_RESPONSE = 0x71
 
@@ -64,15 +71,6 @@ _FREESTYLE_MESSAGE = construct.Struct(
     ),
 )
 
-_FREESTYLE_ENCRYPTED_MESSAGE = construct.Struct(
-    hid_report=construct.Const(0, construct.Byte),
-    message_type=construct.Byte,
-    command=construct.Padded(
-        63,  # command can only be up to 62 bytes, but one is used for length.
-        construct.GreedyBytes,
-    ),
-)
-
 _TEXT_COMPLETION_RE = re.compile(b"CMD (?:OK|Fail!)")
 _TEXT_REPLY_FORMAT = re.compile(
     b"^(?P<message>.*)CKSM:(?P<checksum>[0-9A-F]{8})\r\n"
@@ -121,12 +119,57 @@ class Session:
         encoding: str = "ascii",
     ) -> None:
         self._handle = HidWrapper.open(device_path, ABBOTT_VENDOR_ID, product_id)
-
         self._text_message_type = text_message_type
         self._text_reply_message_type = text_reply_message_type
         self._encoding = encoding
+        self._encrypted_protocol = product_id in [0x3950]
+
+    def encryption_handshake(self):
+        self.send_command(0x05, b"")
+        response = self.read_response()
+        assert response[0] == 0x06
+        serial = response[1][:13]
+
+        crypt = SpeckCMAC(_AUTH_ENC_MASTER_KEY)
+        auth_enc_key = crypt.derive("AuthrEnc".encode(), serial)
+        auth_enc = SpeckEncrypt(auth_enc_key)
+        crypt = SpeckCMAC(_AUTH_MAC_MASTER_KEY)
+        auth_mac_key = crypt.derive("AuthrMAC".encode(), serial)
+        auth_mac = SpeckCMAC(auth_mac_key)
+
+        self.send_command(_ENCRYPTION_SETUP_COMMAND, b"\x11")
+        response = self.read_response()
+        assert response[0] == _ENCRYPTION_SETUP_RESPONSE
+        assert response[1][0] == 0x16
+        reader_rand = response[1][1:9]
+        iv = int.from_bytes(response[1][9:16], 'big', signed=False)
+        driver_rand = random.randbytes(8)
+        resp_enc = auth_enc.encrypt(iv, reader_rand + driver_rand)
+        resp_mac = auth_mac.sign(b"\x14\x1a\x17" + resp_enc + b"\x01")
+        resp_mac = int.to_bytes(resp_mac, 8, byteorder='little', signed=False)
+        self.send_command(_ENCRYPTION_SETUP_COMMAND, b"\x17" + resp_enc + b"\x01" + resp_mac)
+        response = self.read_response()
+        assert response[0] == _ENCRYPTION_SETUP_RESPONSE
+        assert response[1][0] == 0x18
+        mac = auth_mac.sign(b"\x33\x22" + response[1][:24])
+        mac = int.to_bytes(mac, 8, byteorder='little', signed=False)
+        assert mac == response[1][24:32]
+        iv = int.from_bytes(response[1][17:24], 'big', signed=False)
+        resp_dec = auth_enc.decrypt(iv, response[1][1:17])
+        assert resp_dec[:8] == driver_rand
+        assert resp_dec[8:] == reader_rand
+
+        crypt = SpeckCMAC(_SESS_ENC_MASTER_KEY)
+        ses_enc_key = crypt.derive("SessnEnc".encode(), serial + reader_rand + driver_rand)
+        crypt = SpeckCMAC(_SESS_MAC_MASTER_KEY)
+        ses_mac_key = crypt.derive("SessnMAC".encode(), serial + reader_rand + driver_rand)
+        self.crypt_enc = SpeckEncrypt(ses_enc_key)
+        self.crypt_mac = SpeckCMAC(ses_mac_key)
+        #print("HANDSHAKE SUCCESSFUL!")
 
     def connect(self):
+        if self._encrypted_protocol:
+            self.encryption_handshake()
         """Open connection to the device, starting the knocking sequence."""
         self.send_command(_INIT_COMMAND, b"")
         response = self.read_response()
@@ -135,6 +178,26 @@ class Session:
                 f"Connection error: unexpected message %{response[0]:02x}:{response[1].hex()}"
             )
 
+    def encrypt_message(self, packet: bytes):
+        output = bytearray(packet)
+        # 0xFF IV is actually 0, because of some weird padding
+        encrypted = self.crypt_enc.encrypt(0xFF, packet[2:57])
+        output[2:57] = encrypted
+        # Not giving a f**k about the IV counter for now
+        output[57:61] = bytes(4)
+        mac = self.crypt_mac.sign(output[1:61])
+        output[61:65] = int.to_bytes(mac, 8, byteorder='little', signed=False)[4:]
+        return bytes(output)
+
+    def decrypt_message(self, packet: bytes):
+        output = bytearray(packet)
+        mac = self.crypt_mac.sign(packet[:60])
+        mac = int.to_bytes(mac, 8, byteorder='little', signed=False)[4:]
+        assert mac == packet[60:64]
+        iv = int.from_bytes(packet[56:60], 'big', signed=False) << 8
+        output[1:56] = self.crypt_enc.decrypt(iv, packet[1:56])
+        return bytes(output)
+
     def send_command(self, message_type: int, command: bytes, encrypted: bool = False):
         """Send a raw command to the device.
 
@@ -142,16 +205,14 @@ class Session:
           message_type: The first byte sent with the report to the device.
           command: The command to send out the device.
         """
-        if encrypted:
-            assert message_type not in _ALWAYS_UNENCRYPTED_MESSAGES
-            meta_construct = _FREESTYLE_ENCRYPTED_MESSAGE
-        else:
-            meta_construct = _FREESTYLE_MESSAGE
 
-        usb_packet = meta_construct.build(
+        usb_packet = _FREESTYLE_MESSAGE.build(
             {"message_type": message_type, "command": command}
         )
 
+        if self._encrypted_protocol and message_type not in _ALWAYS_UNENCRYPTED_MESSAGES:
+            usb_packet = self.encrypt_message(usb_packet)
+
         logging.debug(f"Sending packet: {usb_packet!r}")
         self._handle.write(usb_packet)
 
@@ -164,12 +225,12 @@ class Session:
         assert usb_packet
         message_type = usb_packet[0]
 
-        if not encrypted or message_type in _ALWAYS_UNENCRYPTED_MESSAGES:
-            message_length = usb_packet[1]
-            message_end_idx = 2 + message_length
-            message_content = usb_packet[2:message_end_idx]
-        else:
-            message_content = usb_packet[1:]
+        if self._encrypted_protocol and message_type not in _ALWAYS_UNENCRYPTED_MESSAGES:
+            usb_packet = self.decrypt_message(usb_packet)
+
+        message_length = usb_packet[1]
+        message_end_idx = 2 + message_length
+        message_content = usb_packet[2:message_end_idx]
 
         # hidapi module returns a list of bytes rather than a bytes object.
         message = (message_type, bytes(message_content))