implement message correlation

apns.py

@@ -29,13 +29,53 @@ def _connect(private_key: str, cert: str) -> tlslite.TLSConnection:
 
     return sock
 
+class IncomingQueue:
+    def __init__(self):
+        self.queue = []
+        self.lock = threading.Lock()
+    
+    def append(self, item):
+        with self.lock:
+            self.queue.append(item)
+        
+    def pop(self, index):
+        with self.lock:
+            return self.queue.pop(index)
+    
+    def __getitem__(self, index):
+        with self.lock:
+            return self.queue[index]
+        
+    def __len__(self):
+        with self.lock:
+            return len(self.queue)
+        
+    def find(self, finder):
+        with self.lock:
+            return next((i for i in self.queue if finder(i)), None)
+        
+    def pop_find(self, finder):
+        with self.lock:
+            found = next((i for i in self.queue if finder(i)), None)
+            if found is not None:
+                # We have the lock, so we can safely remove it
+                self.queue.remove(found)
+            return found
+        
+    def wait_pop_find(self, finder, delay=0.1):
+        found = None
+        while found is None:
+            found = self.pop_find(finder)
+            if found is None:
+                time.sleep(delay)
+        return found
 
 class APNSConnection:
-    incoming_queue = []
+    incoming_queue = IncomingQueue()
 
     # Sink everything in the queue
     def sink(self):
-        self.incoming_queue = []
+        self.incoming_queue = IncomingQueue()
 
     def _queue_filler(self):
         while True and not self.sock.closed:
@@ -48,22 +88,32 @@ class APNSConnection:
 
             if payload is not None:
                 #print("QUEUE: Received payload: " + str(payload))
+                print("QUEUE: Received payload type: " + hex(payload[0]))
                 self.incoming_queue.append(payload)
         # print("QUEUE: Thread ended")
 
-    def _pop_by_id(self, id: int) -> tuple[int, list[tuple[int, bytes]]] | None:
+    # def _pop_by_id(self, id: int) -> tuple[int, list[tuple[int, bytes]]] | None:
-        # print("QUEUE: Looking for id " + str(id) + " in " + str(self.incoming_queue))
+    #     def finder(item):
-        for i in range(len(self.incoming_queue)):
+    #         return item[0] == id
-            if self.incoming_queue[i][0] == id:
+    #     return self.incoming_queue.find(finder)
-                return self.incoming_queue.pop(i)
+    #     # print("QUEUE: Looking for id " + str(id) + " in " + str(self.incoming_queue))
-        return None
+    #     #for i in range(len(self.incoming_queue)):
+    #     #    if self.incoming_queue[i][0] == id:
+    #     #        return self.incoming_queue.pop(i)
+    #     #return None
 
-    def wait_for_packet(self, id: int) -> tuple[int, list[tuple[int, bytes]]]:
+    # def wait_for_packet(self, id: int) -> tuple[int, list[tuple[int, bytes]]]:
-        payload = self._pop_by_id(id)
+    #     found = None
-        while payload is None:
+    #     while found is None:
-            payload = self._pop_by_id(id)
+    #         found = self._pop_by_id(id)
-            time.sleep(0.1)
+    #         if found is None:
-        return payload
+    #             time.sleep(0.1)
+    #     return found
+    
+    # def find_packet(self, finder) -> 
+    
+    #def replace_packet(self, payload: tuple[int, list[tuple[int, bytes]]]):
+    #    self.incoming_queue.append(payload)
 
     def __init__(self, private_key=None, cert=None):
         # Generate the private key and certificate if they're not provided
@@ -96,7 +146,7 @@ class APNSConnection:
 
         self.sock.write(payload)
 
