Merge pull request #148 from malmeloo/feat/better-serialization

Make more objects `Serializable`

Author: malmeloo

examples/_login.py

@@ -1,15 +1,13 @@
-# ruff: noqa: ASYNC230
+from __future__ import annotations
 
 from findmy.reports import (
     AppleAccount,
     AsyncAppleAccount,
-    BaseAnisetteProvider,
     LoginState,
     SmsSecondFactorMethod,
     TrustedDeviceSecondFactorMethod,
 )
-
-ACCOUNT_STORE = "account.json"
+from findmy.reports.anisette import LocalAnisetteProvider, RemoteAnisetteProvider
 
 
 def _login_sync(account: AppleAccount) -> None:
@@ -66,27 +64,45 @@ async def _login_async(account: AsyncAppleAccount) -> None:
         await method.submit(code)
 
 
-def get_account_sync(anisette: BaseAnisetteProvider) -> AppleAccount:
+def get_account_sync(
+    store_path: str,
+    anisette_url: str | None,
+    libs_path: str | None,
+) -> AppleAccount:
     """Tries to restore a saved Apple account, or prompts the user for login otherwise. (sync)"""
-    acc = AppleAccount(anisette=anisette)
-    acc_store = "account.json"
     try:
-        acc.from_json(acc_store)
+        acc = AppleAccount.from_json(store_path, anisette_libs_path=libs_path)
     except FileNotFoundError:
+        ani = (
+            LocalAnisetteProvider(libs_path=libs_path)
+            if anisette_url is None
+            else RemoteAnisetteProvider(anisette_url)
+        )
+        acc = AppleAccount(ani)
         _login_sync(acc)
-        acc.to_json(acc_store)
+
+        acc.to_json(store_path)
 
     return acc
 
 
-async def get_account_async(anisette: BaseAnisetteProvider) -> AsyncAppleAccount:
+async def get_account_async(
+    store_path: str,
+    anisette_url: str | None,
+    libs_path: str | None,
+) -> AsyncAppleAccount:
     """Tries to restore a saved Apple account, or prompts the user for login otherwise. (async)"""
-    acc = AsyncAppleAccount(anisette=anisette)
-    acc_store = "account.json"
     try:
-        acc.from_json(acc_store)
+        acc = AsyncAppleAccount.from_json(store_path, anisette_libs_path=libs_path)
     except FileNotFoundError:
+        ani = (
+            LocalAnisetteProvider(libs_path=libs_path)
+            if anisette_url is None
+            else RemoteAnisetteProvider(anisette_url)
+        )
+        acc = AsyncAppleAccount(ani)
         await _login_async(acc)
-        acc.to_json(acc_store)
+
+        acc.to_json(store_path)
 
     return acc

examples/fetch_reports.py

@@ -4,28 +4,47 @@
 from _login import get_account_sync
 
 from findmy import KeyPair
-from findmy.reports import RemoteAnisetteProvider
 
-# URL to (public or local) anisette server
-ANISETTE_SERVER = "http://localhost:6969"
+# Path where login session will be stored.
+# This is necessary to avoid generating a new session every time we log in.
+STORE_PATH = "account.json"
+
+# URL to LOCAL anisette server. Set to None to use built-in Anisette generator instead (recommended)
+# IF YOU USE A PUBLIC SERVER, DO NOT COMPLAIN THAT YOU KEEP RUNNING INTO AUTHENTICATION ERRORS!
+# If you change this value, make sure to remove the account store file.
+ANISETTE_SERVER = None
+
+# Path where Anisette libraries will be stored.
+# This is only relevant when using the built-in Anisette server.
+# It can be omitted (set to None) to avoid saving to disk,
+# but specifying a path is highly recommended to avoid downloading the bundle on every run.
+ANISETTE_LIBS_PATH = "ani_libs.bin"
 
 logging.basicConfig(level=logging.INFO)
 
 
 def fetch_reports(priv_key: str) -> int:
-    key = KeyPair.from_b64(priv_key)
-    acc = get_account_sync(
-        RemoteAnisetteProvider(ANISETTE_SERVER),
-    )
+    # Step 0: construct an account instance
+    # We use a helper for this to simplify interactive authentication
+    acc = get_account_sync(STORE_PATH, ANISETTE_SERVER, ANISETTE_LIBS_PATH)
 
     print(f"Logged in as: {acc.account_name} ({acc.first_name} {acc.last_name})")
 
