cli: refactor

JJTech0130 · JJTech0130 · May 18, 2024 · May 11, 2024 · May 11, 2024 · May 11, 2024
commit 67b2a4cb21c2ade2338edf1bbd4621c8576a55a0
diff --git a/pyproject.toml b/pyproject.toml
@@ -13,16 +13,17 @@ dependencies = [
 ]
 
 [project.scripts]
-apnsproxy = "pypush.cli.apnsproxy:main"
+pypush = "pypush.cli:main"
 
 [project.optional-dependencies]
 test = [
     "pytest",
     "pytest-asyncio",
 ]
-proxy = [
+cli = [
     "frida",
-    "rich"
+    "rich",
+    "typer"
 ]
 
 [tool.setuptools_scm]

diff --git a/pypush/apns/__init__.py b/pypush/apns/__init__.py
@@ -1,5 +1,5 @@
-__all__ = ["connection", "albert"]
+# __all__ = ["connection", "albert"]
 
-from . import connection
-from . import albert
-from . import protocol
+# from . import connection
+# from . import albert
+# from . import protocol
diff --git a/pypush/apns/new/_protocol.py b/pypush/apns/new/_protocol.py
@@ -47,7 +47,11 @@ def from_packet(cls, packet: Packet):
                 raise TypeError(f"Unsupported field type: {t}")
         # Check for extra fields
         for field in packet.fields:
-            if field.id not in [f.metadata["packet_id"] for f in dataclass_fields(cls) if f.metadata is not None and "packet_id" in f.metadata]:
+            if field.id not in [
+                f.metadata["packet_id"]
+                for f in dataclass_fields(cls)
+                if f.metadata is not None and "packet_id" in f.metadata
+            ]:
                 logging.warning(
                     f"Unexpected field with packet ID {field.id} in packet {packet}"
                 )
@@ -86,7 +90,13 @@ def to_packet(self) -> Packet:
     return cls
 
 
-def fid(packet_id: int, byte_len: int = 1, default: Any = MISSING, default_factory: Any = MISSING, repr: bool = True):
+def fid(
+    packet_id: int,
+    byte_len: int = 1,
+    default: Any = MISSING,
+    default_factory: Any = MISSING,
+    repr: bool = True,
+):
     """
     :param packet_id: The packet ID of the field
     :param byte_len: The length of the field in bytes (for int fields)
@@ -95,8 +105,18 @@ def fid(packet_id: int, byte_len: int = 1, default: Any = MISSING, default_facto
     if not default == MISSING and not default_factory == MISSING:
         raise ValueError("Cannot specify both default and default_factory")
     if not default == MISSING:
-        return field(metadata={"packet_id": packet_id, "packet_bytes": byte_len}, default=default, repr=repr)
+        return field(
+            metadata={"packet_id": packet_id, "packet_bytes": byte_len},
+            default=default,
+            repr=repr,
+        )
     if not default_factory == MISSING:
-        return field(metadata={"packet_id": packet_id, "packet_bytes": byte_len}, default_factory=default_factory, repr=repr)
+        return field(
+            metadata={"packet_id": packet_id, "packet_bytes": byte_len},
+            default_factory=default_factory,
+            repr=repr,
+        )
     else:
-        return field(metadata={"packet_id": packet_id, "packet_bytes": byte_len}, repr=repr)
+        return field(
+            metadata={"packet_id": packet_id, "packet_bytes": byte_len}, repr=repr
+        )
diff --git a/pypush/apns/new/protocol.py b/pypush/apns/new/protocol.py
@@ -30,7 +30,7 @@ def to_packet(self) -> Packet:
 class ConnectCommand(Command):
     PacketType = Packet.Type.Connect
 
-    push_token: bytes = fid(1)
+    push_token: Optional[bytes] = fid(1)
     state: Optional[int] = fid(2)
     flags: int = fid(5, byte_len=4)
     certificate: Optional[bytes] = fid(12)
@@ -60,18 +60,20 @@ class ConnectAck(Command):
     max_message_size: int = fid(4, byte_len=2)
     unknown5: bytes = fid(5)
     capabilities: bytes = fid(6)
-    large_message_size: int = fid(8, byte_len=2)
+    large_message_size: Optional[int] = fid(8, byte_len=2)
     timestamp: int = fid(10, byte_len=8)
     region: Optional[str] = fid(11)
     timestamp2: Optional[int] = fid(12, byte_len=8)
     unknown19: Optional[bytes] = fid(19)
 
+
 @auto_packet
 @dataclass
 class NoStorageCommand(Command):
     PacketType = Packet.Type.NoStorage
     token: bytes = fid(1)
 
+
 @auto_packet
 @dataclass(repr=False)
 class FilterCommand(Command):
@@ -86,8 +88,15 @@ class FilterCommand(Command):
     unknown12: Optional[bytes] = fid(12, default=None)
 
     def _lookup_hashes(self, hashes: Optional[list[bytes]]):
-        return [KNOWN_TOPICS_LOOKUP[hash] if hash in KNOWN_TOPICS_LOOKUP else hash for hash in hashes] if hashes else []
-
+        return (
+            [
+                KNOWN_TOPICS_LOOKUP[hash] if hash in KNOWN_TOPICS_LOOKUP else hash
+                for hash in hashes
+            ]
+            if hashes
+            else []
+        )
+
     @property
     def enabled_topics(self):
         return self._lookup_hashes(self.enabled_topic_hashes)
@@ -103,11 +112,11 @@ def opportunistic_topics(self):
     @property
     def paused_topics(self):
         return self._lookup_hashes(self.paused_topic_hashes)
-    
+
     @property
     def non_waking_topics(self):
         return self._lookup_hashes(self.non_waking_topic_hashes)
-    
+
     def __repr__(self):
         return f"FilterCommand(token={self.token}, enabled_topics={self.enabled_topics}, ignored_topics={self.ignored_topics}, opportunistic_topics={self.opportunistic_topics}, paused_topics={self.paused_topics}, non_waking_topics={self.non_waking_topics})"
 
@@ -126,12 +135,14 @@ class KeepAliveCommand(Command):
     unknown9: Optional[int] = fid(9, default=None, byte_len=1)
     unknown10: Optional[int] = fid(10, default=None, byte_len=1)
 
+
 @auto_packet
 @dataclass
 class KeepAliveAck(Command):
     PacketType = Packet.Type.KeepAliveAck
     unknown: Optional[int] = fid(1)
 
+
 @auto_packet
 @dataclass
 class Unknown29Command(Command):
@@ -143,7 +154,8 @@ class Unknown29Command(Command):
 
     def __repr__(self):
         return f"Unknown29Command(ignored)"
-
+
+
 @auto_packet
 @dataclass
 class Unknown30Command(Command):
@@ -156,6 +168,7 @@ class Unknown30Command(Command):
     def __repr__(self):
         return f"Unknown30Command(ignored)"
 
+
 @auto_packet
 @dataclass
 class Unknown32Command(Command):
@@ -168,7 +181,8 @@ class Unknown32Command(Command):
 
     def __repr__(self):
         return f"Unknown32Command(ignored)"
-
+
+
 @auto_packet
 @dataclass
 class SetStateCommand(Command):
@@ -177,6 +191,7 @@ class SetStateCommand(Command):
     state: int = fid(1)
     unknown2: int = fid(2, byte_len=4)
 
+
 @auto_packet
 @dataclass
 class SendMessageCommand(Command):
@@ -201,30 +216,51 @@ class SendMessageCommand(Command):
 
     _token_topic_1: bytes = fid(1, default=None, repr=False)
     _token_topic_2: bytes = fid(2, default=None, repr=False)
-    
+
     def __post_init__(self):
-        if not (self.topic is not None and self.token is not None and self.outgoing is not None) and not (self._token_topic_1 is not None and self._token_topic_2 is not None):
+        if not (
+            self.topic is not None
+            and self.token is not None
+            and self.outgoing is not None
+        ) and not (self._token_topic_1 is not None and self._token_topic_2 is not None):
             raise ValueError("topic, token, and outgoing must be set.")
-        
+
         if self.outgoing == True:
             assert self.topic and self.token
-            self._token_topic_1 = sha1(self.topic.encode()).digest() if isinstance(self.topic, str) else self.topic
+            self._token_topic_1 = (
+                sha1(self.topic.encode()).digest()
+                if isinstance(self.topic, str)
+                else self.topic
+            )
             self._token_topic_2 = self.token
         elif self.outgoing == False:
             assert self.topic and self.token
             self._token_topic_1 = self.token
-            self._token_topic_2 = sha1(self.topic.encode()).digest() if isinstance(self.topic, str) else self.topic
+            self._token_topic_2 = (
+                sha1(self.topic.encode()).digest()
+                if isinstance(self.topic, str)
+                else self.topic
+            )
         else:
             assert self._token_topic_1 and self._token_topic_2
-            if len(self._token_topic_1) == 20: # SHA1 hash, topic
-                self.topic = KNOWN_TOPICS_LOOKUP[self._token_topic_1] if self._token_topic_1 in KNOWN_TOPICS_LOOKUP else self._token_topic_1
+            if len(self._token_topic_1) == 20:  # SHA1 hash, topic
+                self.topic = (
+                    KNOWN_TOPICS_LOOKUP[self._token_topic_1]
+                    if self._token_topic_1 in KNOWN_TOPICS_LOOKUP
+                    else self._token_topic_1
+                )
                 self.token = self._token_topic_2
                 self.outgoing = True
             else:
-                self.topic = KNOWN_TOPICS_LOOKUP[self._token_topic_2] if self._token_topic_2 in KNOWN_TOPICS_LOOKUP else self._token_topic_2
+                self.topic = (
+                    KNOWN_TOPICS_LOOKUP[self._token_topic_2]
+                    if self._token_topic_2 in KNOWN_TOPICS_LOOKUP
+                    else self._token_topic_2
+                )
                 self.token = self._token_topic_1
                 self.outgoing = False
 
+
 @auto_packet
 @dataclass
 class SendMessageAck(Command):
@@ -272,6 +308,7 @@ def command_from_packet(packet: Packet) -> Command:
     else:
         return UnknownCommand.from_packet(packet)
 
+
 @dataclass
 class CommandStream(ObjectStream[Command]):
     transport_stream: ObjectStream[Packet]
@@ -281,10 +318,9 @@ async def send(self, command: Command) -> None:
 
     async def receive(self) -> Command:
         return command_from_packet(await self.transport_stream.receive())
-
 
     async def aclose(self) -> None:
         await self.transport_stream.aclose()
 
     async def send_eof(self) -> None:
-        await self.transport_stream.send_eof()
+        await self.transport_stream.send_eof()
diff --git a/pypush/apns/new/transport.py b/pypush/apns/new/transport.py
@@ -30,6 +30,8 @@ class Type(Enum):
         KeepAlive = 12
         KeepAliveAck = 13
         NoStorage = 14
+        Unknown17 = 17
+        Unknown18 = 18
         SetState = 20
         Unknown29 = 29
         Unknown30 = 30
@@ -59,6 +61,13 @@ async def create_courier_connection(
     )
 
 
+async def receive_exact(stream: ByteStream, length: int) -> bytes:
+    buffer = b""
+    while len(buffer) < length:
+        buffer += await stream.receive(length - len(buffer))
+    return buffer
+
+
 @dataclass
 class PacketStream(ObjectStream[Packet]):
     transport_stream: ByteStream
@@ -84,11 +93,13 @@ async def send(self, packet: Packet) -> None:
         await self.transport_stream.send(self._serialize_packet(packet))
 
     async def receive(self) -> Packet:
-        packet_id = int.from_bytes(await self.transport_stream.receive(1), "big")
-        packet_length = int.from_bytes(await self.transport_stream.receive(4), "big")
+        packet_id = int.from_bytes(await receive_exact(self.transport_stream, 1), "big")
+        packet_length = int.from_bytes(
+            await receive_exact(self.transport_stream, 4), "big"
+        )
         if packet_length == 0:
             return Packet(Packet.Type(packet_id), [])
-        payload = await self.transport_stream.receive(packet_length)
+        payload = await receive_exact(self.transport_stream, packet_length)
         assert len(payload) == packet_length
         fields = []
         while len(payload) > 0:

diff --git a/pypush/cli/__init__.py b/pypush/cli/__init__.py
@@ -0,0 +1,52 @@
+import typer
+from typing_extensions import Annotated
+import logging
+from rich.logging import RichHandler
+
+logging.basicConfig(level=logging.DEBUG, handlers=[RichHandler()], format="%(message)s")
+
+app = typer.Typer()
+
+
+@app.command()
+def proxy(
+    attach: Annotated[
+        bool, typer.Option(help="Use Frida to attach to the running `apsd`")
+    ] = True,
+    dual: Annotated[
+        bool,
+        typer.Option(
+            help="EXPERIMENTAL: Listen on both 127.0.0.2 and 127.0.0.3, to proxy both production and sandbox connections"
+        ),
+    ] = False,
+):
+    """
+    Proxy APNs traffic between the local machine and the APNs courier
+
+    Attach requires SIP to be disabled and to be running as root
+    """
+    from . import apnsproxy
+
+    apnsproxy.main(attach, dual)
+
+
+@app.command()
+def client(
+    topic: Annotated[str, typer.Argument(help="app topic to listen on")],
+    sandbox: Annotated[
+        bool, typer.Option("--sandbox/--production", help="APNs courier to use")
+    ] = True,
+):
+    """
+    Connect to the APNs courier and listen for app notifications on the given topic
+    """
+    typer.echo("Running APNs client")
+    raise NotImplementedError("Not implemented yet")
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    main()