pymammotion 0.3.8__py3-none-any.whl → 0.4.0a0__py3-none-any.whl

pymammotion/http/encryption.py

@@ -0,0 +1,221 @@
+import base64
+import logging
+import secrets
+import string
+from typing import Optional
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import padding, serialization
+from cryptography.hazmat.primitives.asymmetric import padding as rsa_padding
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class EncryptionUtils:
+    PRIVATE_KEY = """MIICdwIBADANBgkqhkiG9w0BAQEFAASCAmEwggJdAgEAAoGBAOFizbd1fC5XNKJ89u0XNvPZNR/L
+                         0h547iSWjOCuvvMu76ZSaS3/Tu2C1C+XmlnmBWTyY4ON+xECiNUXm/aWQ3P0g+wf60zjPbNzgL2Q
+                         7njXJG6wka4KkbdQxUdS0TTpL256LnV1LsG855bsbJIJiQPbfUq6HbB5xH7sXdrmFu1DAgMBAAEC
+                         gYEAoT2TGE1ncquWjyxBZup1uMvKkp25C23OSMSfslmxZ75LWjyY3HxK1eYDsKyPkwLZFxfFE6du
+                         VwPuKiyCuk1ToPfnb4niTGzXPyC2PbO4SFrWL8n1YZ80M0bfTGI9dMCZvpmZJ41WYUsBaf2374lt
+                         oEiDEHJp7MeXk/970xiKP1ECQQD65rLHk840q+FZS6kZVexJucPZj/YAII6klU1E20ctioe8Pi5m
+                         WSPqclH27/t4FqdvP7tFqaavyXg+CEQpxmxLAkEA5fddDuzcjWgF9pl9fP7/baFMYjUS9z1Vc3gx
+                         CnvAgCnv71wjDQhvsUc6sAiidsBGFDyud06RyyLcOlQchMb36QJBAIui/Xjpn+fciQxjeXcqRNk7
+                         U+6vml+zvu+GUHyz9Uc5RBXWHYjEr6J5gXiHU1MgeIsH0zgQFT7cR9luTFFbp0UCQFIntfogCocG
+                         E6NOoHMoUi5jQnuPRHBJXB69YJ/DKDlhQhN8EhWU3voxXTkITKop9J9EMnvy+MjecljwNaQFxQkC
+                         QB9lz67iDe9Gj8NxSElVZxUm9EfbL1RPqTZPx/lADR06CPB8pP3Bl5/5/5RGzc+UTZ+wX5GWKvC7
+                         zUJaROxQB+E=""".replace(" ", "")
+
+    PUBLIC_KEY_PROD = """MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEApLbeSgOvnwLTWbhaBQWNnnHMtSDAi
+                            Gz0PEDbrtd1tLYoO0hukW5PSa6eHykch0Hc6etiqEx1xziS+vNf+iOXds70I4htaYit6yRToZlQ
+                            Mim3DQxaZX68nIHIZogur0zGv9U8j01v5l/rHRxyDdlVx3+JkBg6Cqx4U1PXEnAJriqcyg0B8Gm
+                            V8Lnmfng+aJLRyq5MkhstYCRv9AsmWu8NpZDJ1ffbkaS02Z9/wpubXTiFP6DG3V2mDw2VvzEcHi
+                            cchw49oXmTi92yui+kBgSYlNygssOAyU6H071AfmRUeH3+TsV5u5rg+bCiKyHemVmcKdd3hhZB+
+                            HjA8o3On6rg5wIDAQAB""".replace(" ", "")
+
+    PUBLIC_KEY_TEST = """MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQC1nAzH31arNBmYKvTlvKgkxI1MIr4HpfLbmM
+                            XPIhd8D/cXB0dYY1ppUq4a/ezq41YShN88e0elyZgqdnFrkhiLpnKWa7jXtVRgXi9eS18PLO8ns
+                            eHude9URaj7relK1AZ0xovKsbLKHd01PpmngLXZfnKA06J2ru/zH+cnpXdy8QIDAQAB""".replace(" ", "")
+
+    def __init__(self) -> None:
+        self.AES_PASW = self.get_aes_key()  # Get from previous implementation
+        self.IV = self.get_iv()  # Get from previous implementation
+        self._public_key = self.load_public_key()
+        self._private_key = self.load_private_key()
+
+    @staticmethod
+    def load_private_key():
+        """Load the private key from base64 encoded string"""
+        try:
+            private_key_bytes = base64.b64decode(EncryptionUtils.PRIVATE_KEY)
+            return serialization.load_der_private_key(private_key_bytes, password=None, backend=default_backend())
+        except Exception as e:
+            raise Exception(f"Failed to load private key: {str(e)}")
+
+    @staticmethod
+    def load_public_key(is_production: bool = True):
+        """Load the public key from base64 encoded string
+
+        Args:
+            is_production (bool): If True, uses production key, else uses test key
+
+        """
+        try:
+            key_string = EncryptionUtils.PUBLIC_KEY_PROD if is_production else EncryptionUtils.PUBLIC_KEY_TEST
+            public_key_bytes = base64.b64decode(key_string)
+            return serialization.load_der_public_key(public_key_bytes, backend=default_backend())
+        except Exception as e:
+            raise Exception(f"Failed to load public key: {str(e)}")
+
+    @staticmethod
+    def encrypt(plaintext: str, key: str, iv: str) -> str:
+        """Encrypt text using AES/CBC/PKCS5Padding
+
+        Args:
+            plaintext (str): Text to encrypt
+            key (str): Encryption key
+            iv (str): Initialization vector
+
+        Returns:
+            str: Base64 encoded encrypted string
+
+        Raises:
+            Exception: If encryption fails
+
+        """
+        try:
+            # Convert strings to bytes
+            plaintext_bytes = plaintext.encode("utf-8")
+            key_bytes = key.encode("utf-8")
+            iv_bytes = iv.encode("utf-8")
+
+            # Create padder
+            padder = padding.PKCS7(128).padder()
+            padded_data = padder.update(plaintext_bytes) + padder.finalize()
+
+            # Create cipher
+            cipher = Cipher(algorithms.AES(key_bytes), modes.CBC(iv_bytes), backend=default_backend())
+
+            # Encrypt
+            encryptor = cipher.encryptor()
+            encrypted_bytes = encryptor.update(padded_data) + encryptor.finalize()
+
+            # Encode to base64
+            return base64.b64encode(encrypted_bytes).decode("utf-8")
+
+        except Exception as e:
+            raise Exception(f"Encryption failed: {str(e)}")
+
+    def encryption_by_aes(self, text: str) -> str:
+        """Encrypt text using AES with class-level key and IV
+
+        Args:
+            text (str): Text to encrypt
+
+        Returns:
+            str: Encrypted text or None if encryption fails
+
+        """
+        try:
+            # Perform encryption
+            encrypted = self.encrypt(text, self.AES_PASW, self.IV)
+
+            return encrypted
+
+        except Exception as e:
+            _LOGGER.error(f"Encryption failed: {str(e)}")
+            return None
+
+    def encrypt_by_public_key(self) -> Optional[str]:
+        """Encrypt data using RSA public key.
+
+        Args:
+
+        Returns:
+            Optional[str]: Base64 encoded encrypted data or None if encryption fails
+
+        """
+
+        data = f"{self.AES_PASW},{self.IV}"
+
+        if not self._public_key:
+            _LOGGER.error("Public key not initialized")
+            return None
+
+        try:
+            # Convert input string to bytes
+            data_bytes = data.encode("utf-8")
+
+            # Encrypt the data padding.PKCS7(128).padder()
+            encrypted_bytes = self._public_key.encrypt(data_bytes, rsa_padding.PKCS1v15())
+
+            # Convert to base64 string
+            encrypted_str = base64.b64encode(encrypted_bytes).decode("utf-8")
+            _LOGGER.debug("Data encrypted successfully")
+
+            return encrypted_str
+
+        except Exception as err:
+            _LOGGER.error("Encryption failed: %s", str(err))
+            return None
+
+    @staticmethod
+    def get_random_string(length: int) -> str:
+        """Generate a random string of specified length using alphanumeric characters.
+
+        Args:
+            length (int): The desired length of the random string
+
+        Returns:
+            str: A random alphanumeric string of specified length
+
+        Raises:
+            ValueError: If length is less than 1
+
+        """
+        if length < 1:
+            raise ValueError("Length must be positive")
+
+        charset = string.ascii_letters + string.digits
+        return "".join(secrets.choice(charset) for _ in range(length))
+
+    @staticmethod
+    def get_random_int(length: int) -> str:
+        """Generate a random string of specified length containing only digits.
+
+        Args:
+            length (int): The desired length of the random number string
+
+        Returns:
+            str: A string of random digits of specified length
+
+        Raises:
+            ValueError: If length is less than 1
+
+        """
+        if length < 1:
+            raise ValueError("Length must be positive")
+
+        return "".join(secrets.choice(string.digits) for _ in range(length))
+
+    @staticmethod
+    def get_aes_key() -> str:
+        """Generate a random AES key of 16 characters using alphanumeric characters.
+        Matches Java implementation behavior.
+
+        Returns:
+            str: A 16-character random string for AES key
+
+        """
+        return EncryptionUtils.get_random_string(16)
+
+    @staticmethod
+    def get_iv() -> str:
+        """Generate a random initialization vector of 16 digits.
+        Matches Java implementation behavior.
+
+        Returns:
+            str: A 16-digit random string for initialization vector
+
+        """
+        return EncryptionUtils.get_random_int(16)

