From 2c77347d56edf36c2e8d16338c7e9e348fc44c24 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Diego=20Elio=20Petten=C3=B2?= <flameeyes@flameeyes.com>
Date: Wed, 2 Aug 2023 23:05:35 +0100
Subject: If present, load the actual keys from freestyle-keys package.

---
 freestyle_hid/_exceptions.py | 11 +++++++++++
 freestyle_hid/_session.py    | 25 +++++++++++++++----------
 setup.cfg                    |  2 ++
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/freestyle_hid/_exceptions.py b/freestyle_hid/_exceptions.py
index 38a822d..2b803f8 100644
--- a/freestyle_hid/_exceptions.py
+++ b/freestyle_hid/_exceptions.py
@@ -16,3 +16,14 @@ class ChecksumError(Exception):
 
 class CommandError(Exception):
     """Errors related to the command stream."""
+
+
+class MissingFreeStyleKeys(Exception):
+    """The freestyle-hid-keys package is missing."""
+
+    def __init__(self):
+        super().__init__(
+            "The freestyle-hid-keys package is missing, please install it from PyPi."
+            " You can install freestyle-hid[encryption] to select the encryption keys"
+            " package as an extra dependency."
+        )
diff --git a/freestyle_hid/_session.py b/freestyle_hid/_session.py
index cdfca42..0debaed 100644
--- a/freestyle_hid/_session.py
+++ b/freestyle_hid/_session.py
@@ -10,16 +10,18 @@ from typing import AnyStr, Callable, Iterator, Optional, Sequence, Tuple
 
 import construct
 
-from ._exceptions import ChecksumError, CommandError
+from ._exceptions import ChecksumError, CommandError, MissingFreeStyleKeys
 from ._freestyle_encryption import SpeckCMAC, SpeckEncrypt
 from ._hidwrapper import HidWrapper
 
-ABBOTT_VENDOR_ID = 0x1A61
+try:
+    from freestyle_keys import libre2 as libre2_keys
+
+    _HAS_LIBRE2_KEYS = True
+except ImportError:
+    _HAS_LIBRE2_KEYS = False
 
-_AUTH_ENC_MASTER_KEY = 0xDEADBEEF
-_AUTH_MAC_MASTER_KEY = 0xDEADBEEF
-_SESS_ENC_MASTER_KEY = 0xDEADBEEF
-_SESS_MAC_MASTER_KEY = 0xDEADBEEF
+ABBOTT_VENDOR_ID = 0x1A61
 
 _INIT_COMMAND = 0x01
 _INIT_RESPONSE = 0x71
@@ -125,15 +127,18 @@ class Session:
         self._encrypted_protocol = product_id in [0x3950]
 
     def encryption_handshake(self):
+        if not _HAS_LIBRE2_KEYS:
+            raise MissingFreeStyleKeys()
+
         self.send_command(0x05, b"")
         response = self.read_response()
         assert response[0] == 0x06
         serial = response[1][:13]
 
-        crypt = SpeckCMAC(_AUTH_ENC_MASTER_KEY)
+        crypt = SpeckCMAC(libre2_keys.AUTHORIZATION_ENCRYPTION_KEY)
         auth_enc_key = crypt.derive("AuthrEnc".encode(), serial)
         auth_enc = SpeckEncrypt(auth_enc_key)
-        crypt = SpeckCMAC(_AUTH_MAC_MASTER_KEY)
+        crypt = SpeckCMAC(libre2_keys.AUTHORIZATION_MAC_KEY)
         auth_mac_key = crypt.derive("AuthrMAC".encode(), serial)
         auth_mac = SpeckCMAC(auth_mac_key)
 
@@ -161,11 +166,11 @@ class Session:
         assert resp_dec[:8] == driver_rand
         assert resp_dec[8:] == reader_rand
 
-        crypt = SpeckCMAC(_SESS_ENC_MASTER_KEY)
+        crypt = SpeckCMAC(libre2_keys.SESSION_ENCRYPTION_KEY)
         ses_enc_key = crypt.derive(
             "SessnEnc".encode(), serial + reader_rand + driver_rand
         )
-        crypt = SpeckCMAC(_SESS_MAC_MASTER_KEY)
+        crypt = SpeckCMAC(libre2_keys.SESSION_MAC_KEY)
         ses_mac_key = crypt.derive(
             "SessnMAC".encode(), serial + reader_rand + driver_rand
         )
diff --git a/setup.cfg b/setup.cfg
index 9b69d56..d9d0089 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -40,6 +40,8 @@ python_requires = ~= 3.7
 [options.extras_require]
 hidapi =
     hidapi
+encryption =
+    freestyle-keys
 tools =
     click
     click_log
-- 
cgit v1.2.3