pyatomberg 0.1.0__py3-none-any.whl

pyatomberg/__init__.py

@@ -0,0 +1,79 @@
+"""pyatomberg: Python library for Atomberg smart home authentication, discovery, and protocol codecs."""
+
+from pyatomberg.auth.cognito import CognitoOAuthHandler
+from pyatomberg.auth.developer import DeveloperTokenHandler
+from pyatomberg.client import AtombergClient
+from pyatomberg.exceptions import (
+    AtombergAuthError,
+    AtombergBLEError,
+    AtombergConnectionError,
+    AtombergCryptoError,
+    AtombergDeviceOfflineError,
+    AtombergError,
+    AtombergPairingError,
+)
+from pyatomberg.fan.device import AtombergFan
+from pyatomberg.fan.udp_listener import AtombergUDPListener
+from pyatomberg.lock.cloud_sync import AtombergCloudSync
+from pyatomberg.lock.dynamic_pin import (
+    AtombergDynamicPinClient,
+    calculate_offline_otp,
+    quantize_expiration_time,
+)
+from pyatomberg.lock.hsj_protocol import (
+    DEFAULT_SLOT_MAPPINGS,
+    LiveEvent,
+    build_hsj_packet,
+    format_ist_time,
+    parse_hsj_packet,
+    parse_live_event,
+    parse_log_record,
+    slot_name,
+)
+from pyatomberg.lock.pairing import (
+    PairingResult,
+    extract_keys_from_btsnoop,
+)
+from pyatomberg.models.device import (
+    AtombergDeviceModel,
+    DeviceType,
+    FanState,
+    LockState,
+    WaterPurifierState,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "PairingResult",
+    "extract_keys_from_btsnoop",
+    "AtombergCloudSync",
+    "AtombergDynamicPinClient",
+    "calculate_offline_otp",
+    "quantize_expiration_time",
+    "LiveEvent",
+    "build_hsj_packet",
+    "parse_hsj_packet",
+    "parse_live_event",
+    "parse_log_record",
+    "slot_name",
+    "format_ist_time",
+    "DEFAULT_SLOT_MAPPINGS",
+    "AtombergClient",
+    "AtombergFan",
+    "AtombergUDPListener",
+    "DeveloperTokenHandler",
+    "CognitoOAuthHandler",
+    "AtombergError",
+    "AtombergAuthError",
+    "AtombergConnectionError",
+    "AtombergDeviceOfflineError",
+    "AtombergBLEError",
+    "AtombergCryptoError",
+    "AtombergPairingError",
+    "AtombergDeviceModel",
+    "DeviceType",
+    "FanState",
+    "LockState",
+    "WaterPurifierState",
+]

pyatomberg/auth/__init__.py

@@ -0,0 +1,6 @@
+"""Authentication handlers for pyatomberg."""
+
+from pyatomberg.auth.cognito import CognitoOAuthHandler
+from pyatomberg.auth.developer import DeveloperTokenHandler
+
+__all__ = ["DeveloperTokenHandler", "CognitoOAuthHandler"]

pyatomberg/auth/cognito.py

