Add experimental key derivation API for real airtags

examples/real_airtag.py

@@ -0,0 +1,48 @@
+"""
+Example showing how to retrieve the primary key of your own AirTag
+(or any other FindMy-accessory).
+This key can be used to retrieve the device's location for a single day.
+"""
+
+from findmy import FindMyAccessory
+
+# PUBLIC key that the accessory is broadcasting or has previously broadcast.
+# For nearby devices, you can use `device_scanner.py` to find it.
+LOOKUP_KEY = "9J5sdEARfh6h0Hr3anfNjy+vnIwETaUodv73ZA=="
+
+# PRIVATE master key. 28 (?) bytes.
+MASTER_KEY = b""
+
+# "Primary" shared secret. 32 bytes.
+SKN = b""
+
+# "Secondary" shared secret. 32 bytes.
+SKS = b""
+
+# Lookahead in time slots. Each time slot is 15 minutes.
+# Should be AT LEAST the time that has passed since you paired the accessory!
+MAX_LOOKAHEAD = 7 * 24 * 4
+
+
+def main() -> None:
+    airtag = FindMyAccessory(MASTER_KEY, SKN, SKS)
+
+    for i in range(MAX_LOOKAHEAD):
+        prim_key, sec_key = airtag.keys_at(i)
+        if LOOKUP_KEY == prim_key.adv_key_b64 or LOOKUP_KEY == prim_key.adv_key_b64:
+            print(f"KEY FOUND!!")
+            print(f"This key was found at index {i}."
+                  f" It was likely paired approximately {i * 15} minutes ago")
+            print()
+            print("KEEP THE BELOW KEY SECRET! IT CAN BE USED TO RETRIEVE THE DEVICE'S LOCATION!")
+            if LOOKUP_KEY == prim_key.adv_key_b64:
+                print(f"PRIMARY key: {prim_key.private_key_b64}")
+            else:
+                print(f"SECONDARY key: {sec_key.private_key_b64}")
+    else:
+        print("No match found! :(")
+        return
+
+
+if __name__ == '__main__':
+    main()

findmy/__init__.py

@@ -1,5 +1,14 @@
 """A package providing everything you need to work with Apple's FindMy network."""
-from . import keys, reports
+from . import keys, reports, scanner
+from .accessory import FindMyAccessory
+from .keys import KeyPair
 from .util import errors
 
-__all__ = ("reports", "keys", "errors")
+__all__ = (
+    "keys",
+    "reports",
+    "scanner",
+    "errors",
+    "FindMyAccessory",
+    "KeyPair",
+)

findmy/accessory.py

