demo.py

#!/usr/bin/env python

import json
import logging
import os
import threading
import platform
import time
from base64 import b64decode, b64encode
from getpass import getpass

from rich.logging import RichHandler

minor_version = int(platform.python_version().split('.')[1])

if minor_version < 10 or minor_version > 11:
    raise Exception(f"Incompatible Python version '{platform.python_version()}'")

import apns
import ids
import imessage

logging.basicConfig(
    level=logging.NOTSET, format="%(message)s", datefmt="[%X]", handlers=[RichHandler()]
)

# Set sane log levels
logging.getLogger("urllib3").setLevel(logging.WARNING)
logging.getLogger("py.warnings").setLevel(logging.ERROR) # Ignore warnings from urllib3
logging.getLogger("asyncio").setLevel(logging.WARNING)
logging.getLogger("jelly").setLevel(logging.INFO)
logging.getLogger("nac").setLevel(logging.INFO)
logging.getLogger("apns").setLevel(logging.INFO)
logging.getLogger("albert").setLevel(logging.INFO)
logging.getLogger("ids").setLevel(logging.DEBUG)
logging.getLogger("bags").setLevel(logging.INFO)
logging.getLogger("imessage").setLevel(logging.DEBUG)

logging.captureWarnings(True)

# Try and load config.json
try:
    with open("config.json", "r") as f:
        CONFIG = json.load(f)
except FileNotFoundError:
    CONFIG = {}

conn = apns.APNSConnection(
    CONFIG.get("push", {}).get("key"), CONFIG.get("push", {}).get("cert")
)


def safe_b64decode(s):
    try:
        return b64decode(s)
    except:
        return None


conn.connect(token=safe_b64decode(CONFIG.get("push", {}).get("token")))
conn.set_state(1)
conn.filter(["com.apple.madrid"])

user = ids.IDSUser(conn)

if CONFIG.get("auth", {}).get("cert") is not None:
    auth_keypair = ids._helpers.KeyPair(CONFIG["auth"]["key"], CONFIG["auth"]["cert"])
    user_id = CONFIG["auth"]["user_id"]
    handles = CONFIG["auth"]["handles"]
    user.restore_authentication(auth_keypair, user_id, handles)
else:
    username = input("Username: ")
    password = getpass("Password: ")

    user.authenticate(username, password)

user.encryption_identity = ids.identity.IDSIdentity(
    encryption_key=CONFIG.get("encryption", {}).get("rsa_key"),
    signing_key=CONFIG.get("encryption", {}).get("ec_key"),
)

if (
    CONFIG.get("id", {}).get("cert") is not None
    and user.encryption_identity is not None
):
    id_keypair = ids._helpers.KeyPair(CONFIG["id"]["key"], CONFIG["id"]["cert"])
    user.restore_identity(id_keypair)
else:
    logging.info("Registering new identity...")
    import emulated.nac

    vd = emulated.nac.generate_validation_data()
    vd = b64encode(vd).decode()

    user.register(vd)

logging.info("Waiting for incoming messages...")

# Write config.json
CONFIG["encryption"] = {
    "rsa_key": user.encryption_identity.encryption_key,
    "ec_key": user.encryption_identity.signing_key,
}
CONFIG["id"] = {
    "key": user._id_keypair.key,
    "cert": user._id_keypair.cert,
}
CONFIG["auth"] = {
    "key": user._auth_keypair.key,
    "cert": user._auth_keypair.cert,
    "user_id": user.user_id,
    "handles": user.handles,
}
CONFIG["push"] = {
    "token": b64encode(user.push_connection.token).decode(),
    "key": user.push_connection.private_key,
    "cert": user.push_connection.cert,
}

with open("config.json", "w") as f:
    json.dump(CONFIG, f, indent=4)

im = imessage.iMessageUser(conn, user)

INPUT_QUEUE = apns.IncomingQueue()

def input_thread():
    from prompt_toolkit import prompt
    while True:

        try:
            msg = prompt('>> ')
        except:
            msg = 'quit'
        INPUT_QUEUE.append(msg)