pymammotion/http/http.py

@@ -10,17 +10,25 @@ from pymammotion.const import (
     MAMMOTION_CLIENT_SECRET,
     MAMMOTION_DOMAIN,
 )
+from pymammotion.http.encryption import EncryptionUtils
 from pymammotion.http.model.http import ErrorInfo, LoginResponseData, Response
 
 
 class MammotionHTTP:
-    def __init__(self, response: Response) -> None:
+    def __init__(self) -> None:
+        self.code = None
+        self.msg = None
+        self.response: Response | None = None
+        self.login_info: LoginResponseData | None = None
         self._headers = {"User-Agent": "okhttp/3.14.9", "App-Version": "google Pixel 2 XL taimen-Android 11,1.11.332"}
-        self.login_info = LoginResponseData.from_dict(response.data) if response.data else None
-        self._headers["Authorization"] = f"Bearer {self.login_info.access_token}" if response.data else None
-        self.response = response
-        self.msg = response.msg
-        self.code = response.code
+        self.encryption_utils = EncryptionUtils()
+
+    @staticmethod
+    def generate_headers(token: str) -> dict:
+        return {"Authorization": f"Bearer {token}"}
+
+    async def login_by_email(self, email: str, password: str) -> Response[LoginResponseData]:
+        return await self.login(email, password)
 
     async def get_all_error_codes(self) -> dict[str, ErrorInfo]:
         async with ClientSession(MAMMOTION_API_DOMAIN) as session:
