pypush-plus-plus/imessage.py

581 lines
19 KiB
Python
Raw Permalink Normal View History

2023-07-27 15:04:57 +00:00
# LOW LEVEL imessage function, decryption etc
# Don't handle APNS etc, accept it already setup
## HAVE ANOTHER FILE TO SETUP EVERYTHING AUTOMATICALLY, etc
# JSON parsing of keys, don't pass around strs??
2023-07-31 23:38:28 +00:00
import base64
2023-07-31 17:03:45 +00:00
import gzip
import logging
2023-07-27 15:52:20 +00:00
import plistlib
2023-07-31 17:03:45 +00:00
import random
from typing import Union
2023-07-31 17:03:45 +00:00
import uuid
from dataclasses import dataclass, field
from hashlib import sha1, sha256
2023-07-27 15:52:20 +00:00
from io import BytesIO
from cryptography.hazmat.primitives import hashes
2023-07-31 17:03:45 +00:00
from cryptography.hazmat.primitives.asymmetric import ec, padding
2023-07-27 15:52:20 +00:00
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
2023-07-31 23:38:28 +00:00
from xml.etree import ElementTree
2023-07-31 17:03:45 +00:00
import apns
import ids
2023-07-27 15:52:20 +00:00
2023-07-27 21:34:38 +00:00
logger = logging.getLogger("imessage")
2023-07-31 17:03:45 +00:00
NORMAL_NONCE = b"\x00" * 15 + b"\x01" # This is always used as the AES nonce
2024-01-30 17:31:28 +00:00
bundled_payloads = []
2023-07-27 15:52:20 +00:00
class BalloonBody:
2023-07-31 17:03:45 +00:00
"""Represents the special parts of message extensions etc."""
def __init__(self, type: str, data: bytes):
self.type = type
self.data = data
# TODO : Register handlers based on type id
2023-07-31 17:03:45 +00:00
2023-07-31 23:38:28 +00:00
class AttachmentFile:
def data(self) -> bytes:
raise NotImplementedError()
@dataclass
class MMCSFile(AttachmentFile):
url: Union[str, None] = None
size: Union[int, None] = None
owner: Union[str, None] = None
signature: Union[bytes, None] = None
decryption_key: Union[bytes, None] = None
2023-07-31 23:38:28 +00:00
def data(self) -> bytes:
import requests
logger.info(requests.get(
url=self.url,
headers={
"User-Agent": f"IMTransferAgent/900 CFNetwork/596.2.3 Darwin/12.2.0 (x86_64) (Macmini5,1)",
# "MMCS-Url": self.url,
# "MMCS-Signature": str(base64.encodebytes(self.signature)),
# "MMCS-Owner": self.owner
},
).headers)
return b""
@dataclass
class InlineFile(AttachmentFile):
_data: bytes
def data(self) -> bytes:
return self._data
@dataclass
class Attachment:
name: str
mime_type: str
versions: list[AttachmentFile]
def __init__(self, message_raw_content: dict, xml_element: ElementTree.Element):
attrs = xml_element.attrib
self.name = attrs["name"] if "name" in attrs else None
self.mime_type = attrs["mime-type"] if "mime-type" in attrs else None
if "inline-attachment" in attrs:
# just grab the inline attachment !
self.versions = [InlineFile(message_raw_content[attrs["inline-attachment"]])]
else:
# suffer
versions = []
for attribute in attrs:
if attribute.startswith("mmcs") or \
attribute.startswith("decryption-key") or \
attribute.startswith("file-size"):
segments = attribute.split('-')
if segments[-1].isnumeric():
index = int(segments[-1])
attribute_name = segments[:-1]
else:
index = 0
attribute_name = attribute
while index >= len(versions):
versions.append(MMCSFile())
val = attrs[attribute_name]
if attribute_name == "mmcs-url":
versions[index].url = val
elif attribute_name == "mmcs-owner":
versions[index].owner = val
elif attribute_name == "mmcs-signature-hex":
versions[index].signature = base64.b16decode(val)
elif attribute_name == "file-size":
versions[index].size = int(val)
elif attribute_name == "decryption-key":
versions[index].decryption_key = base64.b16decode(val)[1:]
2023-07-31 23:38:28 +00:00
self.versions = versions
def __repr__(self):
return f'<Attachment name="{self.name}" type="{self.mime_type}">'
2023-07-31 17:03:45 +00:00
@dataclass
class iMessage:
2023-07-31 17:03:45 +00:00
"""Represents an iMessage"""
2023-07-28 23:20:32 +00:00
text: str = ""
2023-07-31 17:03:45 +00:00
"""Plain text of message, always required, may be an empty string"""
xml: Union[str, None] = None
2023-07-31 17:03:45 +00:00
"""XML portion of message, may be None"""
participants: list[str] = field(default_factory=list)
"""List of participants in the message, including the sender"""
sender: Union[str, None] = None
2023-07-31 17:03:45 +00:00
"""Sender of the message"""
id: Union[uuid.UUID, None] = None
2023-07-31 17:03:45 +00:00
"""ID of the message, will be randomly generated if not provided"""
group_id: Union[uuid.UUID, None] = None
2023-07-31 17:03:45 +00:00
"""Group ID of the message, will be randomly generated if not provided"""
body: Union[BalloonBody, None] = None
2023-07-31 17:03:45 +00:00
"""BalloonBody, may be None"""
effect: Union[str, None] = None
2023-07-31 18:35:34 +00:00
"""iMessage effect sent with this message, may be None"""
2023-07-28 23:20:32 +00:00
_compressed: bool = True
2023-07-31 17:03:45 +00:00
"""Internal property representing whether the message should be compressed"""
2023-07-28 23:20:32 +00:00
_raw: Union[dict, None] = None
2023-07-31 17:03:45 +00:00
"""Internal property representing the original raw message, may be None"""
2023-07-31 23:38:28 +00:00
def attachments(self) -> list[Attachment]:
if self.xml is not None:
return [Attachment(self._raw, elem) for elem in ElementTree.fromstring(self.xml)[0] if elem.tag == "FILE"]
else:
return []
2023-07-31 17:03:45 +00:00
def sanity_check(self):
"""Corrects any missing fields"""
2023-07-31 23:38:28 +00:00
if self.id is None:
self.id = uuid.uuid4()
2023-07-31 17:03:45 +00:00
if self.group_id is None:
self.group_id = uuid.uuid4()
if self.sender is None:
if len(self.participants) > 1:
self.sender = self.participants[-1]
else:
logger.warning(
"Message has no sender, and only one participant, sanity check failed"
)
return False
2023-07-31 17:03:45 +00:00
if self.sender not in self.participants:
self.participants.append(self.sender)
2023-07-31 17:03:45 +00:00
if self.xml != None:
self._compressed = False # XML is never compressed for some reason
2023-07-31 17:03:45 +00:00
return True
def from_raw(message: bytes, sender: Union[str, None] = None) -> "iMessage":
2023-07-31 17:03:45 +00:00
"""Create an `iMessage` from raw message bytes"""
compressed = False
try:
message = gzip.decompress(message)
compressed = True
except:
pass
2023-07-31 17:03:45 +00:00
message = plistlib.loads(message)
return iMessage(
text=message.get("t", ""),
xml=message.get("x"),
participants=message.get("p", []),
2023-07-31 19:23:09 +00:00
sender=sender if sender is not None else message.get("p", [])[-1] if "p" in message else None,
2023-07-31 23:38:28 +00:00
id=uuid.UUID(message.get("r")) if "r" in message else None,
2023-07-31 17:59:46 +00:00
group_id=uuid.UUID(message.get("gid")) if "gid" in message else None,
2023-07-31 23:46:38 +00:00
body=BalloonBody(message["bid"], message["b"]) if "bid" in message and "b" in message else None,
2023-07-31 18:35:34 +00:00
effect=message["iid"] if "iid" in message else None,
2023-07-31 17:03:45 +00:00
_compressed=compressed,
_raw=message,
)
2023-07-28 23:20:32 +00:00
2023-07-31 17:03:45 +00:00
def to_raw(self) -> bytes:
"""Convert an `iMessage` to raw message bytes"""
if not self.sanity_check():
raise ValueError("Message failed sanity check")
2023-07-31 14:58:01 +00:00
d = {
"t": self.text,
"x": self.xml,
"p": self.participants,
2023-07-31 23:38:28 +00:00
"r": str(self.id).upper(),
2023-07-31 17:03:45 +00:00
"gid": str(self.group_id).upper(),
2023-07-30 21:00:04 +00:00
"pv": 0,
2023-07-31 17:03:45 +00:00
"gv": "8",
"v": "1",
2023-07-31 18:47:27 +00:00
"iid": self.effect
}
2023-07-31 17:03:45 +00:00
2023-07-31 14:58:01 +00:00
# Remove keys that are None
2023-07-31 17:03:45 +00:00
d = {k: v for k, v in d.items() if v is not None}
# Serialize as a plist
d = plistlib.dumps(d, fmt=plistlib.FMT_BINARY)
# Compression
if self._compressed:
d = gzip.compress(d, mtime=0)
return d
2023-07-31 18:35:34 +00:00
def to_string(self) -> str:
message_str = f"[{self.sender}] '{self.text}'"
if self.effect is not None:
message_str += f" with effect [{self.effect}]"
return message_str
2023-07-27 15:04:57 +00:00
class iMessageUser:
2023-07-31 17:03:45 +00:00
"""Represents a logged in and connected iMessage user.
This abstraction should probably be reworked into IDS some time..."""
2023-07-27 15:04:57 +00:00
2023-07-27 15:52:20 +00:00
def __init__(self, connection: apns.APNSConnection, user: ids.IDSUser):
self.connection = connection
self.user = user
2023-07-27 15:04:57 +00:00
2023-07-27 15:52:20 +00:00
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
2023-07-27 15:52:20 +00:00
"""
2023-07-31 17:03:45 +00:00
2024-01-30 17:31:28 +00:00
#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
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:
return False
resp_body = plistlib.loads(resp_body)
if "c" not in resp_body or resp_body["c"] != 255:
return False
return True
2023-07-31 17:03:45 +00:00
payload = self.connection.incoming_queue.pop_find(check_response)
if payload is None:
return None
2023-07-27 15:52:20 +00:00
id = apns._get_field(payload[1], 4)
return payload
2023-07-27 15:04:57 +00:00
2023-07-27 15:52:20 +00:00
def _parse_payload(payload: bytes) -> tuple[bytes, bytes]:
payload = BytesIO(payload)
2023-07-27 15:04:57 +00:00
2023-07-27 15:52:20 +00:00
tag = payload.read(1)
2023-07-31 17:08:57 +00:00
#print("TAG", tag)
2023-07-27 15:52:20 +00:00
body_length = int.from_bytes(payload.read(2), "big")
body = payload.read(body_length)
2023-07-31 17:03:45 +00:00
2023-07-27 15:52:20 +00:00
signature_len = payload.read(1)[0]
signature = payload.read(signature_len)
return (body, signature)
2023-07-31 17:03:45 +00:00
2023-07-28 23:20:32 +00:00
def _construct_payload(body: bytes, signature: bytes) -> bytes:
2023-07-31 17:03:45 +00:00
payload = (
b"\x02"
+ len(body).to_bytes(2, "big")
+ body
+ len(signature).to_bytes(1, "big")
+ signature
)
2023-07-28 23:20:32 +00:00
return payload
2023-07-31 14:58:01 +00:00
def _hash_identity(id: bytes) -> bytes:
iden = ids.identity.IDSIdentity.decode(id)
# TODO: Combine this with serialization code in ids.identity
output = BytesIO()
2023-07-31 17:03:45 +00:00
output.write(b"\x00\x41\x04")
output.write(
ids._helpers.parse_key(iden.signing_public_key)
.public_numbers()
.x.to_bytes(32, "big")
)
output.write(
ids._helpers.parse_key(iden.signing_public_key)
.public_numbers()
.y.to_bytes(32, "big")
)
2023-07-31 14:58:01 +00:00
2023-07-31 17:03:45 +00:00
output.write(b"\x00\xAC")
output.write(b"\x30\x81\xA9")
output.write(b"\x02\x81\xA1")
output.write(
ids._helpers.parse_key(iden.encryption_public_key)
.public_numbers()
.n.to_bytes(161, "big")
)
output.write(b"\x02\x03\x01\x00\x01")
2023-07-31 14:58:01 +00:00
return sha256(output.getvalue()).digest()
2023-07-28 23:20:32 +00:00
2023-07-31 17:03:45 +00:00
def _encrypt_sign_payload(
self, key: ids.identity.IDSIdentity, message: bytes
) -> bytes:
2023-07-28 23:20:32 +00:00
# Generate a random AES key
2023-07-31 14:58:01 +00:00
random_seed = random.randbytes(11)
# Create the HMAC
import hmac
2023-07-31 17:03:45 +00:00
hm = hmac.new(
random_seed,
message
+ b"\x02"
+ iMessageUser._hash_identity(self.user.encryption_identity.encode())
+ iMessageUser._hash_identity(key.encode()),
sha256,
).digest()
2023-07-31 14:58:01 +00:00
aes_key = random_seed + hm[:5]
2023-07-31 17:03:45 +00:00
# print(len(aes_key))
2023-07-28 23:20:32 +00:00
# Encrypt the message with the AES key
cipher = Cipher(algorithms.AES(aes_key), modes.CTR(NORMAL_NONCE))
encrypted = cipher.encryptor().update(message)
# Encrypt the AES key with the public key of the recipient
recipient_key = ids._helpers.parse_key(key.encryption_public_key)
rsa_body = recipient_key.encrypt(
2023-07-31 17:03:45 +00:00
aes_key + encrypted[:100],
padding.OAEP(
mgf=padding.MGF1(algorithm=hashes.SHA1()),
algorithm=hashes.SHA1(),
label=None,
),
2023-07-28 23:20:32 +00:00
)
# Construct the payload
body = rsa_body + encrypted[100:]
2023-07-31 17:03:45 +00:00
sig = ids._helpers.parse_key(self.user.encryption_identity.signing_key).sign(
body, ec.ECDSA(hashes.SHA1())
)
2023-07-28 23:20:32 +00:00
payload = iMessageUser._construct_payload(body, sig)
return payload
2023-07-31 17:03:45 +00:00
2023-07-27 15:52:20 +00:00
def _decrypt_payload(self, payload: bytes) -> dict:
payload = iMessageUser._parse_payload(payload)
body = BytesIO(payload[0])
2023-07-31 17:03:45 +00:00
rsa_body = ids._helpers.parse_key(
self.user.encryption_identity.encryption_key
).decrypt(
2023-07-27 15:52:20 +00:00
body.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.read())
2023-07-31 17:08:57 +00:00
2023-07-31 17:03:45 +00:00
return decrypted
2023-07-27 15:52:20 +00:00
def _verify_payload(self, payload: bytes, sender: str, sender_token: str) -> bool:
# Get the public key for the sender
2023-07-31 17:03:45 +00:00
self._cache_keys([sender])
2023-07-27 15:52:20 +00:00
2023-07-31 17:03:45 +00:00
if not sender_token in self.KEY_CACHE:
logger.warning("Unable to find the public key of the sender, cannot verify")
return False
2023-07-27 15:52:20 +00:00
2023-07-31 17:03:45 +00:00
identity_keys = ids.identity.IDSIdentity.decode(self.KEY_CACHE[sender_token][0])
2023-07-27 15:52:20 +00:00
sender_ec_key = ids._helpers.parse_key(identity_keys.signing_public_key)
payload = iMessageUser._parse_payload(payload)
try:
# Verify the signature (will throw an exception if it fails)
sender_ec_key.verify(
payload[1],
payload[0],
ec.ECDSA(hashes.SHA1()),
)
return True
except:
return False
def receive(self) -> Union[iMessage, None]:
"""
Will return the next iMessage in the queue, or None if there are no messages
"""
2023-07-27 15:52:20 +00:00
raw = self._get_raw_message()
if raw is None:
return None
2023-07-27 15:52:20 +00:00
body = apns._get_field(raw[1], 3)
body = plistlib.loads(body)
2023-07-31 17:08:57 +00:00
#print(f"Got body message {body}")
2023-07-27 15:52:20 +00:00
payload = body["P"]
2023-07-31 17:03:45 +00:00
if not self._verify_payload(payload, body['sP'], body["t"]):
raise Exception("Failed to verify payload")
2023-07-27 15:52:20 +00:00
decrypted = self._decrypt_payload(payload)
2023-07-31 17:03:45 +00:00
2023-07-31 19:23:09 +00:00
return iMessage.from_raw(decrypted, body['sP'])
2023-07-31 17:03:45 +00:00
2023-07-31 20:30:06 +00:00
KEY_CACHE_HANDLE: str = ""
2023-07-31 17:03:45 +00:00
KEY_CACHE: dict[bytes, tuple[bytes, bytes]] = {}
"""Mapping of push token : (public key, session token)"""
USER_CACHE: dict[str, list[bytes]] = {}
"""Mapping of handle : [push tokens]"""
2023-07-29 17:57:20 +00:00
def _cache_keys(self, participants: list[str]):
2023-07-31 20:30:06 +00:00
# Clear the cache if the handle has changed
if self.KEY_CACHE_HANDLE != self.user.current_handle:
self.KEY_CACHE_HANDLE = self.user.current_handle
self.KEY_CACHE = {}
self.USER_CACHE = {}
2023-07-31 17:03:45 +00:00
# Check to see if we have cached the keys for all of the participants
if all([p in self.USER_CACHE for p in participants]):
return
2023-07-29 17:57:20 +00:00
# Look up the public keys for the participants, and cache a token : public key mapping
lookup = self.user.lookup(participants)
for key, participant in lookup.items():
if not key in self.USER_CACHE:
self.USER_CACHE[key] = []
2023-07-31 17:03:45 +00:00
for identity in participant["identities"]:
if not "client-data" in identity:
2023-07-29 17:57:20 +00:00
continue
2023-07-31 17:03:45 +00:00
if not "public-message-identity-key" in identity["client-data"]:
2023-07-29 17:57:20 +00:00
continue
2023-07-31 17:03:45 +00:00
if not "push-token" in identity:
2023-07-29 17:57:20 +00:00
continue
2023-07-31 17:03:45 +00:00
if not "session-token" in identity:
2023-07-29 21:14:25 +00:00
continue
2023-07-29 17:57:20 +00:00
2023-07-31 17:03:45 +00:00
self.USER_CACHE[key].append(identity["push-token"])
2023-07-29 17:57:20 +00:00
2023-07-31 17:03:45 +00:00
# print(identity)
self.KEY_CACHE[identity["push-token"]] = (
identity["client-data"]["public-message-identity-key"],
identity["session-token"],
)
2023-07-29 21:14:25 +00:00
def send(self, message: iMessage):
2023-07-28 23:20:32 +00:00
# Set the sender, if it isn't already
if message.sender is None:
2023-07-31 17:03:45 +00:00
message.sender = self.user.handles[0] # TODO : Which handle to use?
2023-07-28 23:20:32 +00:00
2023-07-31 17:03:45 +00:00
message.sanity_check() # Sanity check MUST be called before caching keys, so that the sender is added to the list of participants
2023-07-29 17:57:20 +00:00
self._cache_keys(message.participants)
# Turn the message into a raw message
raw = message.to_raw()
import base64
2023-07-31 17:03:45 +00:00
2023-07-29 17:57:20 +00:00
for participant in message.participants:
2023-12-01 20:32:36 +00:00
participant = participant.lower()
2023-07-29 17:57:20 +00:00
for push_token in self.USER_CACHE[participant]:
2023-07-31 17:57:50 +00:00
if push_token == self.connection.token:
continue # Don't send to ourselves
2023-07-31 17:03:45 +00:00
identity_keys = ids.identity.IDSIdentity.decode(
self.KEY_CACHE[push_token][0]
)
2023-07-29 17:57:20 +00:00
payload = self._encrypt_sign_payload(identity_keys, raw)
2023-07-31 17:03:45 +00:00
bundled_payloads.append(
{
"tP": participant,
"D": not participant
== message.sender, # TODO: Should this be false sometimes? For self messages?
"sT": self.KEY_CACHE[push_token][1],
"P": payload,
"t": push_token,
}
)
2023-07-30 21:00:04 +00:00
msg_id = random.randbytes(4)
2023-07-29 17:57:20 +00:00
body = {
2023-07-31 17:03:45 +00:00
"fcn": 1,
"c": 100,
"E": "pair",
"ua": "[macOS,13.4.1,22F82,MacBookPro18,3]",
"v": 8,
"i": int.from_bytes(msg_id, "big"),
2023-07-31 23:38:28 +00:00
"U": message.id.bytes,
2023-07-31 17:03:45 +00:00
"dtl": bundled_payloads,
"sP": message.sender,
2023-07-29 17:57:20 +00:00
}
2024-01-30 17:31:28 +00:00
2023-07-29 17:57:20 +00:00
body = plistlib.dumps(body, fmt=plistlib.FMT_BINARY)
2023-07-30 21:00:04 +00:00
self.connection.send_message("com.apple.madrid", body, msg_id)
2023-07-29 17:57:20 +00:00
2023-07-31 17:08:57 +00:00
# This code can check to make sure we got a success response, but waiting for the response is annoying,
# so for now we just YOLO it and assume it worked
2024-01-30 17:31:28 +00:00
#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:
#return False
#resp_body = plistlib.loads(resp_body)
#if "c" not in resp_body or resp_body["c"] != 255:
#return False
#return True
num_recv = 0
while True:
if num_recv == len(bundled_payloads):
break
payload = self.connection.incoming_queue.wait_pop_find(check_response)
if payload is None:
continue
resp_body = apns._get_field(payload[1], 3)
resp_body = plistlib.loads(resp_body)
logger.error(resp_body)
num_recv += 1