refactoring, also hopefully fixed the connection dying issue

apns.py

@@ -73,6 +73,10 @@ class IncomingQueue:
                 # We have the lock, so we can safely remove it
                 self.queue.remove(found)
             return found
+        
+    def remove_all(self, id):
+        with self.lock:
+            self.queue = [i for i in self.queue if i[0] != id]
 
     def wait_pop_find(self, finder, delay=0.1):
         found = None
@@ -92,42 +96,16 @@ class APNSConnection:
 
     def _queue_filler(self):
         while True and not self.sock.closed:
-            # print(self.sock.closed)
-            # print("QUEUE: Waiting for payload...")
-            # self.sock.read(1)
-            # print("QUEUE: Got payload?")
             payload = _deserialize_payload(self.sock)
-            # print("QUEUE: Got payload?")
 
             if payload is not None:
-                # print("QUEUE: Received payload: " + str(payload))
+                # Automatically ACK incoming notifications to prevent APNs from getting mad at us
-                # print("QUEUE: Received payload type: " + hex(payload[0]))
+                if payload[0] == 0x0A:
+                    logger.debug("Sending automatic ACK")
+                    self._send_ack(_get_field(payload[1], 4))
                 logger.debug(f"Received payload: {payload}")
                 self.incoming_queue.append(payload)
-        # print("QUEUE: Thread ended")
+                logger.debug(f"Queue length: {len(self.incoming_queue)}")
-
-    # def _pop_by_id(self, id: int) -> tuple[int, list[tuple[int, bytes]]] | None:
-    #     def finder(item):
-    #         return item[0] == id
-    #     return self.incoming_queue.find(finder)
-    #     # print("QUEUE: Looking for id " + str(id) + " in " + str(self.incoming_queue))
-    #     #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]]]:
-    #     found = None
-    #     while found is None:
-    #         found = self._pop_by_id(id)
-    #         if found is None:
-    #             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
@@ -237,6 +215,9 @@ class APNSConnection:
     def keep_alive(self):
         logger.debug("Sending keep alive message")
         self.sock.write(_serialize_payload(0x0C, []))
+        # Remove any keep alive responses we have or missed
+        self.incoming_queue.remove_all(0x0D)
+        
 
     def _send_ack(self, id: bytes):
         logger.debug(f"Sending ACK for message {id}")

demo.py

