diff --git a/demo.py b/demo.py
index e224acf..d42d1ee 100644
--- a/demo.py
+++ b/demo.py
@@ -1,10 +1,6 @@
 import json
 import logging
-import os
-import threading
-import time
 from base64 import b64decode, b64encode
-from getpass import getpass
 from subprocess import PIPE, Popen
 
 from rich.logging import RichHandler
@@ -74,37 +70,7 @@ async def main():
         await conn.filter(["com.apple.madrid"])
 
         user = ids.IDSUser(conn)
-
-        if CONFIG.get("auth", {}).get("cert") is not None:
-            auth_keypair = ids._helpers.KeyPair(CONFIG["auth"]["key"], CONFIG["auth"]["cert"])
-            user_id = CONFIG["auth"]["user_id"]
-            handles = CONFIG["auth"]["handles"]
-            user.restore_authentication(auth_keypair, user_id, handles)
-        else:
-            username = input("Username: ")
-            password = getpass("Password: ")
-
-            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"),
-        )
-
-        if (
-            CONFIG.get("id", {}).get("cert") is not None
-            and user.encryption_identity is not None
-        ):
-            id_keypair = ids._helpers.KeyPair(CONFIG["id"]["key"], CONFIG["id"]["cert"])
-            user.restore_identity(id_keypair)
-        else:
-            logging.info("Registering new identity...")
-            import emulated.nac
-
-            vd = emulated.nac.generate_validation_data()
-            vd = b64encode(vd).decode()
-
-            user.register(vd)
+        user.auth_and_set_encryption_from_config(CONFIG)
 
         # Write config.json
         CONFIG["encryption"] = {
@@ -150,4 +116,4 @@ async def output_task(im: imessage.iMessageUser):
 
 
 if __name__ == "__main__":
-    trio.run(main)
\ No newline at end of file
+    trio.run(main)
diff --git a/emulated/jelly.py b/emulated/jelly.py
index 5bf6668..b4b11ad 100644
--- a/emulated/jelly.py
+++ b/emulated/jelly.py
@@ -353,4 +353,4 @@ def c_string(bytes, start: int = 0) -> str:
         #print(start)
         #print(chr(bytes[i]))
         i += 1
-    return out
\ No newline at end of file
+    return out
diff --git a/emulated/mparser.py b/emulated/mparser.py
index ba34b14..14b544a 100644
--- a/emulated/mparser.py
+++ b/emulated/mparser.py
@@ -23,12 +23,12 @@ from datetime import datetime
 from json import dump
 from math import exp, log
 from os import SEEK_END
-from re import split
 from struct import unpack
 from uuid import UUID
+from typing import Any
 
-#from asn1crypto.cms import ContentInfo
-#from asn1crypto.x509 import DirectoryString
+from asn1crypto.cms import ContentInfo
+from asn1crypto.x509 import DirectoryString
 from plistlib import loads
 
 #import mdictionary as mdictionary
@@ -52,7 +52,7 @@ class Parser():
         self.__is_64_bit = True         # default place-holder
         self.__is_little_endian = True  # ^^
         self.__macho = {}
-        self.__output = {
+        self.__output: dict[str, Any] = {
             'name': 'IMDAppleServices'
         }
 
@@ -931,7 +931,7 @@ class Parser():
 
             n_value = self.get_ll() if self.__is_64_bit else self.get_int()
 
-            symbol = {
+            symbol: dict[str, int | str] = {
                 'n_strx': n_strx,
                 'n_sect': n_sect,
                 'n_desc': n_desc,
@@ -2298,4 +2298,4 @@ class mdictionary:
             2147483648 + 11: 'POWERPC_7450 (LIB64)',
             2147483648 + 100: 'POWERPC_970 (LIB64)'
         }
-    }
\ No newline at end of file
+    }
diff --git a/ids/__init__.py b/ids/__init__.py
index 5c78b55..c21e923 100644
--- a/ids/__init__.py
+++ b/ids/__init__.py
@@ -1,4 +1,6 @@
 from base64 import b64encode
+from getpass import getpass
+import logging
 
 import apns
 
@@ -27,6 +29,9 @@ class IDSUser:
         self._push_keypair = _helpers.KeyPair(
             self.push_connection.credentials.private_key, self.push_connection.credentials.cert
         )
+        # set the encryption_identity to a default randomized value so that
+        # it's still valid if we can't pull it from the config
+        self.encryption_identity: identity.IDSIdentity = identity.IDSIdentity()
 
         self.ec_key = self.rsa_key = None
 
@@ -63,10 +68,6 @@ class IDSUser:
         self.ec_key, self.rsa_key will be set to a randomly gnenerated EC and RSA keypair
         if they are not already set
         """
-        if self.encryption_identity is None:
-            self.encryption_identity = identity.IDSIdentity()
-        
-        
         cert = identity.register(
             b64encode(self.push_connection.credentials.token),
             self.handles,
@@ -81,6 +82,48 @@ class IDSUser:
     def restore_identity(self, id_keypair: _helpers.KeyPair):
         self._id_keypair = id_keypair
 
+    def auth_and_set_encryption_from_config(self, config: dict[str, dict[str, Any]]):
+
+        auth = config.get("auth", {})
+        if (
+			((key := auth.get("key")) is not None) and
+			((cert := auth.get("cert")) is not None) and
+			((user_id := auth.get("user_id")) is not None) and
+			((handles := auth.get("handles")) is not None)
+		):
+            auth_keypair = _helpers.KeyPair(key, cert)
+            self.restore_authentication(auth_keypair, user_id, handles)
+        else:
+            username = input("Username: ")
+            password = getpass("Password: ")
+
+            self.authenticate(username, password)
+
+        encryption: dict[str, str] = config.get("encryption", {})
+        id: dict[str, str] = config.get("id", {})
+
+        if (
+            (rsa_key := encryption.get("rsa_key")) and
+            (signing_key := encryption.get("ec_key")) and
+            (cert := id.get("cert")) and
+            (key := id.get("key"))
+        ):
+            self.encryption_identity = identity.IDSIdentity(
+                encryption_key=rsa_key,
+                signing_key=signing_key,
+            )
+
+            id_keypair = _helpers.KeyPair(key, cert)
+            self.restore_identity(id_keypair)
+        else:
+            logging.info("Registering new identity...")
+            import emulated.nac
+
+            vd = emulated.nac.generate_validation_data()
+            vd = b64encode(vd).decode()
+
+            self.register(vd)
+
     async def lookup(self, uris: list[str], topic: str = "com.apple.madrid") -> Any:
         return await query.lookup(self.push_connection, self.current_handle, self._id_keypair, uris, topic)
         
diff --git a/ids/identity.py b/ids/identity.py
index f7ee5e3..ac46813 100644
--- a/ids/identity.py
+++ b/ids/identity.py
@@ -3,18 +3,26 @@ from base64 import b64decode
 
 import requests
 
-from ._helpers import PROTOCOL_VERSION, USER_AGENT, KeyPair, parse_key, serialize_key
+from ._helpers import PROTOCOL_VERSION, KeyPair, parse_key, serialize_key
 from .signing import add_auth_signature, armour_cert
 
 from io import BytesIO
 
 from cryptography.hazmat.primitives.asymmetric import ec, rsa
 
+from typing import Self
+
 import logging
 logger = logging.getLogger("ids")
 
 class IDSIdentity:
-    def __init__(self, signing_key: str | None = None, encryption_key: str | None = None, signing_public_key: str | None = None, encryption_public_key: str | None = None):
+    def __init__(
+		self,
+		signing_key: str | None = None,
+		encryption_key: str | None = None,
+		signing_public_key: str | None = None,
+		encryption_public_key: str | None = None
+	):
         if signing_key is not None:
             self.signing_key = signing_key
             self.signing_public_key = serialize_key(parse_key(signing_key).public_key())# type: ignore
@@ -36,8 +44,8 @@ class IDSIdentity:
             self.encryption_key = serialize_key(rsa.generate_private_key(65537, 1280))
             self.encryption_public_key = serialize_key(parse_key(self.encryption_key).public_key())# type: ignore
     
-    @staticmethod
-    def decode(inp: bytes) -> 'IDSIdentity':
+    @classmethod
+    def decode(cls, inp: bytes) -> Self:
         input = BytesIO(inp)
 
         assert input.read(5) == b'\x30\x81\xF6\x81\x43' # DER header
diff --git a/ids/profile.py b/ids/profile.py
index f5ae576..6de773e 100644
--- a/ids/profile.py
+++ b/ids/profile.py
@@ -7,13 +7,13 @@ import requests
 from cryptography import x509
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes, serialization
-from cryptography.hazmat.primitives.asymmetric import padding, rsa
+from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.x509.oid import NameOID
 
 import bags
 
 from . import signing
-from ._helpers import PROTOCOL_VERSION, USER_AGENT, KeyPair
+from ._helpers import PROTOCOL_VERSION, KeyPair
 
 import logging
 logger = logging.getLogger("ids")
@@ -50,8 +50,6 @@ def _auth_token_request(username: str, password: str) -> Any:
 def get_auth_token(
     username: str, password: str, factor_gen: Callable | None = None
 ) -> tuple[str, str]:
-    from sys import platform
-    
     result = _auth_token_request(username, password)
     if result["status"] != 0:
         if result["status"] == 5000:
diff --git a/ids/query.py b/ids/query.py
index 716687b..20e2535 100644
--- a/ids/query.py
+++ b/ids/query.py
@@ -5,7 +5,6 @@ from base64 import b64encode
 
 import apns
 import bags
-import logging
 
 from ._helpers import KeyPair, PROTOCOL_VERSION
 from . import signing
diff --git a/ids/signing.py b/ids/signing.py
index 71550ff..f03ef07 100644
--- a/ids/signing.py
+++ b/ids/signing.py
@@ -5,8 +5,7 @@ from datetime import datetime
 from cryptography import x509
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes, serialization
-from cryptography.hazmat.primitives.asymmetric import padding, rsa
-from cryptography.x509.oid import NameOID
+from cryptography.hazmat.primitives.asymmetric import padding
 
 from ._helpers import KeyPair, dearmour
 
@@ -24,8 +23,6 @@ Generates a nonce in this format:
    000001876d008cc5                 # unix time
                    r1r2r3r4r5r6r7r8 # random bytes
 """
-
-
 def generate_nonce() -> bytes:
     return (
         b"\x01"
@@ -33,17 +30,13 @@ def generate_nonce() -> bytes:
         + random.randbytes(8)
     )
 
-
-import typing
-
-
 # Creates a payload from individual parts for signing
 def _create_payload(
     bag_key: str,
     query_string: str,
-    push_token: typing.Union[str, bytes],
+    push_token: str | bytes,
     payload: bytes,
-    nonce: typing.Union[bytes, None] = None,
+    nonce: bytes | None = None,
 ) -> tuple[bytes, bytes]:
     # Generate the nonce
     if nonce is None:
diff --git a/imessage.py b/imessage.py
index 02c146f..55aa7a7 100644
--- a/imessage.py
+++ b/imessage.py
@@ -4,9 +4,10 @@ import logging
 import plistlib
 import random
 import uuid
-from dataclasses import dataclass, field
-from hashlib import sha1, sha256
+from dataclasses import dataclass
+from hashlib import sha256
 from io import BytesIO
+from typing import Any
 
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.asymmetric import ec, padding
@@ -504,15 +505,15 @@ class iMessageUser:
         """
         Will return the next iMessage in the queue, or None if there are no messages
         """
-        body = await self._receive_raw([t for t, _ in MESSAGE_TYPES.items()], [t[0] for _, t in MESSAGE_TYPES.items()])
-        t = MESSAGE_TYPES[body["c"]][1]   
+        body: dict[str, Any] = await self._receive_raw(list(MESSAGE_TYPES.keys()), [t[0] for t in MESSAGE_TYPES.values()])
+        t: type[Message] = MESSAGE_TYPES[body["c"]][1]   
 
         if not await self._verify_payload(body["P"], body["sP"], body["t"]):
             raise Exception("Failed to verify payload")
 
         logger.debug(f"Encrypted body : {body}")
 
-        decrypted = self._decrypt_payload(body["P"])
+        decrypted: bytes = self._decrypt_payload(body["P"])
 
         try:
             return t.from_raw(decrypted, body["sP"])
@@ -627,7 +628,7 @@ class iMessageUser:
 
         await self.connection.send_notification(topic, body, message_id)
 
-    async def _receive_raw(self, c: int | list[int], topics: str | list[str]) -> dict:
+    async def _receive_raw(self, c: int | list[int], topics: str | list[str]) -> dict[str, Any]:
         def check(payload: apns.APNSPayload):
             # Check if the "c" key matches
             body = payload.fields_with_id(3)[0].value
@@ -644,8 +645,8 @@ class iMessageUser:
         
         payload = await self.connection.expect_notification(topics, check)
 
-        body = payload.fields_with_id(3)[0].value
-        body = plistlib.loads(body)
+        body_bytes: bytes = payload.fields_with_id(3)[0].value
+        body: dict[str, Any] = plistlib.loads(body_bytes)
         return body
 
     async def activate_sms(self):
@@ -655,14 +656,12 @@ class iMessageUser:
         Call repeatedly until it returns True
         """
 
-        act_message = await self._receive_raw(145, "com.apple.private.alloy.sms")
-        if act_message is None:
-            return False
+        act_message: dict[str, Any] = await self._receive_raw(145, "com.apple.private.alloy.sms")
         
         logger.info(f"Received SMS activation message : {act_message}")
         # Decrypt the payload
-        act_message = self._decrypt_payload(act_message["P"])
-        act_message = plistlib.loads(maybe_decompress(act_message))
+        act_message_bytes: bytes = self._decrypt_payload(act_message["P"])
+        act_message = plistlib.loads(maybe_decompress(act_message_bytes))
 
         if act_message == {'wc': False, 'ar': True}:
             logger.info("SMS forwarding activated, sending response")
@@ -715,7 +714,7 @@ class iMessageUser:
                 total += 1
 
         while count < total and time.time() - start < 2:
-            resp = await self._receive_raw(255, topic)
+            resp: dict[str, Any] = await self._receive_raw(255, topic)
             #if resp is None:
             #    continue
             count += 1
diff --git a/proxy/printer.py b/proxy/printer.py
index 3eaf822..2763e42 100644
--- a/proxy/printer.py
+++ b/proxy/printer.py
@@ -2,6 +2,7 @@ import plistlib
 import zlib
 from base64 import b64decode, b64encode
 from hashlib import sha1
+from typing import Any
 
 topics = {'com.apple.private.alloy.notes', 'com.apple.icloud-container.clouddocs.F3LWYJ7GM7.com.apple.garageband10', 'com.apple.private.alloy.screentime', 'com.apple.icloud-container.com.apple.appleaccount.custodian', 'com.apple.icloud-container.clouddocs.iCloud.com.apple.configurator.ui', 'com.apple.icloud-container.com.apple.VoiceMemos', 'com.apple.icloud-container.com.apple.SafariShared.Settings', 'com.apple.private.alloy.status.keysharing', 'com.apple.private.alloy.electrictouch', 'com.apple.private.alloy.icloudpairing', 'com.apple.icloud.presence.shared.experience', 'com.apple.icloud-container.com.apple.knowledge-agent', 'com.apple.private.alloy.thumper.keys', 'com.apple.pay.services.ck.zone.prod', 'com.apple.sharedstreams', 'com.apple.jalisco', 'com.apple.private.alloy.ded', 'com.apple.icloud-container.com.apple.cloudpaird', 'com.apple.private.alloy.multiplex1', 'com.apple.private.alloy.nearby', 'com.me.contacts', 'com.apple.TestFlight', 'com.icloud.family', 'com.apple.icloud-container.com.apple.iWork.Pages', 'com.apple.bookassetd', 'com.apple.tv.favoriteTeams', 'com.apple.icloud-container.com.apple.Safari', 'com.apple.mobileme.fmf3', 'com.apple.icloud-container.clouddocs.iCloud.com.apple.iBooks.iTunesU', 'com.apple.private.alloy.applepay', 'com.apple.private.alloy.willow', 'com.apple.idmsauth', 'com.apple.icloud-container.com.apple.iWork.Numbers', 'com.apple.icloud-container.clouddocs.F3LWYJ7GM7.com.apple.mobilegarageband', 'com.apple.private.alloy.maps', 'com.apple.private.alloy.phonecontinuity', 'com.apple.private.alloy.avconference.icloud', 'com.apple.pay.services.apply.prod', 'com.apple.private.alloy.facetime.multi', 'com.apple.icloud-container.clouddocs.com.apple.TextInput', 'com.apple.icloud-container.clouddocs.iCloud.com.reddit.reddit', 'com.apple.icloud-container.clouddocs.com.apple.Numbers', 'com.apple.icloud.fmip.voiceassistantsync', 'com.apple.icloud-container.com.apple.avatarsd', 'com.apple.private.ac', 'company.thebrowser.Browser', 'com.apple.itunesstored', 'com.apple.icloud-container.com.apple.icloud.fmfd', 'com.apple.private.alloy.screentime.invite', 'com.apple.icloud-container.com.apple.donotdisturbd', 'com.apple.icloud-container.clouddocs.com.apple.TextEdit', 'com.apple.appstored', 'com.apple.icloud-container.clouddocs.com.apple.CloudDocs.container-metadata', 'com.apple.private.alloy.screensharing', 'com.apple.private.alloy.accessibility.switchcontrol', 'com.apple.private.alloy.screensharing.qr', 'com.apple.private.alloy.amp.potluck', 'com.apple.icloud-container.com.apple.siriknowledged', 'com.apple.private.alloy.gamecenter', 'com.apple.appstored-testflight', 'com.apple.private.alloy.messagenotification', 'com.apple.passd.usernotifications', 'com.apple.icloud-container.clouddocs.com.apple.Pages', 'com.apple.private.alloy.safeview', 'com.apple.findmy', 'com.apple.pay.auxiliary.registration.requirement.prod', 'com.apple.aa.idms', 'com.apple.private.alloy.ids.cloudmessaging', 'com.apple.icloud-container.com.apple.icloud.searchpartyuseragent', 'com.icloud.quota', 'com.apple.icloud-container.com.apple.upload-request-proxy.com.apple.photos.cloud', 'com.apple.private.alloy.usagetracking', 'com.apple.icloud-container.com.apple.syncdefaultsd', 'com.apple.private.alloy.continuity.tethering', 'com.apple.idmsauthagent', 'com.apple.sagad', 'com.apple.pay.services.ownershipTokens.prod', 'com.apple.private.alloy.sms', 'com.apple.Notes', 'com.apple.icloud-container.com.apple.SafariShared.WBSCloudBookmarksStore', 'com.apple.icloud-container.com.apple.reminders', 'com.apple.private.alloy.classroom', 'com.apple.news', 'com.apple.icloud-container.com.apple.imagent', 'com.apple.pay.services.products.prod', 'com.apple.private.alloy.fmf', 'com.apple.amsaccountsd', 'com.apple.private.alloy.itunes', 'com.apple.icloud-container.clouddocs.iCloud.com.apple.iBooks', 'com.apple.private.alloy.gelato', 'com.apple.icloud-container.com.apple.willowd', 'com.apple.icloud-container.clouddocs.com.apple.CloudDocs', 'com.apple.icloud-container.com.apple.protectedcloudstorage.protectedcloudkeysyncing', 'com.apple.icloud-container.com.apple.Notes', 'com.me.cal', 'com.apple.peerpayment', 'com.apple.icloud-container.clouddocs.iCloud.is.workflow.my.workflows', 'com.apple.private.alloy.facetime.sync', 'com.apple.icloud-container.com.apple.news', 'com.apple.icloud-container.com.apple.TrustedPeersHelper', 'com.apple.private.alloy.home.invite', 'com.apple.private.alloy.coreduet.sync', 'com.apple.private.alloy.contextsync', 'com.apple.private.alloy.fmd', 'com.apple.private.alloy.status.personal', 'com.apple.icloud-container.com.apple.assistant.assistantd', 'com.apple.private.alloy.sleep.icloud', 'com.apple.icloud-container.com.apple.security.cuttlefish', 'com.apple.wallet.sharing', 'com.apple.icloud-container.clouddocs.3L68KQB4HG.com.readdle.CommonDocuments', 'com.apple.pay.provision', 'com.apple.icloud-container.com.apple.StatusKitAgent', 'com.apple.icloud-container.clouddocs.com.apple.Preview', 'com.apple.icloud-container.com.apple.gamed', 'com.apple.askpermissiond', 'com.apple.private.alloy.gamecenter.imessage', 'com.apple.private.alloy.safari.groupactivities', 'com.apple.icloud-container.com.apple.Maps', 'com.apple.private.alloy.willow.stream', 'com.apple.pay.services.devicecheckin.prod.us', 'com.apple.icloud.presence.mode.status', 'com.apple.ess', 'com.apple.private.alloy.accounts.representative', 'com.apple.icloud-container.clouddocs.com.apple.QuickTimePlayerX', 'com.apple.private.alloy.facetime.audio', 'com.apple.private.alloy.continuity.unlock', 'com.apple.icloud-container.clouddocs.iCloud.md.obsidian', 'com.apple.icloud-container.clouddocs.iCloud.com.apple.MobileSMS', 'com.apple.iWork.Numbers', 'com.apple.pay.services.account.prod', 'com.apple.private.alloy.quickrelay', 'com.apple.iBooksX', 'com.apple.madrid', 'com.apple.private.alloy.continuity.activity', 'com.apple.icloud-container.com.apple.keyboardservicesd', 'com.apple.icloud-container.clouddocs.com.apple.CloudDocs.health', 'com.apple.icloud-container.com.apple.suggestd', 'com.apple.icloud-container.clouddocs.com.apple.Keynote', 'com.apple.private.alloy.home', 'com.apple.private.alloy.photostream', 'com.apple.icloud-container.com.apple.iBooksX', 'com.apple.private.alloy.digitalhealth', 'com.apple.icloud-container.clouddocs.iCloud.dk.simonbs.Scriptable', 'com.apple.private.alloy.copresence', 'com.apple.private.alloy.continuity.encryption', 'com.apple.icloud-container.com.apple.passd', 'com.apple.icloud-container.com.apple.findmy', 'com.apple.icloud-container.com.apple.financed', 'com.apple.icloud-container.com.apple.photos.cloud', 'com.apple.private.alloy.proxiedcrashcopier.icloud', 'com.apple.private.alloy.tips', 'com.apple.icloud-container.com.apple.appleaccount.beneficiary.private', 'com.apple.watchList', 'com.apple.icloud-container.com.apple.willowd.homekit', 'com.apple.icloud-container.clouddocs.com.apple.CloudDocs.pp-metadata', 'com.apple.icloud-container.com.apple.SafariShared.CloudTabs', 'com.apple.private.alloy.facetime.lp', 'com.apple.icloud-container.com.apple.appleaccount.beneficiary', 'com.apple.aa.setupservice', 'com.apple.icloud.fmip.app.push', 'com.apple.icloud.presence.channel.management', 'com.apple.icloud-container.clouddocs.com.apple.ScriptEditor2', 'com.apple.private.alloy.facetime.mw', 'com.apple.Maps', 'com.apple.icloud-container.clouddocs.com.apple.mail', 'com.apple.mobileme.fmf2', 'com.me.setupservice', 'paymentpass.com.apple', 'com.apple.music.social', 'com.apple.icloud-container.clouddocs.com.apple.iBooks.cloudData', 'com.apple.iWork.Pages', 'com.apple.private.alloy.carmelsync', 'com.apple.private.alloy.maps.eta', 'com.apple.icloud-container.clouddocs.com.apple.shoebox', 'com.apple.dt.Xcode', 'com.apple.private.alloy.facetime.video', 'com.apple.icloud-container.com.apple.sociallayerd', 'com.apple.private.alloy.keytransparency.accountkey.pinning', 'com.apple.wallet.sharing.qa', 'com.apple.icloud-container.com.apple.appleaccount.custodian.private', 'com.apple.private.alloy.phone.auth', 'com.apple.icloud-container.com.apple.amsengagementd', 'com.apple.amsengagementd.notifications', 'com.apple.maps.icloud', 'com.apple.storekit', 'com.apple.triald', 'com.icloud.askpermission', 'com.apple.private.alloy.biz', 'com.apple.tilt', 'com.apple.icloud-container.com.apple.callhistory.sync-helper', 'com.apple.private.ids', 'com.apple.private.alloy.clockface.sharing', 'com.apple.gamed', 'com.apple.icloud-container.company.thebrowser.Browser', 'com.apple.icloud-container.com.apple.securityd'}
 
@@ -104,7 +105,8 @@ def print_payload(payload: apns.APNSPayload, to_server: bool):
     pretty_print_payload(prefix, (payload.id, f))
 
 def pretty_print_payload(
-    prefix, payload: tuple[int, list[tuple[int, bytes]]]
+    prefix,
+	payload: tuple[int, list[tuple[int, bytes]]],
 ) -> bytes | None:
     id = payload[0]
 
@@ -215,39 +217,34 @@ def pretty_print_payload(
         if topic == "com.apple.madrid" or topic == "com.apple.private.alloy.sms":
             print(f" {bcolors.FAIL}Madrid{bcolors.ENDC}", end="")
             orig_payload = payload
-            payload = plistlib.loads(_get_field(payload[1], 3))
+            payload: dict[str, Any] = plistlib.loads(_get_field(payload[1], 3))
 
             # print(payload)
-            if "cT" in payload and False:
+            if (cT := payload.get("cT")) is not None:
                 # It's HTTP over APNs
-                if "hs" in payload:
-                    print(
-                        f" {bcolors.WARNING}HTTP Response{bcolors.ENDC}: {payload['hs']}",
-                        end="",
-                    )
+                if (hs := payload.get("hs")) is not None:
+                    print(f" {bcolors.WARNING}HTTP Response{bcolors.ENDC}: {hs}")
                 else:
-                    print(f" {bcolors.WARNING}HTTP Request{bcolors.ENDC}", end="")
-                # print(f" {bcolors.WARNING}HTTP{bcolors.ENDC} {payload['hs']}", end="")
-                if "u" in payload:
-                    print(f" {bcolors.OKCYAN}URL{bcolors.ENDC}: {payload['u']}", end="")
-                print(
-                    f" {bcolors.FAIL}Content Type{bcolors.ENDC}: {payload['cT']}",
-                    end="",
-                )
-                if "h" in payload:
-                    print(
-                        f" {bcolors.FAIL}Headers{bcolors.ENDC}: {payload['h']}", end=""
-                    )
-                if "b" in payload:
+                    print(f" {bcolors.WARNING}HTTP Request{bcolors.ENDC}")
+
+                if (u := payload.get("u")) is not None:
+                    print(f" {bcolors.OKCYAN}URL{bcolors.ENDC}: {u}")
+
+                print(f" {bcolors.FAIL}Content Type{bcolors.ENDC}: {cT}")
+
+                if (h := payload.get("h")) is not None:
+                    print(f" {bcolors.FAIL}Headers{bcolors.ENDC}: {h}")
+
+                if (body := payload.get("b")) is not None:
                     # What am I really supposed to put in WBITS? Got this from a random SO answer
                     # print(payload["b"])
-                    body = zlib.decompress(payload["b"], 16 + zlib.MAX_WBITS)
-                    if b"plist" in body:
-                        body = plistlib.loads(body)
-                    print(f" {bcolors.FAIL}Body{bcolors.ENDC}: {body}", end="")
-            #if not "cT" in payload:
-            for key in payload:
-                print(f" {bcolors.OKBLUE}{key}{bcolors.ENDC}: {payload[key]}") # type: ignore
+                    body_bytes: bytes = zlib.decompress(body, 16 + zlib.MAX_WBITS)
+                    if b"plist" in body_bytes:
+                        body_bytes = plistlib.loads(body_bytes)
+                    print(f" {bcolors.FAIL}Body{bcolors.ENDC}: {body_bytes}", end="")
+            else:
+                for (key, value) in payload.items():
+                    print(f" {bcolors.OKBLUE}{key}{bcolors.ENDC}: {value}") # type: ignore
 
             if 'dtl' in payload and False:
                 print("OVERRIDE DTL")
diff --git a/proxy/proxy.py b/proxy/proxy.py
index 882a06d..d8b15be 100644
--- a/proxy/proxy.py
+++ b/proxy/proxy.py
@@ -1,5 +1,6 @@
 import os
 import sys
+import traceback
 
 # setting path so we can import the needed packages
 sys.path.append(os.path.join(sys.path[0], "../"))
@@ -39,8 +40,10 @@ async def handle_proxy(stream: trio.SocketStream):
     try:
         p = APNSProxy(stream)
         await p.start()
-    except Exception as e:
-        logging.error("APNSProxy instance encountered exception: " + str(e))
+    except Exception:
+        logging.error(f"APNSProxy instance encountered exception:")
+        traceback.print_exc()
+
         #raise e
 
 class APNSProxy:
@@ -54,7 +57,7 @@ class APNSProxy:
                 try:
                     apns_server = apns.APNSConnection(nursery)
                     await apns_server._connect_socket()
-                    self.server = apns_server.sock
+                    self.connection = apns_server
 
                     nursery.start_soon(self.proxy, True)
                     nursery.start_soon(self.proxy, False)
@@ -69,10 +72,11 @@ class APNSProxy:
     async def proxy(self, to_server: bool):
         if to_server:
             from_stream = self.client
-            to_stream = self.server
+            to_stream = self.connection.sock
         else:
-            from_stream = self.server
+            from_stream = self.connection.sock
             to_stream = self.client
+
         while True:
             payload = await apns.APNSPayload.read_from_stream(from_stream)
             payload = self.tamper(payload, to_server)
@@ -95,15 +99,15 @@ class APNSProxy:
     def tamper_lookup_keys(self, payload: apns.APNSPayload) -> apns.APNSPayload:
         if payload.id == 0xA: # Notification
             if payload.fields_with_id(2)[0].value == sha1(b"com.apple.madrid").digest(): # Topic
-                if body := payload.fields_with_id(3)[0].value is not None:
+                if (body := payload.fields_with_id(3)[0].value) is not None:
                     body = plistlib.loads(body)
                     if body['c'] == 97: # Lookup response
                         resp = gzip.decompress(body["b"]) # HTTP body
                         resp = plistlib.loads(resp)
 
                         # Replace public keys
-                        for r in resp["results"].keys():
-                            for identity in resp["results"][r]["identities"]:
+                        for result in resp["results"].values():
+                            for identity in result["identities"]:
                                 if "client-data" in identity:
                                     identity["client-data"]["public-message-identity-key"] = b"REDACTED"
                         
@@ -117,4 +121,4 @@ class APNSProxy:
         return payload
 
 if __name__ == "__main__":
-    trio.run(main)
\ No newline at end of file
+    trio.run(main)