threading.Thread(target=input_thread, daemon=True).start()

print("Type 'help' for help")  

def fixup_handle(handle):
    if handle.startswith('tel:+'):
        return handle
    elif handle.startswith('mailto:'):
        return handle
    elif handle.startswith('tel:'):
        return 'tel:+' + handle[4:]
    elif handle.startswith('+'):
        return 'tel:' + handle
    # If the handle starts with a number
    elif handle[0].isdigit():
        # If the handle is 10 digits, assume it's a US number
        if len(handle) == 10:
            return 'tel:+1' + handle
        # If the handle is 11 digits, assume it's a US number with country code
        elif len(handle) == 11:
            return 'tel:+' + handle
    else: # Assume it's an email
        return 'mailto:' + handle

current_participants = []
current_effect = None
while True:
    msg = im.receive()
    if msg is not None:
        # print(f'[{msg.sender}] {msg.text}')
        print(msg.to_string())

        attachments = msg.attachments()
        if len(attachments) > 0:
            attachments_path = f"attachments/{msg.id}/"
            os.makedirs(attachments_path, exist_ok=True)

            for attachment in attachments:
                with open(attachments_path + attachment.name, "wb") as attachment_file:
                    attachment_file.write(attachment.versions[0].data())

            print(f"({len(attachments)} attachment{'s have' if len(attachments) != 1 else ' has'} been downloaded and put "
                  f"in {attachments_path})")
    
    if len(INPUT_QUEUE) > 0:
        msg = INPUT_QUEUE.pop()
        if msg == '': continue
        if msg == 'help' or msg == 'h':
            print('help (h): show this message')
            print('quit (q): quit')
            #print('send (s) [recipient] [message]: send a message')
            print('filter (f) [recipient]: set the current chat')
            print('effect (e): adds an iMessage effect to the next sent message')
            print('note: recipient must start with tel: or mailto: and include the country code')
            print('handle <handle>: set the current handle (for sending messages)')
            print('\\: escape commands (will be removed from message)')
        elif msg == 'quit' or msg == 'q':
            break
        elif msg == 'effect' or msg == 'e' or msg.startswith("effect ") or msg.startswith("e "):
            msg = msg.split(" ")
            if len(msg) < 2 or msg[1] == "":
                print("effect [effect namespace]")
            else:
                print(f"next message will be sent with [{msg[1]}]")
                current_effect = msg[1]
        elif msg == 'filter' or msg == 'f' or msg.startswith('filter ') or msg.startswith('f '):
            # Set the curernt chat
            msg = msg.split(' ')
            if len(msg) < 2 or msg[1] == '':
                print('filter [recipients]')
            else:
                print(f'Filtering to {[fixup_handle(h) for h in msg[1:]]}')
                current_participants = [fixup_handle(h) for h in msg[1:]]
        elif msg == 'handle' or msg.startswith('handle '):
            msg = msg.split(' ')
            if len(msg) < 2 or msg[1] == '':
                print('handle [handle]')
                print('Available handles:')
                for h in user.handles:
                    if h == user.current_handle:
                        print(f'\t{h} (current)')
                    else:
                        print(f'\t{h}')
            else:
                h = msg[1]
                h = fixup_handle(h)
                if h in user.handles:
                    print(f'Using {h} as handle')
                    user.current_handle = h
                else:
                    print(f'Handle {h} not found')

        elif current_participants != []:
            if msg.startswith('\\'):
                msg = msg[1:]
            im.send(imessage.iMessage(
                text=msg,
                participants=current_participants,
                sender=user.current_handle,
                effect=current_effect
            ))
            current_effect = None
        else:
            print('No chat selected, use help for help')

    time.sleep(0.1)
        
        # elif msg.startswith('send') or msg.startswith('s'):
        #     msg = msg.split(' ')
        #     if len(msg) < 3:
        #         print('send [recipient] [message]')
        #     else:
        #         im.send(imessage.iMessage(
        #             text=' '.join(msg[2:]),
        #             participants=[msg[1], user.handles[0]],
        #             #sender=user.handles[0]
        #         ))