@@ -24,7 +24,7 @@ logging.basicConfig(
 logging.getLogger("urllib3").setLevel(logging.WARNING)
 logging.getLogger("jelly").setLevel(logging.INFO)
 logging.getLogger("nac").setLevel(logging.INFO)
-logging.getLogger("apns").setLevel(logging.INFO)
+logging.getLogger("apns").setLevel(logging.DEBUG)
 logging.getLogger("albert").setLevel(logging.INFO)
 logging.getLogger("ids").setLevel(logging.DEBUG)
 logging.getLogger("bags").setLevel(logging.DEBUG)
@@ -89,7 +89,7 @@ logging.info("Waiting for incoming messages...")
 
 def keepalive():
     while True:
-        time.sleep(5)
+        time.sleep(300)
         conn.keep_alive()
 
 
@@ -122,6 +122,10 @@ with open("config.json", "w") as f:
 import imessage
 im = imessage.iMessageUser(conn, user)
 
+#import time
+#time.sleep(4)
+#onn._send_ack(b'\t-\x97\x96')
 while True:
     msg = im.receive()
-    print(f"Got message {msg}")
+    if msg is not None:
+        print(f"Got message {msg}")

ids/query.py

@@ -61,6 +61,8 @@ def lookup(
     resp = plistlib.loads(resp)
     resp = gzip.decompress(resp["b"])
     resp = plistlib.loads(resp)
+    # Acknowledge the message
+    #conn._send_ack(apns._get_field(payload[1], 4))
 
     if resp['status'] != 0:
         raise Exception(f'Query failed: {resp}')

imessage.py

@@ -17,11 +17,67 @@ from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
 import gzip
 
+from hashlib import sha1
 import logging
 logger = logging.getLogger("imessage")
 
 NORMAL_NONCE = b"\x00" * 15 + b"\x01"
 
+class BalloonBody:
+    def __init__(self, type: str, data: bytes):
+        self.type = type
+        self.data = data
+
+        # TODO : Register handlers based on type id
+
+class iMessage:
+    text: str
+    xml: str | None = None
+    participants: list[str]
+    sender: str
+    id: str
+    group_id: str
+    body: BalloonBody | None = None
+
+    _raw: dict | None = None
+
+    def from_raw(message: dict) -> 'iMessage':
+        self = iMessage()
+
+        self._raw = message
+
+        self.text = message.get('t')
+        self.xml = message.get('x')
+        self.participants = message.get('p', [])
+        if self.participants != []:
+            self.sender = self.participants[-1]
+        else:
+            self.sender = None
+
+        self.id = message.get('r')
+        self.group_id = message.get('gid')
+
+        if 'bid' in message:
+            # This is a message extension body
+            self.body = BalloonBody(message['bid'], message['b'])
+
+        return self
+
+    def to_raw(self) -> dict:
+        return {
+            "t": self.text,
+            "x": self.xml,
+            "p": self.participants,
+            "r": self.id,
+            "gid": self.group_id,
+        }
+    
+    def __str__(self):
+        if self._raw is not None:
+            return str(self._raw)
+        else:
+            return f"iMessage({self.text} from {self.sender})"
+
 class iMessageUser:
 
     def __init__(self, connection: apns.APNSConnection, user: ids.IDSUser):
@@ -31,21 +87,27 @@ class iMessageUser:
     def _get_raw_message(self):
         """
         Returns a raw APNs message corresponding to the next conforming notification in the queue
+        Returns None if no conforming notification is found
         """
         def check_response(x):
             if x[0] != 0x0A:
                 return False
+            if apns._get_field(x[1], 2) != sha1("com.apple.madrid".encode()).digest():
+                return False
             resp_body = apns._get_field(x[1], 3)
             if resp_body is None:
+                #logger.debug("Rejecting madrid message with no body")
                 return False
             resp_body = plistlib.loads(resp_body)
             if "P" not in resp_body:
+                #logger.debug(f"Rejecting madrid message with no payload : {resp_body}")
                 return False
             return True
         
-        payload = self.connection.incoming_queue.wait_pop_find(check_response)
+        payload = self.connection.incoming_queue.pop_find(check_response)
+        if payload is None:
+            return None
         id = apns._get_field(payload[1], 4)
-        self.connection._send_ack(id)
 
         return payload
 
@@ -123,8 +185,13 @@ class iMessageUser:
         except:
             return False
 
-    def receive(self) -> 'iMessage':
+    def receive(self) -> iMessage | None:
+        """
+        Will return the next iMessage in the queue, or None if there are no messages
+        """
         raw = self._get_raw_message()
+        if raw is None:
+            return None
         body = apns._get_field(raw[1], 3)
         body = plistlib.loads(body)
         payload = body["P"]
@@ -137,69 +204,5 @@ class iMessageUser:
             return self.receive() # Call again to get the next message
         return iMessage.from_raw(decrypted)
     
-    def send(self, message: dict):
+    def send(self, message: iMessage):
-        logger.error(f"Sending {message}")
+        logger.error(f"Sending {message}")
-
-    def receive_message(self) -> str:
-        pass
-
-    def send_message(self, message: str, to: str):
-        pass
-
-class ExtendedBody:
-    def __init__(self, type: str, data: bytes):
-        self.type = type
-        self.data = data
-
-        # TODO : Register handlers based on type id
-
-class iMessage:
-    text: str
-    xml: str | None = None
-    participants: list[str]
-    sender: str
-    id: str
-    group_id: str
-    body: ExtendedBody | None = None
-
-    _raw: dict | None = None
-
-    def from_raw(message: dict) -> 'iMessage':
-        self = iMessage()
-
-        self._raw = message
-
-        self.text = message.get('t')
-        self.xml = message.get('x')
-        self.participants = message.get('p', [])
-        if self.participants != []:
-            self.sender = self.participants[-1]
-        else:
-            self.sender = None
-
-        self.id = message.get('r')
-        self.group_id = message.get('gid')
-
-        if 'bid' in message:
-            # This is a message extension body
-            self.body = ExtendedBody(message['bid'], message['b'])
-
-        return self
-
-    def to_raw(self) -> dict:
-        return {
-            "t": self.text,
-            "x": self.xml,
-            "p": self.participants,
-            "r": self.id,
-            "gid": self.group_id,
-        }
-    
-    def __str__(self):
-        if self._raw is not None:
-            return str(self._raw)
-        else:
-            return f"iMessage({self.text} from {self.sender})"
-
-        
-