pypush-plus-plus/imessage.py

389 lines
13 KiB
Python
Raw 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??
import apns
import ids
2023-07-27 15:52:20 +00:00
import plistlib
from io import BytesIO
from cryptography.hazmat.primitives.asymmetric import ec, padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
import gzip
2023-07-28 23:20:32 +00:00
import uuid
import random
2023-07-29 17:57:20 +00:00
import time
2023-07-27 15:52:20 +00:00
2023-07-31 14:58:01 +00:00
from hashlib import sha1, sha256
2023-07-27 21:34:38 +00:00
import logging
logger = logging.getLogger("imessage")
2023-07-27 15:52:20 +00:00
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:
2023-07-28 23:20:32 +00:00
text: str = ""
xml: str | None = None
2023-07-28 23:20:32 +00:00
participants: list[str] = []
sender: str | None = None
id: str | None = None
group_id: str | None = None
body: BalloonBody | None = None
2023-07-28 23:20:32 +00:00
_compressed: bool = True
_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'])
2023-07-28 23:20:32 +00:00
if 'compressed' in message: # This is a hack, not a real field
self._compressed = message['compressed']
return self
def to_raw(self) -> dict:
2023-07-31 14:58:01 +00:00
d = {
"t": self.text,
"x": self.xml,
"p": self.participants,
"r": self.id,
"gid": self.group_id,
2023-07-28 23:20:32 +00:00
"compressed": self._compressed,
2023-07-30 21:00:04 +00:00
"pv": 0,
2023-07-31 14:58:01 +00:00
"gv": '8',
"v": '1'
}
2023-07-31 14:58:01 +00:00
# Remove keys that are None
return {k: v for k, v in d.items() if v is not None}
def __str__(self):
if self._raw is not None:
return str(self._raw)
else:
return f"iMessage({self.text} from {self.sender})"
2023-07-27 15:04:57 +00:00
class iMessageUser:
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
"""
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
2023-07-27 15:52:20 +00:00
resp_body = apns._get_field(x[1], 3)
if resp_body is None:
#logger.debug("Rejecting madrid message with no body")
2023-07-27 15:52:20 +00:00
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}")
2023-07-27 15:52:20 +00:00
return False
return True
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-28 23:20:32 +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)
signature_len = payload.read(1)[0]
signature = payload.read(signature_len)
return (body, signature)
2023-07-28 23:20:32 +00:00
def _construct_payload(body: bytes, signature: bytes) -> bytes:
payload = b"\x02" + len(body).to_bytes(2, "big") + body + len(signature).to_bytes(1, "big") + signature
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()
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"))
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')
return sha256(output.getvalue()).digest()
2023-07-30 21:00:04 +00:00
def _encrypt_sign_payload(self, key: ids.identity.IDSIdentity, message: dict) -> bytes:
2023-07-28 23:20:32 +00:00
# Dump the message plist
compressed = message.get('compressed', False)
2023-07-30 21:00:04 +00:00
# Remove the compressed flag from the message
#if 'compressed' in message:
# del message['compressed']
m2 = message.copy()
if 'compressed' in m2:
del m2['compressed']
message = plistlib.dumps(m2, fmt=plistlib.FMT_BINARY)
2023-07-28 23:20:32 +00:00
# Compress the message
if compressed:
message = gzip.compress(message, mtime=0)
# Generate a random AES key
2023-07-31 14:58:01 +00:00
random_seed = random.randbytes(11)
# Create the HMAC
import hmac
hm = hmac.new(random_seed, message + b"\x02" + iMessageUser._hash_identity(self.user.encryption_identity.encode()) + iMessageUser._hash_identity(key.encode()), sha256).digest()
aes_key = random_seed + hm[:5]
#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(
aes_key + encrypted[:100],
padding.OAEP(
mgf=padding.MGF1(algorithm=hashes.SHA1()),
algorithm=hashes.SHA1(),
label=None,
),
)
# Construct the payload
body = rsa_body + encrypted[100:]
sig = ids._helpers.parse_key(self.user.encryption_identity.signing_key).sign(body, ec.ECDSA(hashes.SHA1()))
payload = iMessageUser._construct_payload(body, sig)
return payload
2023-07-27 15:52:20 +00:00
def _decrypt_payload(self, payload: bytes) -> dict:
payload = iMessageUser._parse_payload(payload)
body = BytesIO(payload[0])
rsa_body = ids._helpers.parse_key(self.user.encryption_identity.encryption_key).decrypt(
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())
# Try to gzip decompress the payload
2023-07-28 23:20:32 +00:00
compressed = False
2023-07-27 15:52:20 +00:00
try:
decrypted = gzip.decompress(decrypted)
2023-07-28 23:20:32 +00:00
compressed = True
2023-07-27 15:52:20 +00:00
except:
pass
2023-07-28 23:20:32 +00:00
pl = plistlib.loads(decrypted)
pl['compressed'] = compressed # This is a hack so that messages can be re-encrypted with the same compression
return pl
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
lookup = self.user.lookup([sender])[sender]
sender_iden = None
for identity in lookup['identities']:
if identity['push-token'] == sender_token:
sender_iden = identity
break
identity_keys = sender_iden['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)
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) -> 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-29 17:57:20 +00:00
print(f"Got body message {body}")
2023-07-27 15:52:20 +00:00
payload = body["P"]
decrypted = self._decrypt_payload(payload)
2023-07-27 21:34:38 +00:00
if "p" in decrypted:
2023-07-30 21:00:04 +00:00
if not self._verify_payload(payload, decrypted["p"][-1], body["t"]):
raise Exception("Failed to verify payload")
2023-07-27 21:34:38 +00:00
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)
2023-07-29 21:14:25 +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]):
# 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] = []
for identity in participant['identities']:
if not 'client-data' in identity:
continue
if not 'public-message-identity-key' in identity['client-data']:
continue
if not 'push-token' in identity:
continue
2023-07-29 21:14:25 +00:00
if not 'session-token' in identity:
continue
2023-07-29 17:57:20 +00:00
self.USER_CACHE[key].append(identity['push-token'])
2023-07-30 21:00:04 +00:00
#print(identity)
2023-07-29 21:14:25 +00:00
self.KEY_CACHE[identity['push-token']] = (identity['client-data']['public-message-identity-key'], identity['session-token'])
2023-07-29 17:57:20 +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:
message.sender = self.user.handles[0] # TODO : Which handle to use?
if message.sender not in message.participants:
message.participants.append(message.sender)
2023-07-29 17:57:20 +00:00
self._cache_keys(message.participants)
2023-07-28 23:20:32 +00:00
# Set the group id, if it isn't already
if message.group_id is None:
message.group_id = str(uuid.uuid4()).upper() # TODO: Keep track of group ids?
2023-07-29 17:57:20 +00:00
message_id = uuid.uuid4()
2023-07-28 23:20:32 +00:00
if message.id is None:
2023-07-29 17:57:20 +00:00
message.id = str(message_id).upper()
# Turn the message into a raw message
raw = message.to_raw()
import base64
bundled_payloads = []
for participant in message.participants:
for push_token in self.USER_CACHE[participant]:
2023-07-29 21:14:25 +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)
bundled_payloads.append({
'tP': participant,
'D': not participant == message.sender, # TODO: Should this be false sometimes? For self messages?
2023-07-29 21:14:25 +00:00
'sT': self.KEY_CACHE[push_token][1],
2023-07-29 17:57:20 +00:00
'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 = {
'fcn': 1,
'c': 100,
'E': 'pair',
'ua': '[macOS,13.4.1,22F82,MacBookPro18,3]',
'v': 8,
2023-07-30 21:00:04 +00:00
'i': int.from_bytes(msg_id, 'big'),
2023-07-29 17:57:20 +00:00
'U': message_id.bytes,
'dtl': bundled_payloads,
'sP': message.sender,
2023-07-30 21:00:04 +00:00
#'e': time.time_ns(),
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
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:
2023-07-30 21:00:04 +00:00
if num_recv == len(bundled_payloads):
2023-07-29 17:57:20 +00:00
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)
2023-07-30 21:00:04 +00:00
num_recv += 1