From 5c592c9dee04e16479ff1c5289913b2721036121 Mon Sep 17 00:00:00 2001
From: JJTech0130 <jjtech@jjtech.dev>
Date: Mon, 31 Jul 2023 16:30:06 -0400
Subject: [PATCH] implement handle selection

---
 demo.py         | 22 +++++++++++++++++++---
 ids/__init__.py |  4 +++-
 imessage.py     |  7 +++++++
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/demo.py b/demo.py
index faa412c..45bb3b0 100644
--- a/demo.py
+++ b/demo.py
@@ -143,17 +143,18 @@ while True:
             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.startswith("effect ") or msg.startswith("e "):
+        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.startswith('filter ') or msg.startswith('f '):
+        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] == '':
@@ -161,13 +162,28 @@ while True:
             else:
                 print(f'Filtering to {msg[1:]}')
                 current_participants = 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:
+                    print(f'\t{h}')
+            else:
+                h = msg[1]
+                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.handles[0],
+                sender=user.current_handle,
                 effect=current_effect
             ))
             current_effect = None
diff --git a/ids/__init__.py b/ids/__init__.py
index 5d7226b..b0a74d2 100644
--- a/ids/__init__.py
+++ b/ids/__init__.py
@@ -45,6 +45,8 @@ class IDSUser:
             self._auth_keypair,
             self._push_keypair,
         )
+        self.current_handle = self.handles[0]
+
 
     # Uses an existing authentication keypair
     def restore_authentication(
@@ -79,5 +81,5 @@ class IDSUser:
         self._id_keypair = id_keypair
 
     def lookup(self, uris: list[str], topic: str = "com.apple.madrid") -> any:
-        return query.lookup(self.push_connection, self.handles[0], self._id_keypair, uris, topic)
+        return query.lookup(self.push_connection, self.current_handle, self._id_keypair, uris, topic)
         
diff --git a/imessage.py b/imessage.py
index 45fb528..7314b5d 100644
--- a/imessage.py
+++ b/imessage.py
@@ -343,12 +343,19 @@ class iMessageUser:
         
         return iMessage.from_raw(decrypted, body['sP'])
 
+    KEY_CACHE_HANDLE: str = ""
     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]"""
 
     def _cache_keys(self, participants: list[str]):
+        # 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 = {}
+        
         # 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