-    # It's that simple!
+    # Step 1: construct a key object and get its location reports
+    key = KeyPair.from_b64(priv_key)
     reports = acc.fetch_last_reports(key)
+
+    # Step 2: print the reports!
     for report in sorted(reports):
         print(report)
 
-    return 1
+        # We can save the report to a file if we want
+        report.to_json("last_report.json")
+
+    # Step 3: Make sure to save account state when you're done!
+    acc.to_json(STORE_PATH)
+
+    return 0
 
 
 if __name__ == "__main__":

examples/fetch_reports_async.py

@@ -5,30 +5,49 @@
 from _login import get_account_async
 
 from findmy import KeyPair
-from findmy.reports import RemoteAnisetteProvider
 
-# URL to (public or local) anisette server
-ANISETTE_SERVER = "http://localhost:6969"
+# Path where login session will be stored.
+# This is necessary to avoid generating a new session every time we log in.
+STORE_PATH = "account.json"
+
+# URL to LOCAL anisette server. Set to None to use built-in Anisette generator instead (recommended)
+# IF YOU USE A PUBLIC SERVER, DO NOT COMPLAIN THAT YOU KEEP RUNNING INTO AUTHENTICATION ERRORS!
+# If you change this value, make sure to remove the account store file.
+ANISETTE_SERVER = None
+
+# Path where Anisette libraries will be stored.
+# This is only relevant when using the built-in Anisette server.
+# It can be omitted (set to None) to avoid saving to disk,
+# but specifying a path is highly recommended to avoid downloading the bundle on every run.
+ANISETTE_LIBS_PATH = "ani_libs.bin"
 
 logging.basicConfig(level=logging.INFO)
 
 
 async def fetch_reports(priv_key: str) -> int:
-    key = KeyPair.from_b64(priv_key)
-    acc = await get_account_async(
-        RemoteAnisetteProvider(ANISETTE_SERVER),
-    )
+    # Step 0: construct an account instance
+    # We use a helper for this to simplify interactive authentication
+    acc = await get_account_async(STORE_PATH, ANISETTE_SERVER, ANISETTE_LIBS_PATH)
 
     try:
         print(f"Logged in as: {acc.account_name} ({acc.first_name} {acc.last_name})")
 
-        # It's that simple!
+        # Step 1: construct a key object and get its location reports
+        key = KeyPair.from_b64(priv_key)
         reports = await acc.fetch_last_reports(key)
+
+        # Step 2: print the reports!
         for report in sorted(reports):
             print(report)
+
+            # We can save the report to a file if we want
+            report.to_json("last_report.json")
     finally:
         await acc.close()
 
+        # Make sure to save account state when you're done!
+        acc.to_json(STORE_PATH)
+
     return 0
 
 

examples/real_airtag.py

@@ -11,10 +11,21 @@
 from _login import get_account_sync
 
 from findmy import FindMyAccessory
-from findmy.reports import RemoteAnisetteProvider
 
-# URL to (public or local) anisette server
-ANISETTE_SERVER = "http://localhost:6969"
+# Path where login session will be stored.
+# This is necessary to avoid generating a new session every time we log in.
+STORE_PATH = "account.json"
+
+# URL to LOCAL anisette server. Set to None to use built-in Anisette generator instead (recommended)
+# IF YOU USE A PUBLIC SERVER, DO NOT COMPLAIN THAT YOU KEEP RUNNING INTO AUTHENTICATION ERRORS!
+# If you change this value, make sure to remove the account store file.
+ANISETTE_SERVER = None
+
+# Path where Anisette libraries will be stored.
+# This is only relevant when using the built-in Anisette server.
+# It can be omitted (set to None) to avoid saving to disk,
+# but specifying a path is highly recommended to avoid downloading the bundle on every run.
+ANISETTE_LIBS_PATH = "ani_libs.bin"
 
 logging.basicConfig(level=logging.INFO)
 
@@ -26,8 +37,7 @@ def main(plist_path: str) -> int:
 
     # Step 1: log into an Apple account
     print("Logging into account")
