pypush-plus-plus/ids/query.py

70 lines
1.7 KiB
Python
Raw Normal View History

2023-05-09 14:36:33 -05:00
import gzip
import plistlib
import random
from base64 import b64encode
import apns
import bags
2023-05-09 16:03:27 -05:00
2023-05-09 18:31:09 -05:00
from ._helpers import KeyPair, PROTOCOL_VERSION
2023-05-09 17:01:32 -05:00
from . import signing
2023-05-09 14:36:33 -05:00
2023-05-09 18:29:17 -05:00
def lookup(
2023-05-09 14:36:33 -05:00
conn: apns.APNSConnection,
2023-05-09 17:01:32 -05:00
self_uri: str,
2023-05-09 18:29:17 -05:00
id_keypair: KeyPair,
query: list[str],
topic,
2023-05-09 14:36:33 -05:00
) -> bytes:
2023-05-09 18:29:17 -05:00
BAG_KEY = "id-query"
2023-05-09 18:31:09 -05:00
2023-05-09 18:29:17 -05:00
conn.filter([topic])
2023-05-09 18:31:09 -05:00
body = plistlib.dumps({"uris": query})
2023-05-09 14:36:33 -05:00
body = gzip.compress(body, mtime=0)
push_token = b64encode(conn.token).decode()
headers = {
2023-05-09 17:01:32 -05:00
"x-id-self-uri": self_uri,
2023-05-09 18:31:09 -05:00
"x-protocol-version": PROTOCOL_VERSION,
2023-05-09 14:36:33 -05:00
}
2023-05-09 18:31:09 -05:00
signing.add_id_signature(headers, body, BAG_KEY, id_keypair, push_token)
2023-05-09 14:36:33 -05:00
msg_id = random.randbytes(16)
req = {
"cT": "application/x-apple-plist",
"U": msg_id,
"c": 96,
2023-05-09 18:29:17 -05:00
"u": bags.ids_bag()[BAG_KEY],
2023-05-09 14:36:33 -05:00
"h": headers,
"v": 2,
"b": body,
}
2023-05-09 18:29:17 -05:00
2023-05-09 14:36:33 -05:00
conn.send_message(topic, plistlib.dumps(req, fmt=plistlib.FMT_BINARY))
def check_response(x):
if x[0] != 0x0A:
return False
resp_body = apns._get_field(x[1], 3)
if resp_body is None:
return False
resp_body = plistlib.loads(resp_body)
return resp_body.get('U') == msg_id
2023-05-09 14:36:33 -05:00
# Lambda to check if the response is the one we want
payload = conn.incoming_queue.wait_pop_find(check_response)
resp = apns._get_field(payload[1], 3)
2023-05-09 18:29:17 -05:00
resp = plistlib.loads(resp)
2023-05-09 14:36:33 -05:00
resp = gzip.decompress(resp["b"])
resp = plistlib.loads(resp)
2023-05-09 19:01:22 -05:00
if resp['status'] != 0:
raise Exception(f'Query failed: {resp}')
if not 'results' in resp:
raise Exception(f'No results in response: {resp}')
return resp['results']