Skip to content

Commit f303033

Browse files
polybassaNils Weiss
andauthored
Enhance SuperSocket for AnsweringMachine typing (#4897)
* Enhance SuperSocket for AnsweringMachine typing and compatibility improvements * Refactor imports in supersocket.py to use TYPE_CHECKING for AnsweringMachine. --------- Co-authored-by: Nils Weiss <nils.weiss@dissecto.com>
1 parent a5bc2bb commit f303033

File tree

1 file changed

+18
-8
lines changed

1 file changed

+18
-8
lines changed

scapy/supersocket.py

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,21 @@
4444
Optional,
4545
Tuple,
4646
Type,
47+
TypeVar,
4748
cast,
49+
TYPE_CHECKING,
4850
)
51+
from scapy.compat import Self
52+
53+
if TYPE_CHECKING:
54+
from scapy.ansmachine import AnsweringMachine
55+
4956

5057
# Utils
5158

5259

5360
class _SuperSocket_metaclass(type):
54-
desc = None # type: Optional[str]
61+
desc = None # type: Optional[str]
5562

5663
def __repr__(self):
5764
# type: () -> str
@@ -82,10 +89,13 @@ class tpacket_auxdata(ctypes.Structure):
8289

8390
# SuperSocket
8491

92+
_T = TypeVar("_T", Packet, PacketList)
93+
94+
8595
class SuperSocket(metaclass=_SuperSocket_metaclass):
8696
closed = False # type: bool
8797
nonblocking_socket = False # type: bool
88-
auxdata_available = False # type: bool
98+
auxdata_available = False # type: bool
8999

90100
def __init__(self,
91101
family=socket.AF_INET, # type: int
@@ -271,19 +281,17 @@ def tshark(self, *args, **kargs):
271281
from scapy import sendrecv
272282
sendrecv.tshark(opened_socket=self, *args, **kargs)
273283

274-
# TODO: use 'scapy.ansmachine.AnsweringMachine' when typed
275284
def am(self,
276-
cls, # type: Type[Any]
277-
*args, # type: Any
285+
cls, # type: Type[AnsweringMachine[_T]]
278286
**kwargs # type: Any
279287
):
280-
# type: (...) -> Any
288+
# type: (...) -> AnsweringMachine[_T]
281289
"""
282290
Creates an AnsweringMachine associated with this socket.
283291
284292
:param cls: A subclass of AnsweringMachine to instantiate
285293
"""
286-
return cls(*args, opened_socket=self, socket=self, **kwargs)
294+
return cls(opened_socket=self, socket=self, **kwargs)
287295

288296
@staticmethod
289297
def select(sockets, remain=conf.recv_poll_rate):
@@ -295,6 +303,7 @@ def select(sockets, remain=conf.recv_poll_rate):
295303
:returns: an array of sockets that were selected and
296304
the function to be called next to get the packets (i.g. recv)
297305
"""
306+
inp = [] # type: List[SuperSocket]
298307
try:
299308
inp, _, _ = select(sockets, [], [], remain)
300309
except (IOError, select_error) as exc:
@@ -309,7 +318,7 @@ def __del__(self):
309318
self.close()
310319

311320
def __enter__(self):
312-
# type: () -> SuperSocket
321+
# type: () -> Self
313322
return self
314323

315324
def __exit__(self, exc_type, exc_value, traceback):
@@ -627,6 +636,7 @@ def _iter(obj=cast(SndRcvList, obj)):
627636
s.time = s.sent_time
628637
yield s
629638
yield r
639+
630640
self.iter = _iter()
631641
elif isinstance(obj, (list, PacketList)):
632642
if isinstance(obj[0], bytes):

0 commit comments

Comments
 (0)