Merge pull request #34 from bretterer/vehicle-control

Add method to get OTA details and include 2FA info in get user info request
bretterer · Aug 10, 2023 · a9e41fa · a9e41fa
2 parents e6959f5 + e01d211
commit a9e41fa
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -13,7 +13,7 @@ packages = [
 python = "^3.9"
 aiohttp = ">=3.0.0"
 cryptography = "^41.0.1"
-backports-strenum = { version = "^1.2.4", python = "<3.11" } 
+backports-strenum = { version = "^1.2.4", python = "<3.11" }
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.1.2"

diff --git a/src/rivian/rivian.py b/src/rivian/rivian.py
@@ -285,10 +285,11 @@ async def get_user_information(
 
         vehicles_fragment = "vehicles { id vin name vas { __typename vasVehicleId vehiclePublicKey } roles state createdAt updatedAt vehicle { __typename id vin modelYear make model expectedBuildDate plannedBuildDate expectedGeneralAssemblyStartDate actualGeneralAssemblyDate vehicleState { supportedFeatures { __typename name status } } } }"
         phones_fragment = "enrolledPhones { __typename vas { __typename vasPhoneId publicKey } enrolled { __typename deviceType deviceName vehicleId identityId shortName } }"
+        _2fa_fragment = "registrationChannels { type }"
 
         graphql_json = {
             "operationName": "getUserInfo",
-            "query": f"query getUserInfo {{ currentUser {{ __typename id {vehicles_fragment} {phones_fragment if include_phones else ''} }} }}",
+            "query": f"query getUserInfo {{ currentUser {{ __typename id {vehicles_fragment} {_2fa_fragment} {phones_fragment if include_phones else ''} }} }}",
             "variables": None,
         }
 
@@ -394,6 +395,24 @@ async def get_vehicle_state(
 
         return await self.__graphql_query(headers, url, graphql_json)
 
+    async def get_vehicle_ota_update_details(self, vehicle_id: str) -> ClientResponse:
+        """Get vehicle OTA update details."""
+        url = GRAPHQL_GATEWAY
+        headers = BASE_HEADERS | {
+            "A-Sess": self._app_session_token,
+            "U-Sess": self._user_session_token,
+        }
+
+        graphql_query = "query getOTAUpdateDetails($vehicleId:String!){getVehicle(id:$vehicleId){availableOTAUpdateDetails{url version locale}currentOTAUpdateDetails{url version locale}}}"
+
+        graphql_json = {
+            "operationName": "getOTAUpdateDetails",
+            "query": graphql_query,
+            "variables": {"vehicleId": vehicle_id},
+        }
+
+        return await self.__graphql_query(headers, url, graphql_json)
+
     async def get_live_charging_session(
         self, vin: str, properties: set[str] | None = None
     ) -> ClientResponse: