more refactoring

bags.py

@@ -4,8 +4,13 @@ import requests
 import logging
 logger = logging.getLogger("bags")
 
-
+OLD_APNS_BAG = None
 def apns_init_bag_old():
+    global OLD_APNS_BAG
+
+    if OLD_APNS_BAG is not None:
+        return OLD_APNS_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")
@@ -15,11 +20,19 @@ def apns_init_bag_old():
 
     logger.debug("Received APNs old-style init bag")
 
+    OLD_APNS_BAG = bag
+
     return bag
 
 
 # This is the same as the above, but the response has a signature which we unwrap
+APNS_BAG = None
 def apns_init_bag():
+    global APNS_BAG
+
+    if APNS_BAG is not None:
+        return APNS_BAG
+    
     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")
@@ -29,10 +42,18 @@ def apns_init_bag():
 
     logger.debug("Received APNs new init bag")
 
+    APNS_BAG = bag
+
     return bag
 
 
+IDS_BAG = None
 def ids_bag():
+    global IDS_BAG
+
+    if IDS_BAG is not None:
+        return IDS_BAG
+    
     r = requests.get(
         "https://init.ess.apple.com/WebObjects/VCInit.woa/wa/getBag?ix=3", verify=False
     )
@@ -46,6 +67,8 @@ def ids_bag():
 
     logger.debug("Recieved IDS bag")
 
+    IDS_BAG = bag
+
     return bag
 
 

demo.py

@@ -124,109 +124,4 @@ im = imessage.iMessageUser(conn, user)
 
 while True:
     msg = im.receive()
-    print(f"Got message {msg['t']}")
-    
-# user_rsa_key = ids._helpers.parse_key(user.encryption_identity.encryption_key)
-# NORMAL_NONCE = b"\x00" * 15 + b"\x01"
-
-# def decrypt(payload, sender_token, rsa_key: rsa.RSAPrivateKey = user_rsa_key):
-#     """
-#     iMessage payload format:
-#     0x00 - ?
-#     0x01-0x02 - length of payload
-#     0x03-0xA0 - RSA encrypted payload portion
-#         0x00-0x0F - AES key
-#         0x0F-0x?? - AES encrypted payload portion
-#     0xA1-0xlength of payload+3 - AES encrypted payload portion
-#     0xlength of payload+3 - length of signature
-#     0xLEN+4-0xLEN+4+length of signature - signature
-#     """
-#     from io import BytesIO
-
-
-#     payload = BytesIO(payload)
-#     tag = payload.read(1)
-#     length = int.from_bytes(payload.read(2), "big")
-#     body = payload.read(length)
-#     body_io = BytesIO(body)
-#     rsa_body = rsa_key.decrypt(
-#         body_io.read(160),
-#         padding.OAEP(
-#             mgf=padding.MGF1(algorithm=hashes.SHA1()),
-#             algorithm=hashes.SHA1(),
-#             label=None,
-#         ),
-#     )
-
-#     cipher = Cipher(algorithms.AES(rsa_body[:16]), modes.CTR(NORMAL_NONCE))
-#     decrypted = cipher.decryptor().update(rsa_body[16:] + body_io.read())
-
-#     # Try to gzip decompress the payload
-#     try:
-#         decrypted = gzip.decompress(decrypted)
-#     except:
-#         logging.debug("Failed to decompress payload")
-#         pass
-
-#     decrypted = plistlib.loads(decrypted)
-
-#     signature_len = payload.read(1)[0]
-#     signature = payload.read(signature_len)
-#     #logging.info(f"Signature: {signature}")
-#     #logging.info(f"Decrypted: {decrypted}")
-
-#     # Verify the signature
-#     sender = decrypted["p"][-1]
-#     # Lookup the public key for the sender
-#     lookup = user.lookup([sender])[sender]
-#     #logging.debug(f"Lookup: {lookup}")
-#     sender = None
-#     for identity in lookup['identities']:
-#         if identity['push-token'] == sender_token:
-#             sender = identity
-#             break
-    
-#     if sender is None:
-#         logging.error(f"Failed to find identity for {sender_token}")
-
-#     identity_keys = sender['client-data']['public-message-identity-key']
-#     identity_keys = ids.identity.IDSIdentity.decode(identity_keys)
-
-#     sender_ec_key = ids._helpers.parse_key(identity_keys.signing_public_key)
-
-#     from cryptography.hazmat.primitives.asymmetric import ec
-#     #logging.debug(f"Verifying signature {signature} with key {sender_ec_key.public_numbers()} and data {body}")
-#     # Verify the signature (will throw an exception if it fails)
-#     sender_ec_key.verify(
-#         signature,
-#         body,
-#         ec.ECDSA(hashes.SHA1()),
-#     )
-
-#     return decrypted
-
-
-# while True:
-
-#     def check_response(x):
-#         if x[0] != 0x0A:
-#             return False
-#         resp_body = apns._get_field(x[1], 3)
-#         if resp_body is None:
-#             return False
-#         resp_body = plistlib.loads(resp_body)
-#         if "P" not in resp_body:
-#             return False
-#         return True
-
-#     payload = conn.incoming_queue.wait_pop_find(check_response)
-#     resp_body = apns._get_field(payload[1], 3)
-#     id = apns._get_field(payload[1], 4)
-#     conn._send_ack(id)
-#     resp_body = plistlib.loads(resp_body)
-#     # logging.info(f"Got response: {resp_body}")
-#     logging.debug(f"Got message: {resp_body}")
-#     token = resp_body['t']
-#     payload = resp_body["P"]
-#     payload = decrypt(payload, token)
-#     logging.info(f"Got message: {payload['t']} from {payload['p'][1]}")
+    print(f"Got message {msg}")

imessage.py

@@ -17,6 +17,9 @@ from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
 import gzip
 
+import logging
+logger = logging.getLogger("imessage")
+
 NORMAL_NONCE = b"\x00" * 15 + b"\x01"
 
 class iMessageUser:
@@ -120,12 +123,83 @@ class iMessageUser:
         except:
             return False
 
-    def receive(self) -> dict:
+    def receive(self) -> 'iMessage':
         raw = self._get_raw_message()
         body = apns._get_field(raw[1], 3)
         body = plistlib.loads(body)
         payload = body["P"]
         decrypted = self._decrypt_payload(payload)
-        if not self._verify_payload(payload, decrypted["p"][-1], body["t"]):
-            raise Exception("Failed to verify payload")
-        return decrypted
+        if "p" in decrypted:
+            if not self._verify_payload(payload, decrypted["p"][-1], body["t"]):
+                raise Exception("Failed to verify payload")
+        else:
+            logger.warning("Unable to verify, couldn't determine sender! Dropping message! (TODO work out a way to verify these anyway)")
+            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})"
+
+        
+