- Remove eol whitespace

- Add retry for writing to stream since it's busy sometimes
- add a few more type hints
- add back demo.py tui commands
- add sending with effects

apns.py

@@ -88,9 +88,19 @@ class APNSConnection:
 
     async def _send(self, payload: APNSPayload):
         """Sends a payload to the APNs server"""
+        while True:
+            try:
                 await payload.write_to_stream(self.sock)
+                return
+            except trio.BusyResourceError:
+                print("Can't send payload, stream is busy; trying again in 0.2")
+                await trio.sleep(0.2)
+                continue
+            except Exception as e:
+                print(f"Can't send payload: {e}")
+                return
 
-    async def _receive(self, id: int, filter: Callable | None = None):
+    async def _receive(self, id: int, filter: Callable[[APNSPayload], bool] | None = None):
         """
         Waits for a payload with the given id to be added to the queue, then returns it.
         If filter is not None, it will be called with the payload as an argument, and if it returns False,
@@ -237,7 +247,7 @@ class APNSConnection:
             ],
         )
 
-        if token != b"" and token is not None:
+        if token:
             payload.fields.insert(0, APNSField(0x1, token))
 
         await self._send(payload)
@@ -315,7 +325,7 @@ class APNSConnection:
             return True
 
         r = await self._receive(0xA, f)
-        #await self._send_ack(r.fields_with_id(4)[0].value)
+        # await self._send_ack(r.fields_with_id(4)[0].value)
         return r
 
     async def set_state(self, state: int):
@@ -342,7 +352,7 @@ class APNSConnection:
                 APNSField(8, b"\x00"),
             ],
         )
-        await payload.write_to_stream(self.sock)
+        await self._send(payload)
 
 
 @dataclass
@@ -375,13 +385,11 @@ class APNSPayload:
     @staticmethod
     async def read_from_stream(stream: trio.abc.Stream) -> APNSPayload:
         """Reads a payload from the given stream"""
-        id = await stream.receive_some(1)
-        if id is None or id == b"":
+        if not (id_bytes := await stream.receive_some(1)):
             raise Exception("Unable to read payload id from stream")
-        id = int.from_bytes(id, "big")
+        id: int = int.from_bytes(id_bytes, "big")
 
-        length = await stream.receive_some(4)
-        if length is None:
+        if (length := await stream.receive_some(4)) is None:
             raise Exception("Unable to read payload length from stream")
         length = int.from_bytes(length, "big")
 

demo.py

@@ -104,10 +104,81 @@ async def main():
             nursery.start_soon(output_task, im)
 
 async def input_task(im: imessage.iMessageUser):
+    current_effect: str | None = None
+    current_participants: list[str] = []
+
+    def is_cmd(cmd_str: str, name: str) -> bool:
+        return cmd_str in [name, name[0]] or cmd_str.startswith(f"{name} ") or cmd_str.startswith(f"{name[0]} ")
+
+    def get_parameters(cmd: str, err_msg: str) -> list[str] | None:
+        sections: list[str] = cmd.split(" ")
+        if len(sections) < 2 or len(sections[1]) == 0:
+            print(err_msg)
+            return None
+        else:
+            return sections[1:]
+
+    def fixup_handle(handle: str) -> str:
+        if handle.startswith("tel:+") or handle.startswith("mailto:"):
+            return handle
+        elif handle.startswith("tel:"):
+            return "tel:+" + handle[4:]
+        elif handle.startswith("+"):
+            return "tel:" + handle
+        elif handle[0].isdigit():
+            # if the handle is 10 digits, assume it's a US number
+            if len(handle) == 10:
+                return "tel:+1" + handle
+            # otherwise just assume it's a full phone number
+            return "tel:+" + handle
+        else: # assume it's an email
+            return "mailto:" + handle
+
     while True:
-        cmd = await trio.to_thread.run_sync(input, "> ", cancellable=True)
-        if cmd != "":
-            await im.send(imessage.iMessage.create(im, cmd, [im.user.current_handle]))
+        if (cmd := await trio.to_thread.run_sync(input, ">> ", cancellable=True)) == "":
+            continue
+
+        if is_cmd(cmd, "help"):
+            print("help (h): show this message")
+            print("quit (q): quit")
+            print("filter (f) [recipient]: set the current chat")
+            print("note: recipient must start with tel: or mailto: and include the country code")
+            print("effect (e): adds an iMessage effect to the next sent message")
+            print("handle [handle]: set the current handle (for sending messages with)")
+            print("\\: escape commands (will be removed from message)")
+            print("all other commands will be treated as message text to be sent")
+        elif is_cmd(cmd, "quit"):
+            exit(0)
+        elif is_cmd(cmd, "effect"):
+            if (effect := get_parameters(cmd, "effect [effect namespace]")) is not None:
+                print(f"next message will be sent with [{effect[0]}]")
+                current_effect = effect[0]
+        elif is_cmd(cmd, "filter"):
+            # set the current chat
+            if (participants := get_parameters(cmd, "filter [recipients]")) is not None:
+                fixed_participants: list[str] = list(map(fixup_handle, participants))
+                print(f"Filtering to {fixed_participants}")
+                current_participants = fixed_participants
+        elif is_cmd(cmd, "handle"):
+            handles: list[str] = im.user.handles
+            av_handles: str = "\n".join([f"\t{h}{' (current)' if h == im.user.current_handle else ''}" for h in handles])
+            err_str: str = f"handle [handle]\nAvailable handles:\n{av_handles}"
+
+            if (input_handles := get_parameters(cmd, err_str)) is not None:
+                handle = fixup_handle(input_handles[0])
+                if handle in handles:
+                    print(f"Using {handle} as handle")
+                    im.user.current_handle = handle
+                else:
+                    print(f"Handle {handle} not found")
+        elif len(current_participants) > 0:
+            if cmd.startswith("\\"):
+                cmd = cmd[1:]
+
+            await im.send(imessage.iMessage.create(im, cmd, current_participants, current_effect))
+            current_effect = None
+        else:
+            print("No chat selected")
 
 async def output_task(im: imessage.iMessageUser):
     while True:

emulated/jelly.py

@@ -35,7 +35,6 @@ class VirtualInstructions:
             else:
                 self.push(args[i])
 
-    
     def call(self, address: int, args: list[int] = []):
         logger.debug(f"Calling {hex(address)} with args {args}")
         self.push(STOP_ADDRESS)
@@ -43,7 +42,6 @@ class VirtualInstructions:
         self.uc.emu_start(address, STOP_ADDRESS)
         return self.uc.reg_read(unicorn.x86_const.UC_X86_REG_RAX)
 
-        
 class Jelly:
     # Constants
     UC_ARCH = unicorn.UC_ARCH_X86

ids/__init__.py

@@ -54,9 +54,7 @@ class IDSUser:
 
 
     # Uses an existing authentication keypair
-    def restore_authentication(
-        self, auth_keypair: _helpers.KeyPair, user_id: str, handles: dict
-    ):
+    def restore_authentication(self, auth_keypair: _helpers.KeyPair, user_id: str, handles: list[str]):
         self._auth_keypair = auth_keypair
         self.user_id = user_id
         self.handles = handles

ids/_helpers.py

@@ -36,4 +36,3 @@ def serialize_key(key) -> str:
             encoding=serialization.Encoding.PEM,
             format=serialization.PublicFormat.SubjectPublicKeyInfo,
         ).decode("utf-8").strip()
-    

ids/profile.py

@@ -146,7 +146,7 @@ def get_auth_cert(user_id, token) -> KeyPair:
     )
 
 
-def get_handles(push_token, user_id: str, auth_key: KeyPair, push_key: KeyPair):
+def get_handles(push_token, user_id: str, auth_key: KeyPair, push_key: KeyPair) -> list[str]:
     BAG_KEY = "id-get-handles"
 
     headers = {

imessage.py

@@ -262,10 +262,17 @@ class SMSIncomingImage(Message):
 
 @dataclass
 class iMessage(Message):
+    """
+    An iMessage
+
+    Description of payload keys:
+    t: The sender token. Normally just a big thing of data/bytes, doesn't need to be decoded in any way
+    U: the id of the message (uuid) as bytes
+    """
     effect: str | None = None
 
     @staticmethod
-    def create(user: "iMessageUser", text: str, participants: list[str]) -> "iMessage":
+    def create(user: "iMessageUser", text: str, participants: list[str], effect: str | None) -> "iMessage":
         """Creates a basic outgoing `iMessage` from the given text and participants"""
 
         sender = user.user.current_handle
@@ -277,6 +284,7 @@ class iMessage(Message):
             sender=sender,
             participants=participants,
             id=uuid.uuid4(),
+            effect=effect
         )
 
     @staticmethod