diff --git a/aimlib/__init__.py b/aimlib/__init__.py index a32ecb3..6a77e3b 100644 --- a/aimlib/__init__.py +++ b/aimlib/__init__.py @@ -7,6 +7,7 @@ import asyncio import importlib +import math import os import re from typing import Optional @@ -194,7 +195,10 @@ def _with_google_chrome_brand(rows: list[dict[str, str]], *, label: str) -> list return branded -def _google_chrome_user_agent_override(identity: object) -> Optional[dict]: +def _google_chrome_user_agent_override( + identity: object, + model_override: Optional[str] = None, +) -> Optional[dict]: """Build a coherent ChromePublic to Google Chrome UA-CH override.""" if not isinstance(identity, dict): raise BrowserPolicyError("managed browser identity metadata is unavailable") @@ -216,6 +220,10 @@ def _google_chrome_user_agent_override(identity: object) -> Optional[dict]: raise BrowserPolicyError("managed browser did not present a mobile Android user agent") if ua_platform != "Android" or identity.get("mobile") is not True: raise BrowserPolicyError("managed browser did not present Android client hints") + if model_override is not None and ( + not isinstance(model_override, str) or not 1 <= len(model_override) <= 128 + ): + raise BrowserPolicyError("managed browser footprint model is invalid") brands = _validated_ua_brand_rows(identity.get("brands"), label="brand") full_versions = _validated_ua_brand_rows( @@ -226,7 +234,8 @@ def _google_chrome_user_agent_override(identity: object) -> Optional[dict]: full_has_google = any(row["brand"] == "Google Chrome" for row in full_versions) if low_has_google != full_has_google: raise BrowserPolicyError("managed browser Chrome brand metadata is inconsistent") - if low_has_google: + effective_model = model_override if model_override is not None else identity["model"] + if low_has_google and effective_model == identity["model"]: return None return { @@ -241,7 +250,7 @@ def _google_chrome_user_agent_override(identity: object) -> Optional[dict]: "platform": ua_platform, "platformVersion": identity["platformVersion"], "architecture": identity["architecture"], - "model": identity["model"], + "model": effective_model, "mobile": True, "bitness": identity["bitness"], "wow64": identity.get("wow64") is True, @@ -249,6 +258,76 @@ def _google_chrome_user_agent_override(identity: object) -> Optional[dict]: } +def _validated_footprint_identity(value: object, applied_footprint: str) -> Optional[dict]: + """Validate the public, browser-observable identity for an applied footprint.""" + if not applied_footprint: + return None + if not isinstance(applied_footprint, str) or len(applied_footprint) > 128: + raise BrowserPolicyError("managed browser footprint identifier is invalid") + if not isinstance(value, dict) or value.get("name") != applied_footprint: + raise BrowserPolicyError("managed browser footprint identity is unavailable") + model = value.get("model") + width = value.get("screen_width") + height = value.get("screen_height") + ratio = value.get("device_pixel_ratio") + if not isinstance(model, str) or not 1 <= len(model) <= 128: + raise BrowserPolicyError("managed browser footprint model is invalid") + if ( + isinstance(width, bool) + or not isinstance(width, int) + or not 320 <= width <= 10_000 + or isinstance(height, bool) + or not isinstance(height, int) + or not 320 <= height <= 10_000 + ): + raise BrowserPolicyError("managed browser footprint display is invalid") + if ( + isinstance(ratio, bool) + or not isinstance(ratio, (int, float)) + or not math.isfinite(float(ratio)) + or not 0.5 <= float(ratio) <= 10 + ): + raise BrowserPolicyError("managed browser footprint pixel ratio is invalid") + return { + "name": applied_footprint, + "model": model, + "screen_width": width, + "screen_height": height, + "device_pixel_ratio": float(ratio), + } + + +def _device_metrics_override(identity: Optional[dict]) -> Optional[dict]: + """Convert physical display pixels into the target's portrait CSS viewport.""" + if identity is None: + return None + ratio = identity["device_pixel_ratio"] + css_width = math.ceil(identity["screen_width"] / ratio) + css_height = math.ceil(identity["screen_height"] / ratio) + return { + "width": css_width, + "height": css_height, + "deviceScaleFactor": ratio, + "mobile": True, + "screenWidth": css_width, + "screenHeight": css_height, + "positionX": 0, + "positionY": 0, + "screenOrientation": {"type": "portraitPrimary", "angle": 0}, + } + + +def _identity_verification(identity: Optional[dict], metrics: Optional[dict]) -> Optional[dict]: + if identity is None or metrics is None: + return None + return { + "model": identity["model"], + "screenWidth": metrics["screenWidth"], + "screenHeight": metrics["screenHeight"], + "devicePixelRatio": metrics["deviceScaleFactor"], + } + + _READ_NATIVE_BROWSER_IDENTITY = r"""async () => { const data = navigator.userAgentData; if (!data) return null; @@ -279,20 +358,25 @@ def _google_chrome_user_agent_override(identity: object) -> Optional[dict]: }""" -_VERIFY_GOOGLE_CHROME_BRAND = r"""async () => { +_VERIFY_MANAGED_BROWSER_IDENTITY = r"""async expected => { const data = navigator.userAgentData; // UA-CH is unavailable on Chrome's internal startup page. That is not an override failure: // the same target exposes the metadata after its first secure navigation. if (!data || !Array.isArray(data.brands)) return null; let high = {}; try { - high = await data.getHighEntropyValues(['fullVersionList']); + high = await data.getHighEntropyValues(['fullVersionList', 'model']); } catch (_) { return null; } if (!Array.isArray(high.fullVersionList)) return null; - return data.brands.some(row => row.brand === 'Google Chrome') && + const branded = data.brands.some(row => row.brand === 'Google Chrome') && high.fullVersionList.some(row => row.brand === 'Google Chrome'); + if (!branded || !expected) return branded; + return high.model === expected.model && + screen.width === expected.screenWidth && + screen.height === expected.screenHeight && + Math.abs(devicePixelRatio - expected.devicePixelRatio) < 0.001; }""" @@ -485,7 +569,7 @@ def __init__(self, ai: "Aimlib", device: Device, data: dict): self.max_tabs = int(data.get("max_tabs") or 5) # per-session tab cap (server-configured) self._gave_initial = False self.egress_ip = None - self.fingerprint = {} + self.fingerprint = data.get("resolved_profile") or {} self.desired_footprint = data.get("desired_footprint") or "" self.applied_footprint = data.get("applied_footprint") or "" self.expires_at = data.get("expires_at") @@ -516,6 +600,7 @@ async def wait_until_ready(self, timeout: float = 180): # Do not connect until a requested browser footprint is fully active. footprint_ok = (not self.desired_footprint) or (self.applied_footprint == self.desired_footprint) if self.status in ("ready", "active", "idle") and footprint_ok: + _validated_footprint_identity(self.fingerprint, self.applied_footprint) return if self.status in ("expiring", "gone"): raise SessionExpiredError(f"session {self.status}") @@ -611,9 +696,15 @@ async def set_footprint(self, footprint: Optional[str], timeout: float = 120): r = await self._ai._http.get(f"/v1/devices/{self.device.id}/browser") if r.status_code == 404: raise SessionExpiredError("session not found / expired") + _raise_for_typed(r) data = r.json() + if data.get("session_id") != self.id: + raise SessionExpiredError("session was replaced by a newer browser session") + self.status = data.get("status") self.applied_footprint = data.get("applied_footprint") or "" - if self.applied_footprint == fp and data.get("status") in ("ready", "active", "idle"): + self.fingerprint = data.get("resolved_profile") or {} + if self.applied_footprint == fp and self.status in ("ready", "active", "idle"): + _validated_footprint_identity(self.fingerprint, self.applied_footprint) await self.connect(timeout=timeout) # transparent reconnect to the same session_id return await asyncio.sleep(2) @@ -646,17 +737,31 @@ async def _apply_page_identity(self, page): return cdp_session = None try: - identity = await self._native_identity_for_page(page) - override = _google_chrome_user_agent_override(identity) - if override is not None: + footprint_identity = _validated_footprint_identity( + self.fingerprint, + self.applied_footprint, + ) + identity = await self._native_identity_for_page(page, footprint_identity) + override = _google_chrome_user_agent_override( + identity, + footprint_identity["model"] if footprint_identity else None, + ) + metrics = _device_metrics_override(footprint_identity) + if override is not None or metrics is not None: cdp_session = await page.context.new_cdp_session(page) - await cdp_session.send("Emulation.setUserAgentOverride", override) - verified = await page.evaluate(_VERIFY_GOOGLE_CHROME_BRAND) + if override is not None: + await cdp_session.send("Emulation.setUserAgentOverride", override) + if metrics is not None: + await cdp_session.send("Emulation.setDeviceMetricsOverride", metrics) + verified = await page.evaluate( + _VERIFY_MANAGED_BROWSER_IDENTITY, + _identity_verification(footprint_identity, metrics), + ) # Internal/about:blank startup pages cannot expose UA-CH. The override has already # been validated on the secure bootstrap tab below; a definitive false result on a # page that does expose UA-CH still fails closed. if verified is False: - raise BrowserPolicyError("managed browser Chrome branding did not apply") + raise BrowserPolicyError("managed browser identity did not apply") self._identity_cdp_sessions.append(cdp_session) cdp_session = None self._identity_pages.add(page_key) @@ -671,7 +776,7 @@ async def _apply_page_identity(self, page): except Exception: # noqa: BLE001 pass - async def _native_identity_for_page(self, page): + async def _native_identity_for_page(self, page, footprint_identity: Optional[dict] = None): """Read UA-CH in a secure first-party context when Chrome's startup tab cannot expose it.""" if self._native_browser_identity is not None: return self._native_browser_identity @@ -687,21 +792,32 @@ async def _native_identity_for_page(self, page): timeout=30_000, ) identity = await bootstrap_page.evaluate(_READ_NATIVE_BROWSER_IDENTITY) - override = _google_chrome_user_agent_override(identity) - if override is not None: + override = _google_chrome_user_agent_override( + identity, + footprint_identity["model"] if footprint_identity else None, + ) + metrics = _device_metrics_override(footprint_identity) + if override is not None or metrics is not None: bootstrap_cdp_session = await page.context.new_cdp_session( bootstrap_page ) - await bootstrap_cdp_session.send( - "Emulation.setUserAgentOverride", - override, - ) + if override is not None: + await bootstrap_cdp_session.send( + "Emulation.setUserAgentOverride", + override, + ) + if metrics is not None: + await bootstrap_cdp_session.send( + "Emulation.setDeviceMetricsOverride", + metrics, + ) verified = await bootstrap_page.evaluate( - _VERIFY_GOOGLE_CHROME_BRAND + _VERIFY_MANAGED_BROWSER_IDENTITY, + _identity_verification(footprint_identity, metrics), ) if verified is not True: raise BrowserPolicyError( - "managed browser Chrome branding did not apply" + "managed browser identity did not apply" ) except Exception as exc: # noqa: BLE001 - keep URL/session details out of the error if isinstance(exc, BrowserPolicyError): diff --git a/pyproject.toml b/pyproject.toml index 935deaf..541b149 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "aimlib" -version = "0.4.2" +version = "0.4.3" description = "Python SDK for aimlib mobile proxies and remote-browser sessions" readme = "README.md" requires-python = ">=3.10" diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 3ca0fb5..006e8c6 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -18,9 +18,11 @@ _Operations, _apply_browser_driver_policy, _async_playwright, + _device_metrics_override, _google_chrome_user_agent_override, _raise_for_typed, _ttl_seconds, + _validated_footprint_identity, ) @@ -208,6 +210,50 @@ def test_does_not_rewrite_an_already_branded_google_chrome(self): self.assertIsNone(_google_chrome_user_agent_override(identity)) + def test_rewrites_model_even_when_google_chrome_is_already_branded(self): + identity = self.chromium_identity() + identity["brands"].append({"brand": "Google Chrome", "version": "151"}) + identity["fullVersionList"].append( + {"brand": "Google Chrome", "version": "151.0.7883.0"} + ) + + override = _google_chrome_user_agent_override(identity, "Pixel 6 Pro") + + self.assertEqual(override["userAgentMetadata"]["model"], "Pixel 6 Pro") + self.assertEqual(override["userAgent"], identity["userAgent"]) + + def test_validates_public_identity_and_builds_pixel_6_pro_metrics(self): + profile = _validated_footprint_identity( + { + "name": "pixel-6-pro", + "model": "Pixel 6 Pro", + "screen_width": 1440, + "screen_height": 3120, + "device_pixel_ratio": 3.5, + }, + "pixel-6-pro", + ) + + self.assertEqual( + _device_metrics_override(profile), + { + "width": 412, + "height": 892, + "deviceScaleFactor": 3.5, + "mobile": True, + "screenWidth": 412, + "screenHeight": 892, + "positionX": 0, + "positionY": 0, + "screenOrientation": {"type": "portraitPrimary", "angle": 0}, + }, + ) + + def test_rejects_missing_or_mismatched_selected_identity(self): + for profile in ({}, {"name": "pixel-6a"}): + with self.subTest(profile=profile), self.assertRaises(BrowserPolicyError): + _validated_footprint_identity(profile, "pixel-6-pro") + def test_fails_closed_on_inconsistent_or_non_android_metadata(self): inconsistent = self.chromium_identity() inconsistent["brands"].append({"brand": "Google Chrome", "version": "151"}) @@ -276,7 +322,7 @@ async def test_bootstrap_fails_closed_when_secure_brand_check_fails(self): with self.assertRaisesRegex( BrowserPolicyError, - "managed browser Chrome branding did not apply", + "managed browser identity did not apply", ): await session._apply_page_identity(startup_page) @@ -327,6 +373,60 @@ async def test_does_not_open_a_cdp_session_for_official_chrome(self): page.context.new_cdp_session.assert_not_awaited() + async def test_applies_selected_model_and_display_metrics_to_the_target(self): + session = session_for(MagicMock()) + session.applied_footprint = "pixel-6-pro" + session.fingerprint = { + "name": "pixel-6-pro", + "model": "Pixel 6 Pro", + "screen_width": 1440, + "screen_height": 3120, + "device_pixel_ratio": 3.5, + } + page = MagicMock() + page.evaluate = AsyncMock( + side_effect=(BrowserIdentityPolicyTests.chromium_identity(), True) + ) + cdp_session = MagicMock() + cdp_session.send = AsyncMock() + cdp_session.detach = AsyncMock() + page.context.new_cdp_session = AsyncMock(return_value=cdp_session) + + await session._apply_page_identity(page) + + self.assertEqual(cdp_session.send.await_count, 2) + ua_call, metrics_call = cdp_session.send.await_args_list + self.assertEqual(ua_call.args[0], "Emulation.setUserAgentOverride") + self.assertEqual( + ua_call.args[1]["userAgentMetadata"]["model"], + "Pixel 6 Pro", + ) + self.assertEqual(metrics_call.args[0], "Emulation.setDeviceMetricsOverride") + self.assertEqual( + metrics_call.args[1], + { + "width": 412, + "height": 892, + "deviceScaleFactor": 3.5, + "mobile": True, + "screenWidth": 412, + "screenHeight": 892, + "positionX": 0, + "positionY": 0, + "screenOrientation": {"type": "portraitPrimary", "angle": 0}, + }, + ) + _, expected = page.evaluate.await_args.args + self.assertEqual( + expected, + { + "model": "Pixel 6 Pro", + "screenWidth": 412, + "screenHeight": 892, + "devicePixelRatio": 3.5, + }, + ) + class ModelTests(unittest.TestCase): def test_device_and_proxy_mapping(self): @@ -469,6 +569,56 @@ async def test_ready_rejects_replaced_session(self): with self.assertRaisesRegex(SessionExpiredError, "replaced"): await session.wait_until_ready(timeout=1) + async def test_ready_fails_closed_when_selected_identity_is_missing(self): + http = MagicMock() + http.get = AsyncMock( + return_value=response( + 200, + { + "session_id": "session-1", + "status": "ready", + "desired_footprint": "pixel-6-pro", + "applied_footprint": "pixel-6-pro", + "resolved_profile": {"name": "pixel-6-pro"}, + }, + ) + ) + session = session_for(http) + session.desired_footprint = "pixel-6-pro" + + with self.assertRaisesRegex(BrowserPolicyError, "footprint model is invalid"): + await session.wait_until_ready(timeout=1) + + async def test_set_footprint_refreshes_identity_before_reconnect(self): + http = MagicMock() + http.post = AsyncMock(return_value=response(200, {"status": "accepted"}, "POST")) + profile = { + "name": "pixel-6-pro", + "model": "Pixel 6 Pro", + "screen_width": 1440, + "screen_height": 3120, + "device_pixel_ratio": 3.5, + } + http.get = AsyncMock( + return_value=response( + 200, + { + "session_id": "session-1", + "status": "ready", + "applied_footprint": "pixel-6-pro", + "resolved_profile": profile, + }, + ) + ) + session = session_for(http) + session._disconnect = AsyncMock() + session.connect = AsyncMock() + + await session.set_footprint("pixel-6-pro", timeout=1) + + self.assertEqual(session.fingerprint, profile) + session.connect.assert_awaited_once_with(timeout=1) + async def test_connect_retries_and_stops_failed_playwright_instance(self): http = MagicMock() http.get = AsyncMock(