big typing update

scripts/protobufgen.py

@@ -118,6 +118,11 @@
             if file.name == "services_pb2_grpc.pyi":
                 content = content.replace("grpc.Channel", "grpc.aio.Channel")
 
+            if file.name == "sliver_pb2.pyi":
+                content = content.replace(
+                    "from common_pb2 import", "from ..commonpb.common_pb2 import"
+                )
+
             file.write_text(content)
             console.log(f"Rewrote imports for {file}")
 except Exception as e:

src/sliver/_protocols.py

@@ -0,0 +1,26 @@
+from typing import Protocol
+from .pb.commonpb.common_pb2 import Request
+from .pb.rpcpb.services_pb2_grpc import SliverRPCStub
+
+
+class PbWithRequestProp(Protocol):
+    """Protocol for protobuf with Request field"""
+
+    @property
+    def Request(self) -> Request:
+        ...
+
+
+class InteractiveObject(Protocol):
+    """Protocol for objects with interactive methods"""
+
+    @property
+    def timeout(self) -> int:
+        ...
+
+    @property
+    def _stub(self) -> SliverRPCStub:
+        ...
+
+    def _request(self, pb: PbWithRequestProp) -> PbWithRequestProp:
+        ...

src/sliver/beacon.py

@@ -24,16 +24,15 @@
 from .interactive import BaseInteractiveCommands
 from .pb.rpcpb.services_pb2_grpc import SliverRPCStub
 from .protobuf import client_pb2, common_pb2, sliver_pb2
+from ._protocols import PbWithRequestProp
 
-TIMEOUT = 60
 
-
-class BaseBeacon(object):
+class BaseBeacon:
     def __init__(
         self,
         beacon: client_pb2.Beacon,
         channel: grpc.aio.Channel,
-        timeout: int = TIMEOUT,
+        timeout: int = 60,
     ):
         """Base class for Beacon classes.
 
@@ -137,7 +136,7 @@ def reconnect_interval(self) -> int:
         """Reconnect interval"""
         return self._beacon.ReconnectInterval
 
-    def _request(self, pb):
+    def _request(self, pb: PbWithRequestProp):
         """
         Set request attributes based on current beacon, I'd prefer to return a generic Request
         object, but protobuf for whatever reason doesn't let you assign this type of field directly.
@@ -146,7 +145,7 @@ def _request(self, pb):
 
         :param pb: A protobuf request object.
         """
-        pb.Request.SessionID = self._beacon.ID
+        pb.Request.BeaconID = self._beacon.ID
         pb.Request.Timeout = self.timeout - 1
         pb.Request.Async = True
         return pb
@@ -179,7 +178,7 @@ async def taskresult_events(self):
                 self._log.exception(err)
 
 
-def beacon_taskresult(pb_object):
+def beacon_taskresult(pb_object: Any):
     """
     Wraps a class method to return a future that resolves when the
     beacon task result is available.
@@ -281,10 +280,6 @@ async def get_system(self, *args, **kwargs) -> sliver_pb2.GetSystem:
     async def execute_shellcode(self, *args, **kwargs) -> sliver_pb2.Task:
         return await super().execute_shellcode(*args, **kwargs)
 
-    @beacon_taskresult(sliver_pb2.Task)
-    async def task(self, *args, **kwargs) -> sliver_pb2.Task:
-        return await super().task(*args, **kwargs)
-
     @beacon_taskresult(None)
     async def msf(self, *args, **kwargs) -> None:
         return await super().msf(*args, **kwargs)
@@ -305,10 +300,6 @@ async def migrate(self, *args, **kwargs) -> sliver_pb2.Migrate:
     async def execute(self, *args, **kwargs) -> sliver_pb2.Execute:
         return await super().execute(*args, **kwargs)
 
-    @beacon_taskresult(sliver_pb2.Execute)
-    async def execute_token(self, *args, **kwargs) -> sliver_pb2.Execute:
-        return await super().execute_token(*args, **kwargs)
-
     @beacon_taskresult(sliver_pb2.Sideload)
     async def sideload(self, *args, **kwargs) -> sliver_pb2.Sideload:
         return await super().sideload(*args, **kwargs)