From 125c668ea1896093a1f2b346fbd575a1d7ae4d12 Mon Sep 17 00:00:00 2001
From: JJTech0130 <jjtech@jjtech.dev>
Date: Wed, 12 Apr 2023 13:13:47 -0400
Subject: [PATCH] refactor bag lookups

---
 bags.py    | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 config.py  | 25 -------------------------
 ids.py     | 11 ++++-------
 printer.py |  5 ++++-
 4 files changed, 57 insertions(+), 33 deletions(-)
 create mode 100644 bags.py
 delete mode 100644 config.py

diff --git a/bags.py b/bags.py
new file mode 100644
index 0000000..05c0a4d
--- /dev/null
+++ b/bags.py
@@ -0,0 +1,49 @@
+import plistlib
+
+import requests
+
+
+def apns_init_bag():
+    r = requests.get("https://init.push.apple.com/bag", verify=False)
+    if r.status_code != 200:
+        raise Exception("Failed to get APNs init bag")
+
+    # Parse the config as a plist
+    bag = plistlib.loads(r.content)
+
+    return bag
+
+
+# This is the same as the above, but the response has a signature which we unwrap
+def apns_init_bag_2():
+    r = requests.get("http://init-p01st.push.apple.com/bag", verify=False)
+    if r.status_code != 200:
+        raise Exception("Failed to get APNs init bag 2")
+
+    content = plistlib.loads(r.content)
+    bag = plistlib.loads(content["bag"])
+
+    return bag
+
+
+def ids_bag():
+    r = requests.get(
+        "https://init.ess.apple.com/WebObjects/VCInit.woa/wa/getBag?ix=3", verify=False
+    )
+    if r.status_code != 200:
+        raise Exception("Failed to get IDS bag")
+
+    # Parse the config as a plist
+    content = plistlib.loads(r.content)
+    # Load the inner bag
+    bag = plistlib.loads(content["bag"])
+
+    return bag
+
+
+if __name__ == "__main__":
+    # config = get_config()
+    # print(config)
+    # print(apns_init_bag_2())
+    print(apns_init_bag_2() == apns_init_bag())
+    # print(ids_bag())
diff --git a/config.py b/config.py
deleted file mode 100644
index a5afc5c..0000000
--- a/config.py
+++ /dev/null
@@ -1,25 +0,0 @@
-import plistlib
-
-import requests
-
-# CONFIG_URL = "http://init-p01st.push.apple.com/bag"
-CONFIG_URL = "https://init.push.apple.com/bag"
-
-
-def get_config():
-    r = requests.get(CONFIG_URL, verify=False)
-    if r.status_code != 200:
-        raise Exception("Failed to get config")
-
-    # Parse the config as a plist
-    config = plistlib.loads(r.content)
-
-    # Parse the nested "bag" as a plist
-    # config["bag"] = plistlib.loads(config["bag"])
-
-    return config
-
-
-if __name__ == "__main__":
-    config = get_config()
-    print(config)
diff --git a/ids.py b/ids.py
index 0fa39fd..e8ed02a 100644
--- a/ids.py
+++ b/ids.py
@@ -9,6 +9,7 @@ from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import padding
 
 import apns
+import bags
 
 USER_AGENT = "com.apple.madrid-lookup [macOS,13.2.1,22D68,MacBookPro18,3]"
 # NOTE: The push token MUST be registered with the account for self-uri!
@@ -78,18 +79,14 @@ def sign_payload(
     return sig, nonce
 
 
-BAG_KEYS = {
-    "id-query": "https://query.ess.apple.com/WebObjects/QueryService.woa/wa/query"
-}
-
 global_key, global_cert = load_keys()
 
 
-def _send_request(conn: apns.APNSConnection, type: str, body: bytes) -> bytes:
+def _send_request(conn: apns.APNSConnection, bag_key: str, body: bytes) -> bytes:
     body = zlib.compress(body, wbits=16 + zlib.MAX_WBITS)
 
     # Sign the request
-    signature, nonce = sign_payload(global_key, type, "", PUSH_TOKEN, body)
+    signature, nonce = sign_payload(global_key, bag_key, "", PUSH_TOKEN, body)
 
     headers = {
         "x-id-cert": global_cert.replace("-----BEGIN CERTIFICATE-----", "")
@@ -108,7 +105,7 @@ def _send_request(conn: apns.APNSConnection, type: str, body: bytes) -> bytes:
         "U": b"\x16%D\xd5\xcd:D1\xa1\xa7z6\xa9\xe2\xbc\x8f",  # Just random bytes?
         "c": 96,
         "ua": USER_AGENT,
-        "u": BAG_KEYS[type],
+        "u": bags.ids_bag()[bag_key],
         "h": headers,
         "v": 2,
         "b": body,
diff --git a/printer.py b/printer.py
index 2a859b3..1b53b77 100644
--- a/printer.py
+++ b/printer.py
@@ -221,7 +221,7 @@ def pretty_print_payload(
             print(f" {bcolors.FAIL}Madrid{bcolors.ENDC}", end="")
             payload = plistlib.loads(_get_field(payload[1], 3))
             # print(payload)
-            if "cT" in payload:
+            if "cT" in payload and False:
                 # It's HTTP over APNs
                 if "hs" in payload:
                     print(
@@ -248,6 +248,9 @@ def pretty_print_payload(
                     if b"plist" in body:
                         body = plistlib.loads(body)
                     print(f" {bcolors.FAIL}Body{bcolors.ENDC}: {body}", end="")
+            if not "cT" in payload:
+                for key in payload:
+                    print(f" {bcolors.OKBLUE}{key}{bcolors.ENDC}: {payload[key]}")
 
         print()