@@ -0,0 +1,177 @@
+"""AWS Cognito Hosted UI OAuth2 authentication handler."""
+
+import base64
+import hashlib
+import os
+import secrets
+import time
+import urllib.parse
+from typing import Any
+
+import aiohttp
+
+from pyatomberg.const import (
+    COGNITO_AUTHORIZE_URL,
+    COGNITO_CLIENT_ID,
+    COGNITO_REDIRECT_URI,
+    COGNITO_TOKEN_URL,
+)
+from pyatomberg.exceptions import AtombergAuthError, AtombergConnectionError
+
+
+def generate_pkce_pair() -> tuple[str, str]:
+    """Generate a PKCE (code_verifier, code_challenge) tuple using S256."""
+    verifier = base64.urlsafe_b64encode(os.urandom(32)).decode("utf-8").rstrip("=")
+    challenge_bytes = hashlib.sha256(verifier.encode("utf-8")).digest()
+    challenge = base64.urlsafe_b64encode(challenge_bytes).decode("utf-8").rstrip("=")
+    return verifier, challenge
+
+
+class CognitoOAuthHandler:
+    """Manages OAuth2 authorization code flow with AWS Cognito Hosted UI.
+
+    Supports third-party identity providers such as Google Sign-In and Apple Sign-In
+    federated via auth.atomberg-iot.com.
+    """
+
+    def __init__(
+        self,
+        client_id: str = COGNITO_CLIENT_ID,
+        redirect_uri: str = COGNITO_REDIRECT_URI,
+        refresh_token: str | None = None,
+        session: aiohttp.ClientSession | None = None,
+    ) -> None:
+        """Initialize the Cognito OAuth handler."""
+        self._client_id = client_id
+        self._redirect_uri = redirect_uri
+        self._refresh_token = refresh_token
+        self._id_token: str | None = None
+        self._access_token: str | None = None
+        self._expires_at: float = 0.0
+        self._session = session
+
+        # Active PKCE verifier for pending flow
+        self._pending_verifier: str | None = None
+        self._pending_state: str | None = None
+
+    @property
+    def refresh_token(self) -> str | None:
+        """Return the Cognito refresh token."""
+        return self._refresh_token
+
+    @property
+    def id_token(self) -> str | None:
+        """Return the Cognito ID token used for API requests."""
+        return self._id_token
+
+    @property
+    def access_token(self) -> str | None:
+        """Return the Cognito Access token."""
+        return self._access_token
+
+    @property
+    def is_authenticated(self) -> bool:
+        """Check if an access token is cached and not expired."""
+        return self._access_token is not None and time.time() < (self._expires_at - 60)
+
+    def get_authorization_url(
+        self,
+        redirect_uri: str | None = None,
+        identity_provider: str | None = None,
+        scope: str = "aws.cognito.signin.user.admin openid email",
+    ) -> tuple[str, str, str]:
+        """Generate the Cognito Hosted UI authorization URL with PKCE.
+
+        Returns (auth_url, state, code_verifier).
+        """
+        if redirect_uri:
+            self._redirect_uri = redirect_uri
+
+        verifier, challenge = generate_pkce_pair()
+        state = secrets.token_urlsafe(16)
+        self._pending_verifier = verifier
+        self._pending_state = state
+
+        params: dict[str, str] = {
+            "client_id": self._client_id,
+            "response_type": "code",
+            "scope": scope,
+            "redirect_uri": self._redirect_uri,
+            "state": state,
+            "code_challenge": challenge,
+            "code_challenge_method": "S256",
+        }
+        if identity_provider:
+            params["identity_provider"] = identity_provider
+
+        url = f"{COGNITO_AUTHORIZE_URL}?{urllib.parse.urlencode(params)}"
+        return url, state, verifier
+
+    async def exchange_code_for_tokens(
+        self,
+        code: str,
+        code_verifier: str | None = None,
+        redirect_uri: str | None = None,
+    ) -> dict[str, Any]:
+        """Exchange authorization code for Cognito tokens."""
+        verifier = code_verifier or self._pending_verifier
+        if not verifier:
+            raise AtombergAuthError("No PKCE code_verifier provided or found in pending state")
+
+        r_uri = redirect_uri or self._redirect_uri
+        payload = {
+            "grant_type": "authorization_code",
+            "client_id": self._client_id,
+            "code": code,
+            "redirect_uri": r_uri,
+            "code_verifier": verifier,
+        }
+
+        return await self._post_token_request(payload)
+
+    async def refresh_tokens(self) -> dict[str, Any]:
+        """Refresh tokens using the stored refresh_token."""
+        if not self._refresh_token:
+            raise AtombergAuthError("No refresh token available to refresh Cognito session")
+
+        payload = {
+            "grant_type": "refresh_token",
+            "client_id": self._client_id,
+            "refresh_token": self._refresh_token,
+        }
+        return await self._post_token_request(payload)
+
+    async def _post_token_request(self, payload: dict[str, str]) -> dict[str, Any]:
+        """Post request to Cognito token endpoint."""
+        close_session = False
+        session = self._session
+        if session is None:
+            session = aiohttp.ClientSession()
+            close_session = True
+
+        headers = {"Content-Type": "application/x-www-form-urlencoded"}
+        try:
+            async with session.post(
+                COGNITO_TOKEN_URL,
+                data=payload,
+                headers=headers,
+                timeout=aiohttp.ClientTimeout(total=15),
+            ) as resp:
+                data: dict[str, Any] = await resp.json(content_type=None)
+                if resp.status != 200:
+                    err = data.get("error_description") or data.get("error") or f"HTTP {resp.status}"
+                    raise AtombergAuthError(f"Cognito token request failed: {err}")
+
+                self._access_token = data.get("access_token")
+                self._id_token = data.get("id_token")
+                if "refresh_token" in data:
+                    self._refresh_token = data["refresh_token"]
+
+                expires_in = data.get("expires_in", 3600)
+                self._expires_at = time.time() + float(expires_in)
+                return data
+        except aiohttp.ClientError as err:
+            raise AtombergConnectionError(f"Network error communicating with Cognito: {err}") from err
+        finally:
+            if close_session:
+                await session.close()