@@ -0,0 +1,110 @@
+"""
+Module to interact with accessories that implement Find My.
+
+Accessories could be anything ranging from AirTags to iPhones.
+"""
+from __future__ import annotations
+
+from typing import Generator
+
+from .keys import KeyGenerator, KeyPair
+from .util import crypto
+
+
+class FindMyAccessory:
+    """A findable Find My-accessory using official key rollover."""
+
+    def __init__(self, master_key: bytes, skn: bytes, sks: bytes, name: str | None = None) -> None:
+        """
+        Initialize a FindMyAccessory. These values are usually obtained during pairing.
+
+        :param master_key: The private master key.
+        :param skn: The SKN for the primary key.
+        :param sks: The SKS for the secondary key.
+        """
+        self._primary_gen = AccessoryKeyGenerator(master_key, skn)
+        self._secondary_gen = AccessoryKeyGenerator(master_key, sks)
+
+        self._name = name
+
+    def keys_at(self, ind: int) -> tuple[KeyPair, KeyPair]:
+        """Get the primary and secondary key active at primary key index `ind`."""
+        pkey = self._primary_gen[ind]
+        skey = self._secondary_gen[ind // 96 + 1]
+
+        return pkey, skey
+
+
+class AccessoryKeyGenerator(KeyGenerator[KeyPair]):
+    """KeyPair generator. Uses the same algorithm internally as FindMy accessories do."""
+
+    def __init__(self, master_key: bytes, initial_sk: bytes) -> None:
+        """
+        Initialize the key generator.
+
+        :param master_key: Private master key. Usually obtained during pairing.
+        :param initial_sk: Initial secret key. Can be the SKN to generate primary keys,
+                           or the SKS to generate secondary ones.
+        """
+        if len(master_key) != 28:
+            msg = "The master key must be 28 bytes long"
+            raise ValueError(msg)
+        if len(initial_sk) != 32:
+            msg = "The sk must be 32 bytes long"
+            raise ValueError(msg)
+
+        self._master_key = master_key
+        self._initial_sk = initial_sk
+
+        self._cur_sk = initial_sk
+        self._cur_sk_ind = 0
+
+        self._iter_ind = 0
+
+    def _get_sk(self, ind: int) -> bytes:
+        if ind < self._cur_sk_ind:  # behind us; need to reset :(
+            self._cur_sk = self._initial_sk
+            self._cur_sk_ind = 0
+
+        for _ in range(self._cur_sk_ind, ind):
+            self._cur_sk = crypto.x963_kdf(self._cur_sk, b"update", 32)
+            self._cur_sk_ind += 1
+        return self._cur_sk
+
+    def _get_keypair(self, ind: int) -> KeyPair:
+        sk = self._get_sk(ind)
+        privkey = crypto.derive_ps_key(self._master_key, sk)
+        return KeyPair(privkey)
+
+    def _generate_keys(self, start: int, stop: int | None) -> Generator[KeyPair, None, None]:
+        ind = start
+        while stop is None or ind < stop:
+            yield self._get_keypair(ind)
+
+            ind += 1
+
+    def __iter__(self) -> KeyGenerator:
+        self._iter_ind = -1
+        return self
+
+    def __next__(self) -> KeyPair:
+        self._iter_ind += 1
+
+        return self._get_keypair(self._iter_ind)
+
+    def __getitem__(self, val: int | slice) -> KeyPair | Generator[KeyPair, None, None]:
+        if isinstance(val, int):
+            if val < 0:
+                msg = "The key index must be non-negative"
+                raise ValueError(msg)
+
+            return self._get_keypair(val)
+        if isinstance(val, slice):
+            start, stop = val.start or 0, val.stop
+            if start < 0 or (stop is not None and stop < 0):
+                msg = "The key index must be non-negative"
+                raise ValueError(msg)
+
+            return self._generate_keys(start, stop)
+
+        return NotImplemented

findmy/keys.py

@@ -1,13 +1,16 @@
 """Module to work with private and public keys as used in FindMy accessories."""
+from __future__ import annotations
 
 import base64
 import hashlib
 import secrets
 from abc import ABC, abstractmethod
+from typing import Generator, Generic, TypeVar, overload
 
-from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives.asymmetric import ec
 
+from .util import crypto
+
 
 class HasPublicKey(ABC):
     """
@@ -43,20 +46,19 @@ class KeyPair(HasPublicKey):
 
     def __init__(self, private_key: bytes) -> None:
         """Initialize the `KeyPair` with the private key bytes."""
-        priv_int = int.from_bytes(private_key, "big")
+        priv_int = crypto.bytes_to_int(private_key)
         self._priv_key = ec.derive_private_key(
             priv_int,
             ec.SECP224R1(),
-            default_backend(),
         )
 
     @classmethod
-    def generate(cls) -> "KeyPair":
+    def new(cls) -> KeyPair:
         """Generate a new random `KeyPair`."""
         return cls(secrets.token_bytes(28))
 
     @classmethod
-    def from_b64(cls, key_b64: str) -> "KeyPair":
+    def from_b64(cls, key_b64: str) -> KeyPair:
         """
         Import an existing `KeyPair` from its base64-encoded representation.
 
@@ -88,3 +90,33 @@ def adv_key_bytes(self) -> bytes:
     def dh_exchange(self, other_pub_key: ec.EllipticCurvePublicKey) -> bytes:
         """Do a Diffie-Hellman key exchange using another EC public key."""
         return self._priv_key.exchange(ec.ECDH(), other_pub_key)
+
+    def __repr__(self) -> str:
+        return f'KeyPair(public_key="{self.adv_key_b64}")'
+
+
+K = TypeVar("K")
+
+
+class KeyGenerator(ABC, Generic[K]):
+    """KeyPair generator."""
+
+    @abstractmethod
+    def __iter__(self) -> KeyGenerator:
+        return NotImplemented
+
+    @abstractmethod
+    def __next__(self) -> K:
+        return NotImplemented
+
+    @overload
+    def __getitem__(self, val: int) -> K:
+        ...
+
+    @overload
+    def __getitem__(self, slc: slice) -> Generator[K, None, None]:
+        ...
+
+    @abstractmethod
+    def __getitem__(self, val: int | slice) -> K | Generator[K, None, None]:
+        return NotImplemented

findmy/util/crypto.py

@@ -0,0 +1,106 @@
+"""Pure-python NIST P-224 Elliptic Curve cryptography. Used for some Apple algorithms."""
+
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives.hashes import SHA1
+from cryptography.hazmat.primitives.kdf.x963kdf import X963KDF
+
+ECPoint = tuple[float, float]
+
+P224_A = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFE
+P224_N = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFF16A2E0B8F03E13DD29455C5C2A3D
+P224_G = (
+    0xB70E0CBD6BB4BF7F321390B94A03C1D356C21122343280D6115C1D21,
+    0xBD376388B5F723FB4C22DFE6CD4375A05A07476444D5819985007E34,
+)
+P224_P = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF000000000000000000000001
+
+
+# Old code. Remove if replacement is confirmed to be working.
+# ruff: noqa: ERA001
+#
+# def _ec_add_points(p1: ECPoint, p2: ECPoint) -> ECPoint:
+#     """
+#     Add two points on a P-224 elliptic curve. (0, 0) is identity.
+#
+#     https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication#Point_addition
+#     """
+#     (x1, y1), (x2, y2) = p1, p2
+#
+#     if p1 == (0, 0):  # identity case 1
+#         return p2
+#     if p2 == (0, 0):  # identity case 2
+#         return p1
+#     if x1 == x2 and y1 == -1 * y2:  # additive inverse
+#         return 0, 0
+#
+#     if p1 == p2:  # point doubling using limit
+#         slope = (3 * x1 ** 2 + P224_A) / (2 * y1)
+#     else:
+#         slope = (y2 - y1) / (x2 - x1)
+#
+#     x = slope ** 2 - x1 - x2
+#     y = slope * (x1 - x) - y1
+#
+#     return x, y
+#
+#
+# def _ec_scalar_mul(scalar: int, p: ECPoint) -> ECPoint:
+#     """
+#     Scalar multiplication on a point on a P-224 elliptic curve.
+#
+#     https://en.wikipedia.org/wiki/Elliptic_curve_point_multiplication#Double-and-add
+#     """
+#     res = (0, 0)
+#     cur = p
+#     while scalar > 0:
+#         if scalar & 1:
+#             res = _ec_add_points(res, cur)
+#         cur = _ec_add_points(cur, cur)
+#         scalar >>= 1
+#     return res
+#
+#
+# def derive_ps_key(pubkey: ECPoint, sk: bytes) -> ECPoint:
+#     at = _x963_kdf(sk, b"diversify", 72)
+#     u = int.from_bytes(at[:36], "big") % (P224_N - 1) + 1
+#     v = int.from_bytes(at[36:], "big") % (P224_N - 1) + 1
+#
+#     return _ec_add_points(_ec_scalar_mul(u, pubkey), _ec_scalar_mul(v, P224_G))
+
+
+def x963_kdf(value: bytes, si: bytes, length: int) -> bytes:
+    """Single pass of X9.63 KDF with SHA1."""
+    return X963KDF(
+        algorithm=SHA1(),  # noqa: S303
+        sharedinfo=si,
+        length=length,
+    ).derive(value)
+
+
+def bytes_to_int(value: bytes) -> int:
+    """Convert bytes in big-endian format to int."""
+    return int.from_bytes(value, "big")
+
+
+def derive_ps_key(privkey: bytes, sk: bytes) -> bytes:
+    """
+    Derive a primary or secondary key used by an accessory.
+
+    :param pubkey: Public key generated during pairing
+    :param sk: Current secret key for this time period.
+               Use SKN to derive the primary key, SKS for secondary.
+    """
+    priv_int = bytes_to_int(privkey)
+
+    at = x963_kdf(sk, b"diversify", 72)
+    u = bytes_to_int(at[:36]) % (P224_N - 1) + 1
+    v = bytes_to_int(at[36:]) % (P224_N - 1) + 1
+
+    key = (u * priv_int + v) % P224_N
+    return key.to_bytes(28, "big")
+
+
+def _get_pubkey(privkey: int) -> ECPoint:
+    key = ec.derive_private_key(privkey, ec.SECP224R1())
+    pubkey = key.public_key().public_numbers()
+    return pubkey.x, pubkey.y