@@ -36,15 +44,55 @@ class MammotionHTTP:
                     codes[error_info.code] = error_info
                 return codes
 
-    async def oauth_check(self) -> None:
+    async def oauth_check(self) -> Response:
         """Check if token is valid.
 
         Returns 401 if token is invalid. We then need to re-authenticate, can try to refresh token first
         """
         async with ClientSession(MAMMOTION_API_DOMAIN) as session:
-            async with session.post("/user-server/v1/user/oauth/check") as resp:
+            async with session.post("/user-server/v1/user/oauth/check", headers=self._headers) as resp:
+                data = await resp.json()
+                return Response.from_dict(data)
+
+    async def pair_devices_mqtt(self, mower_name: str, rtk_name: str) -> Response:
+        async with ClientSession(MAMMOTION_API_DOMAIN) as session:
+            async with session.post(
+                "/device-server/v1/iot/device/pairing",
+                headers=self._headers,
+                json={"mowerName": mower_name, "rtkName": rtk_name},
+            ) as resp:
+                data = await resp.json()
+                if data.get("status") == 200:
+                    print(data)
+                    return Response.from_dict(data)
+                else:
+                    print(data)
+
+    async def unpair_devices_mqtt(self, mower_name: str, rtk_name: str) -> Response:
+        async with ClientSession(MAMMOTION_API_DOMAIN) as session:
+            async with session.post(
+                "/device-server/v1/iot/device/unpairing",
+                headers=self._headers,
+                json={"mowerName": mower_name, "rtkName": rtk_name},
+            ) as resp:
+                data = await resp.json()
+                if data.get("status") == 200:
+                    print(data)
+                    return Response.from_dict(data)
+                else:
+                    print(data)
+
+    async def net_rtk_enable(self, device_id: str) -> Response:
+        async with ClientSession(MAMMOTION_API_DOMAIN) as session:
+            async with session.post(
+                "/device-server/v1/iot/net-rtk/enable", headers=self._headers, json={"deviceId": device_id}
+            ) as resp:
                 data = await resp.json()
-                response = Response.from_dict(data)
+                if data.get("status") == 200:
+                    print(data)
+                    return Response.from_dict(data)
+                else:
+                    print(data)
 
     async def get_stream_subscription(self, iot_id: str) -> Response[StreamSubscriptionResponse]:
         """Get agora.io data for view camera stream"""
@@ -63,27 +111,37 @@ class MammotionHTTP:
                 # Assuming the data format matches the expected structure
                 return Response[StreamSubscriptionResponse].from_dict(data)
 