pyatomberg/auth/developer.py

@@ -0,0 +1,121 @@
+"""Developer API token authentication handler."""
+
+import time
+from typing import Any
+
+import aiohttp
+
+from pyatomberg.const import DEVELOPER_API_BASE
+from pyatomberg.exceptions import AtombergAuthError, AtombergConnectionError
+
+
+class DeveloperTokenHandler:
+    """Manages authentication using Atomberg Developer Refresh Tokens.
+
+    Allows credential-free operation for fans by exchanging the refresh token
+    obtained from the Atomberg mobile app for short-lived bearer access tokens.
+    """
+
+    def __init__(
+        self,
+        refresh_token: str,
+        api_key: str | None = None,
+        session: aiohttp.ClientSession | None = None,
+        base_url: str = DEVELOPER_API_BASE,
+    ) -> None:
+        """Initialize the developer token handler."""
+        if not refresh_token or not refresh_token.strip():
+            raise AtombergAuthError("Developer refresh token cannot be empty")
+
+        self._refresh_token = refresh_token.strip()
+        self._api_key = api_key.strip() if api_key else None
+        self._session = session
+        self._base_url = base_url.rstrip("/")
+        self._access_token: str | None = None
+        self._expires_at: float = 0.0
+
+    @property
+    def refresh_token(self) -> str:
+        """Return the configured refresh token."""
+        return self._refresh_token
+
+    @property
+    def api_key(self) -> str | None:
+        """Return the configured API key, if any."""
+        return self._api_key
+
+    @property
+    def has_valid_access_token(self) -> bool:
+        """Check if cached access token is valid (with 60s safety buffer)."""
+        return self._access_token is not None and time.time() < (self._expires_at - 60)
+
+    async def get_access_token(self, force_refresh: bool = False) -> str:
+        """Retrieve a valid access token, refreshing if necessary."""
+        if not force_refresh and self.has_valid_access_token:
+            assert self._access_token is not None
+            return self._access_token
+
+        await self.refresh_access_token()
+        assert self._access_token is not None
+        return self._access_token
+
+    async def refresh_access_token(self) -> str:
+        """Exchange developer refresh token for a new access token."""
+        url = f"{self._base_url}/get_access_token"
+        headers: dict[str, str] = {
+            "Authorization": f"Bearer {self._refresh_token}",
+        }
+        if self._api_key:
+            headers["x-api-key"] = self._api_key
+
+        close_session = False
+        session = self._session
+        if session is None:
+            session = aiohttp.ClientSession()
+            close_session = True
+
+        try:
+            # First attempt standard GET as documented in Developer Portal v2
+            async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=15)) as resp:
+                if resp.status == 405:
+                    # Fallback to POST for v1 compatibility
+                    async with session.post(
+                        url,
+                        json={"refresh_token": self._refresh_token},
+                        headers=headers,
+                        timeout=aiohttp.ClientTimeout(total=15),
+                    ) as post_resp:
+                        return await self._parse_token_response(post_resp)
+                return await self._parse_token_response(resp)
+        except aiohttp.ClientError as err:
+            raise AtombergConnectionError(f"Network error refreshing access token: {err}") from err
+        finally:
+            if close_session:
+                await session.close()
+
+    async def _parse_token_response(self, resp: aiohttp.ClientResponse) -> str:
+        """Parse access token JSON response from developer endpoint."""
+        data: dict[str, Any] = await resp.json(content_type=None)
+        if resp.status != 200 or not data.get("status", True):
+            error_msg = data.get("message") or data.get("error") or f"HTTP {resp.status}"
+            raise AtombergAuthError(f"Failed to refresh developer token: {error_msg}")
+
+        token = data.get("access_token") or data.get("token") or data.get("data", {}).get("access_token")
+        if not token:
+            raise AtombergAuthError("Response did not contain an access_token")
+
+        expires_in = data.get("expires_in", 3600)
+        self._access_token = str(token)
+        self._expires_at = time.time() + float(expires_in)
+        return self._access_token
+
+    async def get_auth_headers(self) -> dict[str, str]:
+        """Return HTTP authorization headers with a valid Bearer token and API key."""
+        token = await self.get_access_token()
+        headers: dict[str, str] = {
+            "Authorization": f"Bearer {token}",
+            "Content-Type": "application/json",
+        }
+        if self._api_key:
+            headers["x-api-key"] = self._api_key
+        return headers

