mirror of
https://github.com/Sneed-Group/pypush-plus-plus
synced 2024-12-24 20:02:31 -06:00
2960b7c63c
- 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
428 lines
14 KiB
Python
428 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import random
|
|
import ssl
|
|
import time
|
|
from base64 import b64encode
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from hashlib import sha1
|
|
from typing import Callable
|
|
|
|
import trio
|
|
from cryptography import x509
|
|
from cryptography.hazmat.primitives import hashes, serialization
|
|
from cryptography.hazmat.primitives.asymmetric import padding
|
|
|
|
import albert
|
|
import bags
|
|
|
|
logger = logging.getLogger("apns")
|
|
|
|
# Pick a random courier server from 01 to APNSCourierHostcount
|
|
try:
|
|
COURIER_HOST = f"{random.randint(1, bags.apns_init_bag()['APNSCourierHostcount'])}-{bags.apns_init_bag()['APNSCourierHostname']}"
|
|
except:
|
|
COURIER_HOST = "01-courier.push.apple.com"
|
|
COURIER_PORT = 5223
|
|
ALPN = [b"apns-security-v3"]
|
|
|
|
|
|
async def apns_test():
|
|
async with APNSConnection.start() as connection:
|
|
print(b64encode(connection.credentials.token).decode())
|
|
while True:
|
|
await trio.sleep(1)
|
|
print(".")
|
|
# await connection.set_state(1)
|
|
|
|
print("Finished")
|
|
|
|
|
|
def main():
|
|
from rich.logging import RichHandler
|
|
|
|
logging.basicConfig(
|
|
level=logging.NOTSET,
|
|
format="%(message)s",
|
|
datefmt="[%X]",
|
|
handlers=[RichHandler()],
|
|
)
|
|
|
|
# Set sane log levels
|
|
logging.getLogger("urllib3").setLevel(logging.WARNING)
|
|
logging.getLogger("py.warnings").setLevel(
|
|
logging.ERROR
|
|
) # Ignore warnings from urllib3
|
|
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("albert").setLevel(logging.INFO)
|
|
logging.getLogger("ids").setLevel(logging.DEBUG)
|
|
logging.getLogger("bags").setLevel(logging.INFO)
|
|
logging.getLogger("imessage").setLevel(logging.DEBUG)
|
|
|
|
logging.captureWarnings(True)
|
|
|
|
print("APNs Test:")
|
|
|
|
trio.run(apns_test)
|
|
|
|
|
|
@dataclass
|
|
class PushCredentials:
|
|
private_key: str = ""
|
|
cert: str = ""
|
|
token: bytes = b""
|
|
|
|
|
|
class APNSConnection:
|
|
"""A connection to the APNs server"""
|
|
|
|
_incoming_queue: list[APNSPayload] = []
|
|
"""A queue of payloads that have been received from the APNs server"""
|
|
_queue_park: trio.Event = trio.Event()
|
|
"""An event that is set when a new payload is added to the queue"""
|
|
|
|
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[[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,
|
|
the payload will be ignored and another will be waited for.
|
|
|
|
NOTE: It is not defined what happens if receive is called twice with the same id and filter,
|
|
as the first payload will be removed from the queue, so the second call might never return
|
|
"""
|
|
|
|
# Check if anything currently in the queue matches the id
|
|
for payload in self._incoming_queue:
|
|
if payload.id == id:
|
|
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}")
|
|
# Check if the new payload matches the id
|
|
if len(self._incoming_queue) == 0:
|
|
continue # all payloads have been removed by someone else
|
|
if self._incoming_queue[-1].id != id:
|
|
continue
|
|
if filter is not None:
|
|
if not filter(self._incoming_queue[-1]):
|
|
continue
|
|
return self._incoming_queue.pop()
|
|
# Otherwise, wait for another payload to be added to the queue
|
|
|
|
async def _queue_filler(self):
|
|
"""Fills the queue with payloads from the APNs socket"""
|
|
while True:
|
|
payload = await APNSPayload.read_from_stream(self.sock)
|
|
|
|
logger.debug(f"Received payload: {payload}")
|
|
|
|
self._incoming_queue.append(payload)
|
|
|
|
# TODO: Hack: Send an ACK if this is a notification
|
|
# We do this because as of now pypush does not handle all incoming notifications
|
|
# and if you do not ACK a notification, APNs will keep resending it and eventually kill the connection
|
|
if payload.id == 0xA:
|
|
await self._send_ack(payload.fields_with_id(4)[0].value)
|
|
|
|
# Signal to any waiting tasks that we have a new payload
|
|
self._queue_park.set()
|
|
self._queue_park = trio.Event() # Reset the event
|
|
|
|
logger.debug(f"Queue length: {len(self._incoming_queue)}")
|
|
|
|
async def _keep_alive(self):
|
|
"""Sends keep alive messages to the APNs server every 5 minutes"""
|
|
while True:
|
|
await trio.sleep(300)
|
|
logger.debug("Sending keep alive message")
|
|
await self._send(APNSPayload(0x0C, []))
|
|
await self._receive(0x0D)
|
|
logger.debug("Got keep alive response")
|
|
|
|
@asynccontextmanager
|
|
@staticmethod
|
|
async def start(credentials: PushCredentials = PushCredentials()):
|
|
"""Sets up a nursery and connection and yields the connection"""
|
|
async with trio.open_nursery() as nursery:
|
|
connection = APNSConnection(nursery, credentials)
|
|
await connection.connect()
|
|
yield connection
|
|
nursery.cancel_scope.cancel() # Cancel heartbeat and queue filler tasks
|
|
await connection.sock.aclose() # Close the socket
|
|
|
|
def __init__(
|
|
self, nursery: trio.Nursery, credentials: PushCredentials = PushCredentials()
|
|
):
|
|
"""Creates a raw APNSConnection. Make sure to call aclose() on the socket and cancel the nursery when you're done with it"""
|
|
self._nursery = nursery
|
|
self.credentials = credentials
|
|
|
|
async def _connect_socket(self):
|
|
sock = await trio.open_tcp_stream(COURIER_HOST, COURIER_PORT)
|
|
|
|
context = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
|
|
context.set_alpn_protocols(["apns-security-v3"])
|
|
|
|
# Turn off certificate verification, for the proxy
|
|
context.check_hostname = False
|
|
context.verify_mode = ssl.CERT_NONE
|
|
|
|
self.sock = trio.SSLStream(sock, context, server_hostname=COURIER_HOST)
|
|
|
|
await self.sock.do_handshake()
|
|
|
|
async def connect(self):
|
|
"""Connects to the APNs server and starts the keep alive and queue filler tasks"""
|
|
await self._connect_socket()
|
|
|
|
logger.info(f"Connected to APNs ({COURIER_HOST})")
|
|
|
|
if self.credentials.cert == "" or self.credentials.private_key == "":
|
|
(
|
|
self.credentials.private_key,
|
|
self.credentials.cert,
|
|
) = albert.generate_push_cert()
|
|
|
|
# Start the queue filler and keep alive tasks
|
|
self._nursery.start_soon(self._queue_filler)
|
|
self._nursery.start_soon(self._keep_alive)
|
|
|
|
self.credentials.token = await self._connect(self.credentials.token)
|
|
|
|
async def _connect(self, token: bytes | None = None, root: bool = False) -> bytes:
|
|
"""Sends the APNs connect message"""
|
|
|
|
cert = x509.load_pem_x509_certificate(self.credentials.cert.encode())
|
|
private_key = serialization.load_pem_private_key(
|
|
self.credentials.private_key.encode(), password=None
|
|
)
|
|
|
|
if token is None:
|
|
logger.debug(f"Sending connect message without token (root={root})")
|
|
else:
|
|
logger.debug(
|
|
f"Sending connect message with token {b64encode(token).decode()} (root={root})"
|
|
)
|
|
flags = 0b01000001
|
|
if root:
|
|
flags |= 0b0100
|
|
|
|
cert = cert.public_bytes(serialization.Encoding.DER)
|
|
nonce = (
|
|
b"\x00" + int(time.time() * 1000).to_bytes(8, "big") + random.randbytes(8)
|
|
)
|
|
signature = b"\x01\x01" + private_key.sign(nonce, padding.PKCS1v15(), hashes.SHA1()) # type: ignore
|
|
|
|
payload = APNSPayload(
|
|
7,
|
|
[
|
|
APNSField(0x2, b"\x01"),
|
|
APNSField(0x5, flags.to_bytes(4, "big")),
|
|
APNSField(0xC, cert),
|
|
APNSField(0xD, nonce),
|
|
APNSField(0xE, signature),
|
|
],
|
|
)
|
|
|
|
if token:
|
|
payload.fields.insert(0, APNSField(0x1, token))
|
|
|
|
await self._send(payload)
|
|
|
|
payload = await self._receive(8)
|
|
|
|
if payload.fields_with_id(1)[0].value != b"\x00":
|
|
raise Exception("Failed to connect")
|
|
|
|
if len(payload.fields_with_id(3)) > 0:
|
|
new_token = payload.fields_with_id(3)[0].value
|
|
else:
|
|
if token is None:
|
|
raise Exception("No token received")
|
|
new_token = token
|
|
|
|
logger.debug(
|
|
f"Received connect response with token {b64encode(new_token).decode()}"
|
|
)
|
|
|
|
return new_token
|
|
|
|
old_topics = []
|
|
async def filter(self, topics: list[str]):
|
|
"""Sends the APNs filter message"""
|
|
if topics == self.old_topics:
|
|
return
|
|
topics = list(set(topics + self.old_topics))
|
|
self.old_topics = topics
|
|
logger.debug(f"Sending filter message with topics {topics}")
|
|
|
|
payload = APNSPayload(9, [APNSField(1, self.credentials.token)])
|
|
|
|
for topic in topics:
|
|
payload.fields.append(APNSField(2, sha1(topic.encode()).digest()))
|
|
|
|
await payload.write_to_stream(self.sock)
|
|
|
|
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)
|
|
|
|
p = APNSPayload(
|
|
0xA,
|
|
[
|
|
APNSField(4, id),
|
|
APNSField(1, sha1(topic.encode()).digest()),
|
|
APNSField(2, self.credentials.token),
|
|
APNSField(3, payload),
|
|
],
|
|
)
|
|
|
|
await self._send(p)
|
|
|
|
# Wait for ACK
|
|
r = await self._receive(0xB)
|
|
|
|
# TODO: Check ACK code
|
|
|
|
async def expect_notification(self, topics: str | list[str], filter: Callable | None = None):
|
|
"""Waits for a notification to be received, and acks it"""
|
|
|
|
if isinstance(topics, list):
|
|
topic_hashes = [sha1(topic.encode()).digest() for topic in topics]
|
|
else:
|
|
topic_hashes = [sha1(topics.encode()).digest()]
|
|
|
|
def f(payload: APNSPayload):
|
|
if payload.fields_with_id(2)[0].value not in topic_hashes:
|
|
return False
|
|
if filter is not None:
|
|
return filter(payload)
|
|
return True
|
|
|
|
r = await self._receive(0xA, f)
|
|
# await self._send_ack(r.fields_with_id(4)[0].value)
|
|
return r
|
|
|
|
async def set_state(self, state: int):
|
|
"""Sends the APNs state message"""
|
|
logger.debug(f"Sending state message with state {state}")
|
|
await self._send(
|
|
APNSPayload(
|
|
0x14,
|
|
[
|
|
APNSField(1, state.to_bytes(1, "big")),
|
|
APNSField(2, 0x7FFFFFFF.to_bytes(4, "big")),
|
|
],
|
|
)
|
|
)
|
|
|
|
async def _send_ack(self, id: bytes):
|
|
"""Sends an ACK for a notification with the given id"""
|
|
logger.debug(f"Sending ACK for message {id}")
|
|
payload = APNSPayload(
|
|
0xB,
|
|
[
|
|
APNSField(1, self.credentials.token),
|
|
APNSField(4, id),
|
|
APNSField(8, b"\x00"),
|
|
],
|
|
)
|
|
await self._send(payload)
|
|
|
|
|
|
@dataclass
|
|
class APNSField:
|
|
"""A field in an APNS payload"""
|
|
|
|
id: int
|
|
value: bytes
|
|
|
|
@staticmethod
|
|
def from_buffer(stream: bytes) -> APNSField:
|
|
id = int.from_bytes(stream[:1], "big")
|
|
length = int.from_bytes(stream[1:3], "big")
|
|
value = stream[3 : 3 + length]
|
|
return APNSField(id, value)
|
|
|
|
def to_buffer(self) -> bytes:
|
|
return (
|
|
self.id.to_bytes(1, "big") + len(self.value).to_bytes(2, "big") + self.value
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class APNSPayload:
|
|
"""An APNS payload"""
|
|
|
|
id: int
|
|
fields: list[APNSField]
|
|
|
|
@staticmethod
|
|
async def read_from_stream(stream: trio.abc.Stream) -> APNSPayload:
|
|
"""Reads a payload from the given stream"""
|
|
if not (id_bytes := await stream.receive_some(1)):
|
|
raise Exception("Unable to read payload id from stream")
|
|
id: int = int.from_bytes(id_bytes, "big")
|
|
|
|
if (length := await stream.receive_some(4)) is None:
|
|
raise Exception("Unable to read payload length from stream")
|
|
length = int.from_bytes(length, "big")
|
|
|
|
if length == 0:
|
|
return APNSPayload(id, [])
|
|
|
|
buffer = await stream.receive_some(length)
|
|
if buffer is None:
|
|
raise Exception("Unable to read payload from stream")
|
|
fields = []
|
|
|
|
while len(buffer) > 0:
|
|
field = APNSField.from_buffer(buffer)
|
|
fields.append(field)
|
|
buffer = buffer[3 + len(field.value) :]
|
|
|
|
return APNSPayload(id, fields)
|
|
|
|
async def write_to_stream(self, stream: trio.abc.Stream):
|
|
"""Writes the payload to the given stream"""
|
|
payload = b""
|
|
|
|
for field in self.fields:
|
|
payload += field.to_buffer()
|
|
|
|
buffer = self.id.to_bytes(1, "big") + len(payload).to_bytes(4, "big") + payload
|
|
|
|
await stream.send_all(buffer)
|
|
|
|
def fields_with_id(self, id: int):
|
|
"""Returns all fields with the given id"""
|
|
return [field for field in self.fields if field.id == id]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|