refactoring, also hopefully fixed the connection dying issue

This commit is contained in:
JJTech0130 2023-07-27 20:03:39 -04:00
parent 3f1777a2c4
commit f6f2c304af
No known key found for this signature in database
GPG key ID: 23C92EBCCF8F93D6
4 changed files with 93 additions and 103 deletions

43
apns.py
View file

@ -74,6 +74,10 @@ class IncomingQueue:
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
while found is 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}")

10
demo.py
View file

@ -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}")

View file

@ -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}')

View file

@ -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}")
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})"