pyatomberg/auth/interactive_oauth.py

@@ -0,0 +1,100 @@
+"""Interactive Google OAuth setup script for Atomberg accounts.
+
+Enables users to sign in with Google in their browser and automatically
+exchanges the resulting authorization code for permanent Cognito refresh and ID tokens.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+import urllib.parse
+
+from pyatomberg.auth.cognito import CognitoOAuthHandler
+from pyatomberg.const import COGNITO_CLIENT_ID, COGNITO_REDIRECT_URI
+from pyatomberg.lock.cloud_sync import AtombergCloudSync
+
+
+async def run_interactive_login() -> tuple[dict[str, str], list[Any]]:
+    """Run interactive CLI Google login and return tokens and discovered locks."""
+    handler = CognitoOAuthHandler(
+        client_id=COGNITO_CLIENT_ID,
+        redirect_uri=COGNITO_REDIRECT_URI,
+    )
+
+    auth_url, state, verifier = handler.get_authorization_url(identity_provider=None)
+
+    print("\n" + "=" * 70)
+    print(" ATOMBERG UNIVERSAL ACCOUNT LOGIN (EMAIL / GOOGLE / APPLE / FACEBOOK)")
+    print("=" * 70)
+    print("\n1. Open the following URL in your web browser:\n")
+    print(auth_url)
+    print("\n2. Sign in with your Atomberg email & password, Google, Apple, or Facebook.")
+    print("3. After signing in, your browser will redirect to https://app.atomberg-iot.com/?code=...")
+    print("4. Copy the entire URL (or just the code value) from your browser address bar.")
+    print("=" * 70 + "\n")
+
+    user_response = input("Paste the redirected URL or authorization code here: ").strip()
+    if not user_response:
+        print("[!] No input provided. Aborting.")
+        return {}, []
+
+    # Extract authorization code
+    auth_code = user_response
+    if "code=" in user_response:
+        parsed = urllib.parse.urlparse(user_response)
+        params = urllib.parse.parse_qs(parsed.query)
+        code_list = params.get("code")
+        if code_list:
+            auth_code = code_list[0]
+        else:
+            # Maybe in fragment or raw string
+            for part in user_response.split("&"):
+                if part.startswith("code=") or "code=" in part:
+                    auth_code = part.split("code=")[-1]
+                    break
+
+    print(f"\n[*] Exchanging authorization code for Cognito tokens...")
+    tokens = await handler.exchange_code_for_tokens(
+        code=auth_code,
+        code_verifier=verifier,
+        redirect_uri=COGNITO_REDIRECT_URI,
+    )
+
+    access_token = tokens.get("access_token", "")
+    id_token = tokens.get("id_token", "")
+    refresh_token = tokens.get("refresh_token", "")
+
+    print("\n[+] Authentication successful!")
+    print(f"    Refresh Token (first 20 chars): {refresh_token[:20]}...")
+    print(f"    Access Token (first 20 chars):  {access_token[:20]}...")
+    print(f"    ID Token (first 20 chars):      {id_token[:20]}...")
+
+    print("\n[*] Querying Atomberg cloud account (server.atomberg-iot.com/app_open) for smart locks...")
+    sync_client = AtombergCloudSync(access_token=access_token)
+    try:
+        locks = await sync_client.fetch_account_locks()
+        if not locks:
+            print("[-] No smart locks found in this Atomberg account.")
+        else:
+            print(f"[+] Discovered {len(locks)} smart lock(s) registered in your account:\n")
+            for idx, lock in enumerate(locks, 1):
+                print(f"    Lock #{idx}:")
+                print(f"      Name:        {lock.name or 'Atomberg Smart Lock'}")
+                print(f"      Model:       {lock.model or 'smart_lock'}")
+                print(f"      MAC Address: {lock.mac_address}")
+                print(f"      Master Key:  {lock.master_key_hex}")
+                print(f"      Salt:        {lock.salt_hex}")
+        return tokens, locks
+    except Exception as err:
+        print(f"[!] Error syncing account locks: {err}")
+        return tokens, []
+
+
+def main() -> None:
+    """Entry point for CLI command."""
+    asyncio.run(run_interactive_login())
+
+
+if __name__ == "__main__":
+    main()