pypush-plus-plus/apns.py

324 lines
9.7 KiB
Python
Raw Normal View History

2023-04-05 23:52:14 +00:00
from __future__ import annotations
2023-04-11 16:23:04 +00:00
import random
import socket
2023-04-08 02:32:00 +00:00
import threading
import time
2023-04-11 16:23:04 +00:00
from hashlib import sha1
2023-07-24 13:18:21 +00:00
from base64 import b64encode, b64decode
2023-07-24 01:28:29 +00:00
import logging
logger = logging.getLogger("apns")
2023-04-11 16:23:04 +00:00
import tlslite
2023-07-23 19:29:16 +00:00
if tlslite.__version__ != "0.8.0-alpha43":
2023-07-24 01:28:29 +00:00
logger.warning("tlslite-ng is not the correct version!")
logger.warning("Please install tlslite-ng==0.8.0a43 or you will experience issues!")
2023-04-11 16:23:04 +00:00
import albert
2023-05-10 00:01:22 +00:00
import bags
2023-05-10 00:01:22 +00:00
#COURIER_HOST = "windows.courier.push.apple.com" # TODO: Get this from config
# Pick a random courier server from 01 to APNSCourierHostcount
COURIER_HOST = f"{random.randint(1, bags.apns_init_bag()['APNSCourierHostcount'])}-{bags.apns_init_bag()['APNSCourierHostname']}"
COURIER_PORT = 5223
ALPN = [b"apns-security-v2"]
# Connect to the courier server
def _connect(private_key: str, cert: str) -> tlslite.TLSConnection:
# Connect to the courier server
sock = socket.create_connection((COURIER_HOST, COURIER_PORT))
# Wrap the socket in TLS
sock = tlslite.TLSConnection(sock)
# Parse the certificate and private key
cert = tlslite.X509CertChain([tlslite.X509().parse(cert)])
private_key = tlslite.parsePEMKey(private_key, private=True)
# Handshake with the server
sock.handshakeClientCert(cert, private_key, alpn=ALPN)
2023-07-24 13:18:21 +00:00
logger.info(f"Connected to APNs ({COURIER_HOST})")
return sock
2023-04-11 16:23:04 +00:00
2023-05-03 00:53:18 +00:00
2023-05-03 00:51:02 +00:00
class IncomingQueue:
def __init__(self):
self.queue = []
self.lock = threading.Lock()
2023-05-03 00:53:18 +00:00
2023-05-03 00:51:02 +00:00
def append(self, item):
with self.lock:
self.queue.append(item)
2023-05-03 00:53:18 +00:00
2023-07-28 21:31:27 +00:00
def pop(self, index = -1):
2023-05-03 00:51:02 +00:00
with self.lock:
return self.queue.pop(index)
2023-05-03 00:53:18 +00:00
2023-05-03 00:51:02 +00:00
def __getitem__(self, index):
with self.lock:
return self.queue[index]
2023-05-03 00:53:18 +00:00
2023-05-03 00:51:02 +00:00
def __len__(self):
with self.lock:
return len(self.queue)
2023-05-03 00:53:18 +00:00
2023-05-03 00:51:02 +00:00
def find(self, finder):
with self.lock:
return next((i for i in self.queue if finder(i)), None)
2023-05-03 00:53:18 +00:00
2023-05-03 00:51:02 +00:00
def pop_find(self, finder):
with self.lock:
found = next((i for i in self.queue if finder(i)), None)
if found is not None:
# We have the lock, so we can safely remove it
self.queue.remove(found)
return found
def remove_all(self, id):
with self.lock:
self.queue = [i for i in self.queue if i[0] != id]
2023-05-03 00:53:18 +00:00
2023-05-03 00:51:02 +00:00
def wait_pop_find(self, finder, delay=0.1):
found = None
while found is None:
found = self.pop_find(finder)
if found is None:
time.sleep(delay)
return found
2023-05-03 00:53:18 +00:00
class APNSConnection:
2023-05-03 00:51:02 +00:00
incoming_queue = IncomingQueue()
2023-04-07 23:53:21 +00:00
# Sink everything in the queue
def sink(self):
2023-05-03 00:51:02 +00:00
self.incoming_queue = IncomingQueue()
2023-04-07 23:53:21 +00:00
def _queue_filler(self):
while True and not self.sock.closed:
payload = _deserialize_payload(self.sock)
if payload is not None:
# Automatically ACK incoming notifications to prevent APNs from getting mad at us
if payload[0] == 0x0A:
logger.debug("Sending automatic ACK")
self._send_ack(_get_field(payload[1], 4))
2023-07-24 20:37:53 +00:00
logger.debug(f"Received payload: {payload}")
2023-04-07 23:53:21 +00:00
self.incoming_queue.append(payload)
logger.debug(f"Queue length: {len(self.incoming_queue)}")
2023-04-11 16:23:04 +00:00
2023-07-28 15:53:13 +00:00
def _keep_alive_loop(self):
while True and not self.sock.closed:
time.sleep(300)
self._keep_alive()
def __init__(self, private_key=None, cert=None):
# Generate the private key and certificate if they're not provided
if private_key is None or cert is None:
2023-07-24 13:18:21 +00:00
logger.debug("APNs needs a new push certificate")
self.private_key, self.cert = albert.generate_push_cert()
2023-04-06 01:01:07 +00:00
else:
self.private_key, self.cert = private_key, cert
2023-04-06 01:01:07 +00:00
self.sock = _connect(self.private_key, self.cert)
2023-04-06 01:01:07 +00:00
2023-04-11 16:23:04 +00:00
self.queue_filler_thread = threading.Thread(
target=self._queue_filler, daemon=True
)
2023-04-07 23:53:21 +00:00
self.queue_filler_thread.start()
2023-07-28 15:53:13 +00:00
self.keep_alive_thread = threading.Thread(
target=self._keep_alive_loop, daemon=True
)
self.keep_alive_thread.start()
2023-04-07 20:24:05 +00:00
def connect(self, root: bool = True, token: bytes = None):
2023-07-24 13:18:21 +00:00
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})")
2023-04-07 20:24:05 +00:00
flags = 0b01000001
if root:
flags |= 0b0100
2023-04-11 16:23:04 +00:00
if token is None:
2023-04-11 16:23:04 +00:00
payload = _serialize_payload(
2023-05-03 00:53:18 +00:00
7, [(2, 0x01.to_bytes(1, "big")), (5, flags.to_bytes(4, "big"))]
2023-04-11 16:23:04 +00:00
)
else:
2023-04-11 16:23:04 +00:00
payload = _serialize_payload(
2023-05-03 00:53:18 +00:00
7,
[
(1, token),
(2, 0x01.to_bytes(1, "big")),
(5, flags.to_bytes(4, "big")),
],
2023-04-11 16:23:04 +00:00
)
self.sock.write(payload)
2023-05-03 00:51:02 +00:00
payload = self.incoming_queue.wait_pop_find(lambda i: i[0] == 8)
2023-04-11 16:23:04 +00:00
if (
payload == None
or payload[0] != 8
2023-05-03 00:53:18 +00:00
or _get_field(payload[1], 1) != 0x00.to_bytes(1, "big")
2023-04-11 16:23:04 +00:00
):
2023-04-06 01:01:07 +00:00
raise Exception("Failed to connect")
2023-04-11 16:23:04 +00:00
new_token = _get_field(payload[1], 3)
if new_token is not None:
self.token = new_token
elif token is not None:
self.token = token
else:
raise Exception("No token")
2023-07-24 13:18:21 +00:00
logger.debug(f"Recieved connect response with token {b64encode(self.token).decode()}")
2023-04-06 01:01:07 +00:00
2023-04-07 20:24:05 +00:00
return self.token
2023-04-06 14:38:29 +00:00
def filter(self, topics: list[str]):
2023-07-24 13:18:21 +00:00
logger.debug(f"Sending filter message with topics {topics}")
fields = [(1, self.token)]
for topic in topics:
fields.append((2, sha1(topic.encode()).digest()))
payload = _serialize_payload(9, fields)
self.sock.write(payload)
2023-04-11 16:23:04 +00:00
def send_message(self, topic: str, payload: str, id=None):
2023-07-24 13:18:21 +00:00
logger.debug(f"Sending message to topic {topic} with payload {payload}")
if id is None:
id = random.randbytes(4)
2023-04-11 16:23:04 +00:00
payload = _serialize_payload(
0x0A,
[
(4, id),
(1, sha1(topic.encode()).digest()),
(2, self.token),
(3, payload),
],
)
2023-04-07 05:48:07 +00:00
self.sock.write(payload)
2023-05-03 00:51:02 +00:00
# Wait for ACK
payload = self.incoming_queue.wait_pop_find(lambda i: i[0] == 0x0B)
2023-04-08 02:32:00 +00:00
2023-05-03 00:53:18 +00:00
if payload[1][0][1] != 0x00.to_bytes(1, "big"):
2023-04-11 16:23:04 +00:00
raise Exception("Failed to send message")
2023-04-07 05:48:07 +00:00
2023-04-07 23:53:21 +00:00
def set_state(self, state: int):
2023-07-24 13:18:21 +00:00
logger.debug(f"Sending state message with state {state}")
2023-04-11 16:23:04 +00:00
self.sock.write(
_serialize_payload(
2023-05-03 00:53:18 +00:00
0x14,
[(1, state.to_bytes(1, "big")), (2, 0x7FFFFFFF.to_bytes(4, "big"))],
2023-04-11 16:23:04 +00:00
)
)
2023-04-07 23:53:21 +00:00
2023-07-28 15:53:13 +00:00
def _keep_alive(self):
2023-07-24 13:18:21 +00:00
logger.debug("Sending keep alive message")
2023-04-11 16:23:04 +00:00
self.sock.write(_serialize_payload(0x0C, []))
# Remove any keep alive responses we have or missed
self.incoming_queue.remove_all(0x0D)
2023-04-07 05:48:07 +00:00
2023-07-25 23:07:29 +00:00
def _send_ack(self, id: bytes):
logger.debug(f"Sending ACK for message {id}")
payload = _serialize_payload(0x0B, [(1, self.token), (4, id), (8, b"\x00")])
self.sock.write(payload)
2023-05-03 00:51:02 +00:00
# #self.sock.write(_serialize_payload(0x0B, [(4, id)])
# #pass
# def recieve_message(self):
# payload = self.incoming_queue.wait_pop_find(lambda i: i[0] == 0x0A)
# # Send ACK
# self._send_ack(_get_field(payload[1], 4))
# return _get_field(payload[1], 3)
2023-04-07 05:48:07 +00:00
# TODO: Find a way to make this non-blocking
2023-04-11 16:23:04 +00:00
# def expect_message(self) -> tuple[int, list[tuple[int, bytes]]] | None:
2023-04-08 02:32:00 +00:00
# return _deserialize_payload(self.sock)
def _serialize_field(id: int, value: bytes) -> bytes:
2023-05-03 00:53:18 +00:00
return id.to_bytes(1, "big") + len(value).to_bytes(2, "big") + value
2023-04-08 02:32:00 +00:00
def _serialize_payload(id: int, fields: list[(int, bytes)]) -> bytes:
payload = b""
for fid, value in fields:
if fid is not None:
payload += _serialize_field(fid, value)
2023-05-03 00:53:18 +00:00
return id.to_bytes(1, "big") + len(payload).to_bytes(4, "big") + payload
2023-04-08 02:32:00 +00:00
def _deserialize_field(stream: bytes) -> tuple[int, bytes]:
id = int.from_bytes(stream[:1], "big")
length = int.from_bytes(stream[1:3], "big")
value = stream[3 : 3 + length]
return id, value
2023-04-11 16:23:04 +00:00
2023-04-08 02:32:00 +00:00
# Note: Takes a stream, not a buffer, as we do not know the length of the payload
# WILL BLOCK IF THE STREAM IS EMPTY
def _deserialize_payload(stream) -> tuple[int, list[tuple[int, bytes]]] | None:
id = int.from_bytes(stream.read(1), "big")
if id == 0x0:
return None
length = int.from_bytes(stream.read(4), "big")
buffer = stream.read(length)
fields = []
while len(buffer) > 0:
fid, value = _deserialize_field(buffer)
fields.append((fid, value))
buffer = buffer[3 + len(value) :]
return id, fields
2023-04-11 16:23:04 +00:00
def _deserialize_payload_from_buffer(
buffer: bytes,
) -> tuple[int, list[tuple[int, bytes]]] | None:
2023-04-08 02:32:00 +00:00
id = int.from_bytes(buffer[:1], "big")
if id == 0x0:
return None
length = int.from_bytes(buffer[1:5], "big")
buffer = buffer[5:]
if len(buffer) < length:
raise Exception("Buffer is too short")
fields = []
while len(buffer) > 0:
fid, value = _deserialize_field(buffer)
fields.append((fid, value))
buffer = buffer[3 + len(value) :]
return id, fields
# Returns the value of the first field with the given id
def _get_field(fields: list[tuple[int, bytes]], id: int) -> bytes:
for field_id, value in fields:
if field_id == id:
return value
2023-04-11 16:23:04 +00:00
return None