From f6f2c304af43d191f7e3fd2f7af22b1721246bcb Mon Sep 17 00:00:00 2001 From: JJTech0130 Date: Thu, 27 Jul 2023 20:03:39 -0400 Subject: [PATCH] refactoring, also hopefully fixed the connection dying issue --- apns.py | 43 +++++----------- demo.py | 10 ++-- ids/query.py | 2 + imessage.py | 141 ++++++++++++++++++++++++++------------------------- 4 files changed, 93 insertions(+), 103 deletions(-) diff --git a/apns.py b/apns.py index f4acf60..ef92717 100644 --- a/apns.py +++ b/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)) - # print("QUEUE: Received payload type: " + hex(payload[0])) + # Automatically ACK incoming notifications to prevent APNs from getting mad at us + 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") - - # 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) + logger.debug(f"Queue length: {len(self.incoming_queue)}") 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}") diff --git a/demo.py b/demo.py index 7ac93b5..82915a0 100644 --- a/demo.py +++ b/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}") \ No newline at end of file + if msg is not None: + print(f"Got message {msg}") \ No newline at end of file diff --git a/ids/query.py b/ids/query.py index 0c462eb..875372b 100644 --- a/ids/query.py +++ b/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}') diff --git a/imessage.py b/imessage.py index 5d481fc..9bcee4d 100644 --- a/imessage.py +++ b/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): - 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})" - - - + def send(self, message: iMessage): + logger.error(f"Sending {message}") \ No newline at end of file