-    anisette = RemoteAnisetteProvider(ANISETTE_SERVER)
-    acc = get_account_sync(anisette)
+    acc = get_account_sync(STORE_PATH, ANISETTE_SERVER, ANISETTE_LIBS_PATH)
 
     # step 2: fetch reports!
     print("Fetching reports")
@@ -39,6 +49,9 @@ def main(plist_path: str) -> int:
     for report in sorted(reports):
         print(f" - {report}")
 
+    # step 4: save current account state to disk
+    acc.to_json(STORE_PATH)
+
     return 0
 
 

findmy/accessory.py

@@ -6,25 +6,40 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import plistlib
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, overload
+from typing import IO, TYPE_CHECKING, Literal, TypedDict, overload
 
 from typing_extensions import override
 
+from findmy.util.abc import Serializable
+from findmy.util.files import read_data_json, save_and_return_json
+
 from .keys import KeyGenerator, KeyPair, KeyType
 from .util import crypto
 
 if TYPE_CHECKING:
-    from collections.abc import Generator, Mapping
+    from collections.abc import Generator
 
 logger = logging.getLogger(__name__)
 
 
+class FindMyAccessoryMapping(TypedDict):
+    """JSON mapping representing state of a FindMyAccessory instance."""
+
+    type: Literal["accessory"]
+    master_key: str
+    skn: str
+    sks: str
+    paired_at: str
+    name: str | None
+    model: str | None
+    identifier: str | None
+
+
 class RollingKeyPairSource(ABC):
     """A class that generates rolling `KeyPair`s."""
 
@@ -67,7 +82,7 @@ def keys_between(self, start: int | datetime, end: int | datetime) -> set[KeyPai
         return keys
 
 
-class FindMyAccessory(RollingKeyPairSource):
+class FindMyAccessory(RollingKeyPairSource, Serializable[FindMyAccessoryMapping]):
     """A findable Find My-accessory using official key rollover."""
 
     def __init__(  # noqa: PLR0913
@@ -242,9 +257,10 @@ def from_plist(
             identifier=identifier,
         )
 
-    def to_json(self, path: str | Path | None = None) -> dict[str, str | int | None]:
-        """Convert the accessory to a JSON-serializable dictionary."""
-        d = {
+    @override
+    def to_json(self, path: str | Path | None = None, /) -> FindMyAccessoryMapping:
+        res: FindMyAccessoryMapping = {
+            "type": "accessory",
             "master_key": self._primary_gen.master_key.hex(),
             "skn": self.skn.hex(),
             "sks": self.sks.hex(),
@@ -253,23 +269,32 @@ def to_json(self, path: str | Path | None = None) -> dict[str, str | int | None]
             "model": self.model,
             "identifier": self.identifier,
         }
-        if path is not None:
-            Path(path).write_text(json.dumps(d, indent=4))
-        return d
+
+        return save_and_return_json(res, path)
 
     @classmethod
-    def from_json(cls, json_: str | Path | Mapping, /) -> FindMyAccessory:
-        """Create a FindMyAccessory from a JSON file."""
-        data = json.loads(Path(json_).read_text()) if isinstance(json_, (str, Path)) else json_
-        return cls(
-            master_key=bytes.fromhex(data["master_key"]),
-            skn=bytes.fromhex(data["skn"]),
-            sks=bytes.fromhex(data["sks"]),
-            paired_at=datetime.fromisoformat(data["paired_at"]),
-            name=data["name"],
-            model=data["model"],
-            identifier=data["identifier"],
-        )
+    @override
+    def from_json(
+        cls,
+        val: str | Path | FindMyAccessoryMapping,
+        /,
+    ) -> FindMyAccessory:
+        val = read_data_json(val)
+        assert val["type"] == "accessory"
+
+        try:
+            return cls(
+                master_key=bytes.fromhex(val["master_key"]),
+                skn=bytes.fromhex(val["skn"]),
+                sks=bytes.fromhex(val["sks"]),
+                paired_at=datetime.fromisoformat(val["paired_at"]),
+                name=val["name"],
+                model=val["model"],
+                identifier=val["identifier"],
+            )
+        except KeyError as e:
+            msg = f"Failed to restore account data: {e}"
+            raise ValueError(msg) from None
 
 
 class AccessoryKeyGenerator(KeyGenerator[KeyPair]):