From 7158532c668ee30f5ab8d77c586bc28fad1ec0ca Mon Sep 17 00:00:00 2001
From: JJTech0130 <jjtech@jjtech.dev>
Date: Tue, 9 May 2023 18:01:32 -0400
Subject: [PATCH] trying to fix query after refactor

---
 demo.py         |  2 ++
 ids/__init__.py |  6 +++++-
 ids/query.py    | 13 ++++++++-----
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/demo.py b/demo.py
index ba21f4c..98d16ce 100644
--- a/demo.py
+++ b/demo.py
@@ -51,6 +51,8 @@ else:
     vd = input_multiline("Enter validation data: ")
     user.register(vd)
 
+print(user.lookup(["mailto:textgpt@icloud.com"]))
+
 # Write config.json
 CONFIG["id"] = {
     "key": user._id_keypair.key,
diff --git a/ids/__init__.py b/ids/__init__.py
index ddf1bae..fa316a3 100644
--- a/ids/__init__.py
+++ b/ids/__init__.py
@@ -2,7 +2,7 @@ from base64 import b64encode
 
 import apns
 
-from . import _helpers, identity, profile
+from . import _helpers, identity, profile, query
 
 
 class IDSUser:
@@ -66,3 +66,7 @@ class IDSUser:
 
     def restore_identity(self, id_keypair: _helpers.KeyPair):
         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, topic, uris)
+        
diff --git a/ids/query.py b/ids/query.py
index 36b991a..9d3bcb1 100644
--- a/ids/query.py
+++ b/ids/query.py
@@ -6,7 +6,8 @@ from base64 import b64encode
 import apns
 import bags
 
-from . import USER_AGENT, KeyPair, signing
+from ._helpers import USER_AGENT, KeyPair
+from . import signing
 
 
 def _send_request(
@@ -15,7 +16,7 @@ def _send_request(
     topic: str,
     body: bytes,
     keypair: KeyPair,
-    username: str,
+    self_uri: str,
 ) -> bytes:
     body = gzip.compress(body, mtime=0)
 
@@ -25,10 +26,11 @@ def _send_request(
     # signature, nonce = _sign_payload(keypair.key, bag_key, "", push_token, body)
 
     headers = {
-        "x-id-self-uri": "mailto:" + username,
+        "x-id-self-uri": self_uri,
         "User-Agent": USER_AGENT,
         "x-protocol-version": "1630",
     }
+    print(headers)
     signing.add_id_signature(headers, body, bag_key, keypair, push_token)
 
     # print(headers)
@@ -46,6 +48,7 @@ def _send_request(
         "b": body,
     }
 
+    print(req)
     conn.send_message(topic, plistlib.dumps(req, fmt=plistlib.FMT_BINARY))
     # resp = conn.wait_for_packet(0x0A)
 
@@ -73,11 +76,11 @@ def _send_request(
 # topic: the IDS topic to query
 # query: a list of URIs to query
 def lookup(
-    conn: apns.APNSConnection, self: str, keypair: KeyPair, topic: str, query: list[str]
+    conn: apns.APNSConnection, self_uri: str, id_keypair: KeyPair, topic: str, query: list[str]
 ) -> any:
     conn.filter([topic])
     query = {"uris": query}
-    resp = _send_request(conn, "id-query", topic, plistlib.dumps(query), keypair, self)
+    resp = _send_request(conn, "id-query", topic, plistlib.dumps(query), id_keypair, self_uri)
     # resp = plistlib.loads(resp)
     # print(resp)
     resp = gzip.decompress(resp["b"])