-        payload = self.wait_for_packet(8)
+        payload = self.incoming_queue.wait_pop_find(lambda i: i[0] == 8)
 
         if (
             payload == None
@@ -141,7 +191,8 @@ class APNSConnection:
 
         self.sock.write(payload)
 
-        payload = self.wait_for_packet(0x0B)
+        # Wait for ACK
+        payload = self.incoming_queue.wait_pop_find(lambda i: i[0] == 0x0B)
 
         if payload[1][0][1] != 0x00.to_bytes(1, 'big'):
             raise Exception("Failed to send message")
@@ -156,6 +207,19 @@ class APNSConnection:
     def keep_alive(self):
         self.sock.write(_serialize_payload(0x0C, []))
 
+    # def _send_ack(self, id: bytes):
+    #     print(f"Sending ACK for message {id}")
+    #     payload = _serialize_payload(0x0B, [(1, self.token), (4, id), (8, b"\x00")])
+    #     self.sock.write(payload)
+    #     #self.sock.write(_serialize_payload(0x0B, [(4, id)])
+    #     #pass
+
+    # def recieve_message(self):
+    #     payload = self.incoming_queue.wait_pop_find(lambda i: i[0] == 0x0A)
+    #     # Send ACK
+    #     self._send_ack(_get_field(payload[1], 4))
+    #     return _get_field(payload[1], 3)
+
     # TODO: Find a way to make this non-blocking
     # def expect_message(self) -> tuple[int, list[tuple[int, bytes]]] | None:
     #   return _deserialize_payload(self.sock)

demo.py

@@ -140,6 +140,7 @@ def lookup(topic:str, users: list[str]):
     print(f"Looking up users {users} for topic {topic}...")
     resp = ids.lookup(conn, CONFIG['username'], ids_keypair, topic, users)
 
+    #print(resp)
     #r = list(resp['results'].values())[0]
     for k, v in resp['results'].items():
         print(f"Result for user {k} topic {topic}:")
@@ -156,20 +157,23 @@ def lookup(topic:str, users: list[str]):
 
 # Hack to make sure that the requests and responses match up
 # This filter MUST contain all the topics you are looking up
-conn.filter(['com.apple.madrid', 'com.apple.private.alloy.facetime.multi', 'com.apple.private.alloy.multiplex1'])
+#conn.filter(['com.apple.madrid', 'com.apple.private.alloy.facetime.multi', 'com.apple.private.alloy.multiplex1', 'com.apple.private.alloy.screensharing'])
-import time
+#import time
-print("...waiting for queued messages... (this is a hack)")
+#print("...waiting for queued messages... (this is a hack)")
-time.sleep(5) # Let the server send us any messages it was holding
+#time.sleep(5) # Let the server send us any messages it was holding
-conn.sink() # Dump the messages
+#conn.sink() # Dump the messages
 
-lookup("com.apple.madrid", ["mailto:jjtech@jjtech.dev"])
+#lookup("com.apple.madrid", ["mailto:jjtech@jjtech.dev"])
-lookup("com.apple.private.alloy.facetime.multi", ["mailto:jjtech@jjtech.dev"])
+#lookup("com.apple.private.alloy.facetime.multi", ["mailto:jjtech@jjtech.dev"])
 
-lookup("com.apple.private.alloy.facetime.multi", ["mailto:user_test2@icloud.com"])
+# lookup("com.apple.private.alloy.facetime.multi", ["mailto:user_test2@icloud.com"])
-lookup("com.apple.madrid", ["mailto:user_test2@icloud.com"])
+# lookup("com.apple.madrid", ["mailto:user_test2@icloud.com"])
 
-lookup("com.apple.private.alloy.multiplex1", ["mailto:user_test2@icloud.com"])
+# lookup("com.apple.private.alloy.multiplex1", ["mailto:user_test2@icloud.com"])
 
+lookup("com.apple.private.alloy.screensharing", ["mailto:user_test2@icloud.com"])
+
+#time.sleep(4)
 # Save config
 with open("config.json", "w") as f:
     json.dump(CONFIG, f, indent=4)

ids.py

@@ -100,9 +100,11 @@ def _send_request(conn: apns.APNSConnection, bag_key: str, topic: str, body: byt
 
     #print(headers)
 
+    msg_id = random.randbytes(16)
+
     req = {
         "cT": "application/x-apple-plist",
-        "U": b"\x16%C\xd5\xcd:D1\xa1\xa7z6\xa9\xe2\xbc\x8f",  # Just random bytes?
+        "U": msg_id,
         "c": 96,
         "ua": USER_AGENT,
         "u": bags.ids_bag()[bag_key],
@@ -112,14 +114,23 @@ def _send_request(conn: apns.APNSConnection, bag_key: str, topic: str, body: byt
     }
 
     conn.send_message(topic, plistlib.dumps(req, fmt=plistlib.FMT_BINARY))
-    resp = conn.wait_for_packet(0x0A)
+    #resp = conn.wait_for_packet(0x0A)
-
-    resp_body = apns._get_field(resp[1], 3)
 
+    def check_response(x):
+        if x[0] != 0x0A:
+            return False
+        resp_body = apns._get_field(x[1], 3)
         if resp_body is None:
-        raise (Exception(f"Got invalid response: {resp}"))
+            return False
+        resp_body = plistlib.loads(resp_body)
+        return resp_body['U'] == msg_id
     
-    return resp_body
+    # Lambda to check if the response is the one we want
+    #conn.incoming_queue.find(check_response)
+    payload = conn.incoming_queue.wait_pop_find(check_response)
+    #conn._send_ack(apns._get_field(payload[1], 4))
+    resp = apns._get_field(payload[1], 3)
+    return plistlib.loads(resp)
 
 
 # Performs an IDS lookup
@@ -132,7 +143,8 @@ def lookup(conn: apns.APNSConnection, self: str, keypair: KeyPair, topic: str, q
     conn.filter([topic])
     query = {"uris": query}
     resp = _send_request(conn, "id-query", topic, plistlib.dumps(query), keypair, self)
-    resp = plistlib.loads(resp)
+    #resp = plistlib.loads(resp)
+    #print(resp)
     resp = gzip.decompress(resp["b"])
     resp = plistlib.loads(resp)
     return resp