-    @classmethod
-    async def login(cls, session: ClientSession, username: str, password: str) -> Response[LoginResponseData]:
-        async with session.post(
-            "/oauth/token",
-            headers={"User-Agent": "okhttp/3.14.9", "App-Version": "google Pixel 2 XL taimen-Android 11,1.11.332"},
-            params=dict(
-                username=username,
-                password=password,
-                client_id=MAMMOTION_CLIENT_ID,
-                client_secret=MAMMOTION_CLIENT_SECRET,
-                grant_type="password",
-            ),
-        ) as resp:
-            data = await resp.json()
-            response = Response[LoginResponseData].from_dict(data)
-            # TODO catch errors from mismatch user / password elsewhere
-            # Assuming the data format matches the expected structure
-            return response
-
-
-async def connect_http(username: str, password: str) -> MammotionHTTP:
-    async with ClientSession(MAMMOTION_DOMAIN) as session:
-        login_response = await MammotionHTTP.login(session, username, password)
-        return MammotionHTTP(login_response)
+    async def login(self, username: str, password: str) -> Response[LoginResponseData]:
+        async with ClientSession(MAMMOTION_DOMAIN) as session:
+            async with session.post(
+                "/oauth/token",
+                headers={
+                    "User-Agent": "okhttp/3.14.9",
+                    "App-Version": "google Pixel 2 XL taimen-Android 11,1.11.332",
+                    "Encrypt-Key": self.encryption_utils.encrypt_by_public_key(),
+                    "Decrypt-Type": "3",
+                    "Ec-Version": "v1",
+                },
+                params=dict(
+                    username=self.encryption_utils.encryption_by_aes(username),
+                    password=self.encryption_utils.encryption_by_aes(password),
+                    client_id=self.encryption_utils.encryption_by_aes(MAMMOTION_CLIENT_ID),
+                    client_secret=self.encryption_utils.encryption_by_aes(MAMMOTION_CLIENT_SECRET),
+                    grant_type=self.encryption_utils.encryption_by_aes("password"),
+                ),
+            ) as resp:
+                if resp.status != 200:
+                    print(resp.json())
+                    return Response.from_dict({"status": resp.status, "msg": "Login failed"})
+                data = await resp.json()
+                login_response = Response[LoginResponseData].from_dict(data)
+                self.login_info = LoginResponseData.from_dict(login_response.data)
+                self._headers["Authorization"] = (
+                    f"Bearer {self.login_info.access_token}" if login_response.data else None
+                )
+                self.response = login_response
+                self.msg = login_response.msg
+                self.code = login_response.code
+                # TODO catch errors from mismatch user / password elsewhere
+                # Assuming the data format matches the expected structure
+                return login_response

pymammotion/mammotion/devices/base.py

@@ -8,6 +8,7 @@ import betterproto
 from pymammotion.aliyun.model.dev_by_account_response import Device
 from pymammotion.data.model import RegionData
 from pymammotion.data.model.device import MowingDevice
+from pymammotion.data.model.raw_data import RawMowerData
 from pymammotion.data.state_manager import StateManager
 from pymammotion.proto.luba_msg import LubaMsg
 from pymammotion.proto.mctrl_nav import NavGetCommDataAck, NavGetHashListAck, SvgMessageAckT
@@ -36,8 +37,9 @@ class MammotionBaseDevice:
     def __init__(self, state_manager: StateManager, cloud_device: Device | None = None) -> None:
         """Initialize MammotionBaseDevice."""
         self.loop = asyncio.get_event_loop()
-        self._raw_data = LubaMsg().to_dict(casing=betterproto.Casing.SNAKE)
         self._state_manager = state_manager
+        self._raw_data = dict()
+        self._raw_mower_data: RawMowerData = RawMowerData()
         self._state_manager.gethash_ack_callback = self.datahash_response
         self._state_manager.get_commondata_ack_callback = self.commdata_response
         self._notify_future: asyncio.Future[bytes] | None = None
@@ -108,7 +110,7 @@ class MammotionBaseDevice:
             case "ota":
                 self._update_ota_data(tmp_msg)
 
-        self.mower.update_raw(self._raw_data)
+        self._raw_mower_data.update_raw(self._raw_data)
 
     def _update_nav_data(self, tmp_msg) -> None:
         """Update navigation data."""

pymammotion/mammotion/devices/mammotion.py

@@ -98,7 +98,7 @@ class MammotionMixedDeviceManager:
         return self._ble_device is not None
 
 
-class MammotionDevices:
+class MammotionDeviceManager:
     devices: dict[str, MammotionMixedDeviceManager] = {}
 
     def add_device(self, mammotion_device: MammotionMixedDeviceManager) -> None:
@@ -151,7 +151,7 @@ async def create_devices(
 class Mammotion:
     """Represents a Mammotion account and its devices."""
 
-    devices = MammotionDevices()
+    device_manager = MammotionDeviceManager()
     mqtt_list: dict[str, MammotionCloud] = dict()
 
     _instance: Mammotion = None