clean up the demo a bit

apns.py

@@ -107,6 +107,11 @@ class APNSConnection:
                 self.incoming_queue.append(payload)
                 logger.debug(f"Queue length: {len(self.incoming_queue)}")
 
+    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:
@@ -117,12 +122,17 @@ class APNSConnection:
 
         self.sock = _connect(self.private_key, self.cert)
 
-        # Start the queue filler thread
         self.queue_filler_thread = threading.Thread(
             target=self._queue_filler, daemon=True
         )
         self.queue_filler_thread.start()
 
+        self.keep_alive_thread = threading.Thread(
+            target=self._keep_alive_loop, daemon=True
+        )
+        self.keep_alive_thread.start()
+
+
     def connect(self, root: bool = True, token: bytes = None):
         if token is None:
             logger.debug(f"Sending connect message without token (root={root})")
@@ -212,7 +222,7 @@ class APNSConnection:
             )
         )
 
-    def keep_alive(self):
+    def _keep_alive(self):
         logger.debug("Sending keep alive message")
         self.sock.write(_serialize_payload(0x0C, []))
         # Remove any keep alive responses we have or missed

demo.py

@@ -1,20 +1,15 @@
-import gzip
 import json
 import logging
-import plistlib
 import threading
 import time
 from base64 import b64decode, b64encode
 from getpass import getpass
 
-from cryptography.hazmat.primitives import hashes, serialization
-from cryptography.hazmat.primitives.asymmetric import padding, rsa
-from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
-from cryptography.hazmat.primitives.serialization import load_pem_private_key
 from rich.logging import RichHandler
 
 import apns
 import ids
+import imessage
 
 logging.basicConfig(
     level=logging.NOTSET, format="%(message)s", datefmt="[%X]", handlers=[RichHandler()]
@@ -65,7 +60,10 @@ else:
 
     user.authenticate(username, password)
 
-user.encryption_identity = ids.identity.IDSIdentity(encryption_key=CONFIG.get("encryption", {}).get("rsa_key"), signing_key=CONFIG.get("encryption", {}).get("ec_key"))
+user.encryption_identity = ids.identity.IDSIdentity(
+    encryption_key=CONFIG.get("encryption", {}).get("rsa_key"),
+    signing_key=CONFIG.get("encryption", {}).get("ec_key"),
+)
 
 if (
     CONFIG.get("id", {}).get("cert") is not None
@@ -84,17 +82,6 @@ else:
 
 logging.info("Waiting for incoming messages...")
 
-# Create a thread to send keepalive messages
-
-
-def keepalive():
-    while True:
-        time.sleep(300)
-        conn.keep_alive()
-
-
-threading.Thread(target=keepalive, daemon=True).start()
-
 # Write config.json
 CONFIG["encryption"] = {
     "rsa_key": user.encryption_identity.encryption_key,
@@ -119,12 +106,8 @@ CONFIG["push"] = {
 with open("config.json", "w") as f:
     json.dump(CONFIG, f, indent=4)
 
-import imessage
 im = imessage.iMessageUser(conn, user)
 
-#import time
-#time.sleep(4)
-#onn._send_ack(b'\t-\x97\x96')
 while True:
     msg = im.receive()
     if msg is not None: