From 9822d6f66fd775b1c129c023f54bbaa47908ed9e Mon Sep 17 00:00:00 2001
From: JJTech0130 <jjtech@jjtech.dev>
Date: Thu, 17 Aug 2023 22:05:18 -0400
Subject: [PATCH] squash more bugs

---
 apns.py      |  9 +++++++--
 demo.py      |  6 +++---
 ids/query.py |  1 -
 imessage.py  | 10 ++++------
 4 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/apns.py b/apns.py
index 45f1042..e3883fb 100644
--- a/apns.py
+++ b/apns.py
@@ -100,7 +100,11 @@ class APNSConnection:
         # Check if anything currently in the queue matches the id
         for payload in self._incoming_queue:
             if payload.id == id:
-                return payload
+                if filter is not None:
+                    if filter(payload):
+                        return payload
+                else:
+                    return payload
         while True:
             await self._queue_park.wait()  # Wait for a new payload to be added to the queue
             logger.debug(f"Woken by event, checking for {id}")
@@ -256,6 +260,7 @@ class APNSConnection:
 
     async def send_notification(self, topic: str, payload: bytes, id=None):
         """Sends a notification to the APNs server"""
+        logger.debug(f"Sending notification to topic {topic}")
         if id is None:
             id = random.randbytes(4)
 
@@ -280,7 +285,7 @@ class APNSConnection:
         """Waits for a notification to be received, and acks it"""
 
         def f(payload: APNSPayload):
-            if payload.fields_with_id(1)[0].value != sha1(topic.encode()).digest():
+            if payload.fields_with_id(2)[0].value != sha1(topic.encode()).digest():
                 return False
             if filter is not None:
                 return filter(payload)
diff --git a/demo.py b/demo.py
index f6c0797..cf42096 100644
--- a/demo.py
+++ b/demo.py
@@ -23,15 +23,15 @@ logging.getLogger("py.warnings").setLevel(logging.ERROR)  # Ignore warnings from
 logging.getLogger("asyncio").setLevel(logging.WARNING)
 logging.getLogger("jelly").setLevel(logging.INFO)
 logging.getLogger("nac").setLevel(logging.INFO)
-logging.getLogger("apns").setLevel(logging.DEBUG)
+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.getLogger("imessage").setLevel(logging.INFO)
 
 logging.captureWarnings(True)
 
-process = Popen(["git", "rev-parse", "HEAD"], stdout=PIPE)
+process = Popen(["git", "rev-parse", "HEAD"], stdout=PIPE) # type: ignore
 (commit_hash, err) = process.communicate()
 exit_code = process.wait()
 commit_hash = commit_hash.decode().strip()
diff --git a/ids/query.py b/ids/query.py
index 63e2527..716687b 100644
--- a/ids/query.py
+++ b/ids/query.py
@@ -55,7 +55,6 @@ async def lookup(
         if body is None:
             return False
         body = plistlib.loads(body)
-        logging.warning(body.get('U'))
         return body.get('U') == msg_id
 
     payload = await conn.expect_notification(topic, check)
diff --git a/imessage.py b/imessage.py
index ac76369..9f37bae 100644
--- a/imessage.py
+++ b/imessage.py
@@ -626,21 +626,19 @@ class iMessageUser:
 
         await self.connection.send_notification(topic, body, message_id)
 
-    async def _receive_raw(self, c: int, topic: str) -> dict | None:
+    async def _receive_raw(self, c: int, topic: str) -> dict:
         def check(payload: apns.APNSPayload):
             # Check if the "c" key matches
             body = payload.fields_with_id(3)[0].value
             if body is None:
                 return False
             body = plistlib.loads(body)
-            if not "c" in body or "c" != c:
+            if not "c" in body or body["c"] != c:
                 return False
             return True
         
         payload = await self.connection.expect_notification(topic, check)
 
-        if payload is None:
-            return None
         body = payload.fields_with_id(3)[0].value
         body = plistlib.loads(body)
         return body
@@ -713,8 +711,8 @@ class iMessageUser:
 
         while count < total and time.time() - start < 2:
             resp = await self._receive_raw(255, topic)
-            if resp is None:
-                continue
+            #if resp is None:
+            #    continue
             count += 1
 
             logger.debug(f"Received response : {resp}")