From d30c397300f3ca7dbe97bf5b87069bd27308b14b Mon Sep 17 00:00:00 2001
From: JJTech0130 <jjtech@jjtech.dev>
Date: Fri, 7 Apr 2023 10:33:03 -0400
Subject: [PATCH] implement overrides

---
 apns.py        |  3 ++-
 printer.py     | 27 +++++++++++++++++++++------
 proxy/proxy.py | 27 +++++++++++++++++----------
 3 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/apns.py b/apns.py
index a66e85e..5592aa7 100644
--- a/apns.py
+++ b/apns.py
@@ -12,7 +12,8 @@ def _serialize_payload(id: int, fields: list[(int, bytes)]) -> bytes:
     payload = b""
 
     for fid, value in fields:
-        payload += _serialize_field(fid, value)
+        if fid is not None:
+            payload += _serialize_field(fid, value)
 
     return id.to_bytes() + len(payload).to_bytes(4, "big") + payload
 
diff --git a/printer.py b/printer.py
index 11884e2..48740cd 100644
--- a/printer.py
+++ b/printer.py
@@ -93,10 +93,11 @@ def _p_filter(prefix, fields: list[tuple[int, bytes]]):
     if len(paused) > 100:
         paused = paused[:100] + "..."
     # (Token: {token.decode()})
-    print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Filter{bcolors.ENDC} {bcolors.WARNING}Enabled:{bcolors.ENDC} {enabled} {bcolors.FAIL}Ignored:{bcolors.ENDC} {ignored} {bcolors.OKBLUE}Oppertunistic:{bcolors.ENDC} {oppertunistic} {bcolors.OKGREEN}Paused:{bcolors.ENDC} {paused}")       
+    print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Filter{bcolors.ENDC} {bcolors.WARNING}Enabled{bcolors.ENDC}: {enabled} {bcolors.FAIL}Ignored{bcolors.ENDC}: {ignored} {bcolors.OKBLUE}Oppertunistic{bcolors.ENDC}: {oppertunistic} {bcolors.OKGREEN}Paused{bcolors.ENDC}: {paused}")       
 
+import apns
 
-def pretty_print_payload(prefix, payload: tuple[int, list[tuple[int, bytes]]]):
+def pretty_print_payload(prefix, payload: tuple[int, list[tuple[int, bytes]]]) -> bytes | None:
     id = payload[0]
 
     if id == 9:
@@ -104,10 +105,24 @@ def pretty_print_payload(prefix, payload: tuple[int, list[tuple[int, bytes]]]):
     elif id == 8:
         token_str = ""
         if _get_field(payload[1], 3):
-            token_str = f"New token: {b64encode(_get_field(payload[1], 3)).decode()}"
-        print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Connected{bcolors.ENDC} {token_str}")
+            token_str = f"{bcolors.WARNING}Token{bcolors.ENDC}: {b64encode(_get_field(payload[1], 3)).decode()}"
+        print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Connected{bcolors.ENDC} {token_str} {bcolors.OKBLUE}{_get_field(payload[1], 1).hex()}{bcolors.ENDC}")
     elif id == 7:
-        print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Connect Request{bcolors.ENDC}")
+        print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Connect Request{bcolors.ENDC}", end="")
+        if _get_field(payload[1], 1):
+            print(f" {bcolors.WARNING}Token{bcolors.ENDC}: {b64encode(_get_field(payload[1], 1)).decode()}", end="")
+        if _get_field(payload[1], 0x0c):
+            print(f" {bcolors.OKBLUE}SIGNED{bcolors.ENDC}", end="")
+        print()
+        
+        for i in range(len(payload[1])):
+            if payload[1][i][0] == 1:
+                pass
+                #payload[1][i] = (None, None)
+                #payload[1][i] = (1, b64decode("D3MtN3e18QE8rve3n92wp+CwK7u/bWk/5WjQUOBN640="))
+
+        out = apns._serialize_payload(payload[0], payload[1])
+        return out
     elif id == 0xc:
         print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Keep Alive{bcolors.ENDC}")
     elif id == 0xd:
@@ -126,7 +141,7 @@ def pretty_print_payload(prefix, payload: tuple[int, list[tuple[int, bytes]]]):
         else: 
             print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Notification{bcolors.ENDC}")
     elif id == 0xb:
-        print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Notification Ack{bcolors.ENDC}")
+        print(f"{bcolors.OKGREEN}{prefix}{bcolors.ENDC}: {bcolors.OKCYAN}Notification Ack{bcolors.ENDC} {bcolors.OKBLUE}{_get_field(payload[1], 8).hex()}{bcolors.ENDC}")
     else:
         print(prefix, f"Payload ID: {hex(payload[0])}")
         for field in payload[1]:
diff --git a/proxy/proxy.py b/proxy/proxy.py
index b8fd5a9..4852188 100644
--- a/proxy/proxy.py
+++ b/proxy/proxy.py
@@ -37,18 +37,25 @@ import printer
 
 
 def proxy(conn1: tlslite.TLSConnection, conn2: tlslite.TLSConnection, prefix: str = ""):
-    while True:
-        # Read data from the first connection
-        data = conn1.read()
-        # If there is no data, the connection has closed
-        if not data:
-            break
+    try:
+        while True:
+            # Read data from the first connection
+            data = conn1.read()
+            # If there is no data, the connection has closed
+            if not data:
+                break
 
-        printer.pretty_print_payload(prefix, apns._deserialize_payload_from_buffer(data))
+            override = printer.pretty_print_payload(prefix, apns._deserialize_payload_from_buffer(data))
+            if override is not None:
+                data = override
+                print("OVERRIDE: ", end="")
+                printer.pretty_print_payload(prefix, apns._deserialize_payload_from_buffer(data))
 
-        #print(prefix, data)
-        # Write the data to the second connection
-        conn2.write(data)
+            #print(prefix, data)
+            # Write the data to the second connection
+            conn2.write(data)
+    except Exception as e:
+        pass # Probably a connection closed error
     print("Connection closed")
     # Close the connections
     conn1.close()