From ed9eb28097c95cfb93fbb6ce836c5af71a648524 Mon Sep 17 00:00:00 2001
From: JJTech0130 <jjtech@jjtech.dev>
Date: Sat, 8 Apr 2023 16:01:03 -0400
Subject: [PATCH] getting ready for imessage decryption...

---
 printer.py     | 7 +++++++
 proxy/proxy.py | 9 ++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/printer.py b/printer.py
index f0323c6..402d422 100644
--- a/printer.py
+++ b/printer.py
@@ -153,8 +153,15 @@ def pretty_print_payload(prefix, payload: tuple[int, list[tuple[int, bytes]]]) -
         # if it has apsd -> APNs in the prefix, it's an outgoing notification
         if "apsd -> APNs" in prefix:
             print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKBLUE}OUTGOING Notification{bcolors.ENDC}", end="")
+            topic = _lookup_topic(_get_field(payload[1], 1))
             if b"bplist" in _get_field(payload[1], 3):
                 print(f" {bcolors.OKCYAN}Binary{bcolors.ENDC}", end="")
+            if topic == "com.apple.madrid":
+                print(f" {bcolors.FAIL}Madrid{bcolors.ENDC}", end="")
+                import plistlib
+                plist = plistlib.loads(_get_field(payload[1], 3))
+                for key in plist:
+                    print(f" {bcolors.OKBLUE}{key}{bcolors.ENDC}: {plist[key]}", end="")
             print(f" {bcolors.WARNING}Topic{bcolors.ENDC}: {_lookup_topic(_get_field(payload[1], 1))}")
             
         else: 
diff --git a/proxy/proxy.py b/proxy/proxy.py
index 72f86fa..3cec0db 100644
--- a/proxy/proxy.py
+++ b/proxy/proxy.py
@@ -61,7 +61,14 @@ def proxy(conn1: tlslite.TLSConnection, conn2: tlslite.TLSConnection, prefix: st
                 print(prefix, "Connection closed due to no data")
                 break
 
-            override = printer.pretty_print_payload(prefix, apns._deserialize_payload_from_buffer(data))
+            try:
+                override = printer.pretty_print_payload(prefix, apns._deserialize_payload_from_buffer(data))
+            except Exception as e:
+                if e == "Buffer is too short":
+                    print("Buffer is too short to parse, ignoring")
+                   # print("Buffer is too short, trying to read more")
+                   # data += conn1.read()
+                    pass
             if override is not None:
                 data = override
                 print("OVERRIDE: ", end="")