diff --git a/examples/fetch_reports.py b/examples/fetch_reports.py index 321a775..d4fb28a 100644 --- a/examples/fetch_reports.py +++ b/examples/fetch_reports.py @@ -24,7 +24,7 @@ def fetch_reports(lookup_key: KeyPair) -> None: print(f"Logged in as: {acc.account_name} ({acc.first_name} {acc.last_name})") # It's that simple! - reports = acc.fetch_last_reports([lookup_key])[lookup_key] + reports = acc.fetch_last_reports(lookup_key) for report in sorted(reports): print(report) diff --git a/examples/fetch_reports_async.py b/examples/fetch_reports_async.py index 754a05f..1e73b4e 100644 --- a/examples/fetch_reports_async.py +++ b/examples/fetch_reports_async.py @@ -27,8 +27,9 @@ async def fetch_reports(lookup_key: KeyPair) -> None: print(f"Logged in as: {acc.account_name} ({acc.first_name} {acc.last_name})") # It's that simple! - reports = await acc.fetch_last_reports([lookup_key]) - print(reports) + reports = await acc.fetch_last_reports(lookup_key) + for report in sorted(reports): + print(report) finally: await acc.close() diff --git a/examples/real_airtag.py b/examples/real_airtag.py index e669050..a259e35 100644 --- a/examples/real_airtag.py +++ b/examples/real_airtag.py @@ -50,30 +50,18 @@ def main() -> None: # Step 0: create an accessory key generator airtag = FindMyAccessory(MASTER_KEY, SKN, SKS, PAIRED_AT) - # Step 1: Generate the accessory's private keys, - # starting from 7 days ago until now (12 hour margin) - fetch_to = datetime.now(tz=timezone.utc).astimezone() + timedelta(hours=12) - fetch_from = fetch_to - timedelta(days=8) - - print(f"Generating keys from {fetch_from} to {fetch_to} ...") - lookup_keys = _gen_keys(airtag, fetch_from, fetch_to) - - print(f"Generated {len(lookup_keys)} keys") - - # Step 2: log into an Apple account + # Step 1: log into an Apple account print("Logging into account") anisette = RemoteAnisetteProvider(ANISETTE_SERVER) acc = get_account_sync(anisette) - # step 3: fetch reports! + # step 2: fetch reports! print("Fetching reports") - reports = acc.fetch_reports(list(lookup_keys), fetch_from, fetch_to) + reports = acc.fetch_last_reports(airtag) - # step 4: print 'em - # reports are in {key: [report]} format, but we only really care about the reports + # step 3: print 'em print() print("Location reports:") - reports = sorted([r for rs in reports.values() for r in rs]) for report in reports: print(f" - {report}") diff --git a/findmy/reports/account.py b/findmy/reports/account.py index 1efb967..c6ff3ca 100644 --- a/findmy/reports/account.py +++ b/findmy/reports/account.py @@ -20,6 +20,7 @@ from typing import ( TypedDict, TypeVar, cast, + overload, ) import bs4 @@ -44,6 +45,7 @@ from .twofactor import ( ) if TYPE_CHECKING: + from findmy.accessory import RollingKeyPairSource from findmy.keys import KeyPair from findmy.util.types import MaybeCoro @@ -215,6 +217,17 @@ class BaseAppleAccount(Closable, ABC): """ raise NotImplementedError + @overload + @abstractmethod + def fetch_reports( + self, + keys: KeyPair, + date_from: datetime, + date_to: datetime | None, + ) -> MaybeCoro[list[LocationReport]]: + ... + + @overload @abstractmethod def fetch_reports( self, @@ -222,6 +235,25 @@ class BaseAppleAccount(Closable, ABC): date_from: datetime, date_to: datetime | None, ) -> MaybeCoro[dict[KeyPair, list[LocationReport]]]: + ... + + @overload + @abstractmethod + def fetch_reports( + self, + keys: RollingKeyPairSource, + date_from: datetime, + date_to: datetime | None, + ) -> MaybeCoro[list[LocationReport]]: + ... + + @abstractmethod + def fetch_reports( + self, + keys: KeyPair | Sequence[KeyPair] | RollingKeyPairSource, + date_from: datetime, + date_to: datetime | None, + ) -> MaybeCoro[list[LocationReport] | dict[KeyPair, list[LocationReport]]]: """ Fetch location reports for a sequence of `KeyPair`s between `date_from` and `date_end`. @@ -229,12 +261,39 @@ class BaseAppleAccount(Closable, ABC): """ raise NotImplementedError + @overload + @abstractmethod + def fetch_last_reports( + self, + keys: KeyPair, + hours: int = 7 * 24, + ) -> MaybeCoro[list[LocationReport]]: + ... + + @overload @abstractmethod def fetch_last_reports( self, keys: Sequence[KeyPair], hours: int = 7 * 24, ) -> MaybeCoro[dict[KeyPair, list[LocationReport]]]: + ... + + @overload + @abstractmethod + def fetch_last_reports( + self, + keys: RollingKeyPairSource, + hours: int = 7 * 24, + ) -> MaybeCoro[list[LocationReport]]: + ... + + @abstractmethod + def fetch_last_reports( + self, + keys: KeyPair | Sequence[KeyPair] | RollingKeyPairSource, + hours: int = 7 * 24, + ) -> MaybeCoro[list[LocationReport] | dict[KeyPair, list[LocationReport]]]: """ Fetch location reports for a sequence of `KeyPair`s for the last `hours` hours. @@ -539,14 +598,41 @@ class AsyncAppleAccount(BaseAppleAccount): return resp - @require_login_state(LoginState.LOGGED_IN) - @override + @overload + async def fetch_reports( + self, + keys: KeyPair, + date_from: datetime, + date_to: datetime | None, + ) -> list[LocationReport]: + ... + + @overload async def fetch_reports( self, keys: Sequence[KeyPair], date_from: datetime, date_to: datetime | None, ) -> dict[KeyPair, list[LocationReport]]: + ... + + @overload + async def fetch_reports( + self, + keys: RollingKeyPairSource, + date_from: datetime, + date_to: datetime | None, + ) -> list[LocationReport]: + ... + + @require_login_state(LoginState.LOGGED_IN) + @override + async def fetch_reports( + self, + keys: KeyPair | Sequence[KeyPair] | RollingKeyPairSource, + date_from: datetime, + date_to: datetime | None, + ) -> list[LocationReport] | dict[KeyPair, list[LocationReport]]: """See `BaseAppleAccount.fetch_reports`.""" date_to = date_to or datetime.now().astimezone() @@ -556,13 +642,37 @@ class AsyncAppleAccount(BaseAppleAccount): keys, ) - @require_login_state(LoginState.LOGGED_IN) - @override + @overload + async def fetch_last_reports( + self, + keys: KeyPair, + hours: int = 7 * 24, + ) -> list[LocationReport]: + ... + + @overload async def fetch_last_reports( self, keys: Sequence[KeyPair], hours: int = 7 * 24, ) -> dict[KeyPair, list[LocationReport]]: + ... + + @overload + async def fetch_last_reports( + self, + keys: RollingKeyPairSource, + hours: int = 7 * 24, + ) -> list[LocationReport]: + ... + + @require_login_state(LoginState.LOGGED_IN) + @override + async def fetch_last_reports( + self, + keys: KeyPair | Sequence[KeyPair] | RollingKeyPairSource, + hours: int = 7 * 24, + ) -> list[LocationReport] | dict[KeyPair, list[LocationReport]]: """See `BaseAppleAccount.fetch_last_reports`.""" end = datetime.now(tz=timezone.utc) start = end - timedelta(hours=hours) @@ -894,23 +1004,74 @@ class AppleAccount(BaseAppleAccount): coro = self._asyncacc.td_2fa_submit(code) return self._evt_loop.run_until_complete(coro) - @override + @overload + def fetch_reports( + self, + keys: KeyPair, + date_from: datetime, + date_to: datetime | None, + ) -> list[LocationReport]: + ... + + @overload def fetch_reports( self, keys: Sequence[KeyPair], date_from: datetime, date_to: datetime | None, ) -> dict[KeyPair, list[LocationReport]]: + ... + + @overload + def fetch_reports( + self, + keys: RollingKeyPairSource, + date_from: datetime, + date_to: datetime | None, + ) -> list[LocationReport]: + ... + + @override + def fetch_reports( + self, + keys: KeyPair | Sequence[KeyPair] | RollingKeyPairSource, + date_from: datetime, + date_to: datetime | None, + ) -> list[LocationReport] | dict[KeyPair, list[LocationReport]]: """See `AsyncAppleAccount.fetch_reports`.""" coro = self._asyncacc.fetch_reports(keys, date_from, date_to) return self._evt_loop.run_until_complete(coro) - @override + @overload + def fetch_last_reports( + self, + keys: KeyPair, + hours: int = 7 * 24, + ) -> list[LocationReport]: + ... + + @overload def fetch_last_reports( self, keys: Sequence[KeyPair], hours: int = 7 * 24, ) -> dict[KeyPair, list[LocationReport]]: + ... + + @overload + def fetch_last_reports( + self, + keys: RollingKeyPairSource, + hours: int = 7 * 24, + ) -> list[LocationReport]: + ... + + @override + def fetch_last_reports( + self, + keys: KeyPair | Sequence[KeyPair] | RollingKeyPairSource, + hours: int = 7 * 24, + ) -> list[LocationReport] | dict[KeyPair, list[LocationReport]]: """See `AsyncAppleAccount.fetch_last_reports`.""" coro = self._asyncacc.fetch_last_reports(keys, hours) return self._evt_loop.run_until_complete(coro) diff --git a/findmy/reports/reports.py b/findmy/reports/reports.py index a8a83b9..f853506 100644 --- a/findmy/reports/reports.py +++ b/findmy/reports/reports.py @@ -5,7 +5,7 @@ import base64 import hashlib import logging import struct -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Sequence, overload from cryptography.hazmat.backends import default_backend @@ -13,6 +13,7 @@ from cryptography.hazmat.primitives.asymmetric import ec from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes from typing_extensions import override +from findmy.accessory import RollingKeyPairSource from findmy.keys import KeyPair if TYPE_CHECKING: @@ -192,11 +193,20 @@ class LocationReportsFetcher: ) -> dict[KeyPair, list[LocationReport]]: ... + @overload async def fetch_reports( self, date_from: datetime, date_to: datetime, - device: KeyPair | Sequence[KeyPair], + device: RollingKeyPairSource, + ) -> list[LocationReport]: + ... + + async def fetch_reports( + self, + date_from: datetime, + date_to: datetime, + device: KeyPair | Sequence[KeyPair] | RollingKeyPairSource, ) -> list[LocationReport] | dict[KeyPair, list[LocationReport]]: """ Fetch location reports for a certain device. @@ -210,13 +220,28 @@ class LocationReportsFetcher: if isinstance(device, KeyPair): return await self._fetch_reports(date_from, date_to, [device]) + # KeyPair generator + # add 12h margin to the generator + if isinstance(device, RollingKeyPairSource): + keys = list( + device.keys_between( + date_from - timedelta(hours=12), + date_to + timedelta(hours=12), + ), + ) + else: + keys = device + # sequence of KeyPairs (fetch 256 max at a time) reports: list[LocationReport] = [] - for key_offset in range(0, len(device), 256): - chunk = device[key_offset : key_offset + 256] + for key_offset in range(0, len(keys), 256): + chunk = keys[key_offset : key_offset + 256] reports.extend(await self._fetch_reports(date_from, date_to, chunk)) - res: dict[KeyPair, list[LocationReport]] = {key: [] for key in device} + if isinstance(device, RollingKeyPairSource): + return reports + + res: dict[KeyPair, list[LocationReport]] = {key: [] for key in keys} for report in reports: res[report.key].append(report) return res