diff --git a/scapy/sendrecv.py b/scapy/sendrecv.py index acb5b3cf67a..d7251449d02 100644 --- a/scapy/sendrecv.py +++ b/scapy/sendrecv.py @@ -31,7 +31,7 @@ from scapy.modules import six from scapy.modules.six.moves import map from scapy.sessions import DefaultSession -from scapy.supersocket import SuperSocket +from scapy.supersocket import SuperSocket, IterSocket if conf.route is None: # unused import, only to initialize conf.route and conf.iface* @@ -867,30 +867,42 @@ def _run(self, if offline is not None: flt = karg.get('filter') + if isinstance(offline, str): + # Single file + offline = [offline] if isinstance(offline, list) and \ all(isinstance(elt, str) for elt in offline): + # List of files sniff_sockets.update((PcapReader( fname if flt is None else - tcpdump(fname, args=["-w", "-"], flt=flt, getfd=True) + tcpdump(fname, + args=["-w", "-"], + flt=flt, + getfd=True, + quiet=quiet) ), fname) for fname in offline) elif isinstance(offline, dict): + # Dict of files sniff_sockets.update((PcapReader( fname if flt is None else - tcpdump(fname, args=["-w", "-"], flt=flt, getfd=True) + tcpdump(fname, + args=["-w", "-"], + flt=flt, + getfd=True, + quiet=quiet) ), label) for fname, label in six.iteritems(offline)) + elif isinstance(offline, (Packet, PacketList, list)): + # Iterables (list of packets, PacketList..) + offline = IterSocket(offline) + sniff_sockets[offline if flt is None else PcapReader( + tcpdump(offline, + args=["-w", "-"], + flt=flt, + getfd=True, + quiet=quiet) + )] = offline else: - # Write Scapy Packet objects to a pcap file - def _write_to_pcap(packets_list): - filename = get_temp_file(autoext=".pcap") - wrpcap(filename, offline) - return filename, filename - - if isinstance(offline, Packet): - tempfile_written, offline = _write_to_pcap([offline]) - elif isinstance(offline, (list, PacketList)) and \ - all(isinstance(elt, Packet) for elt in offline): - tempfile_written, offline = _write_to_pcap(offline) - + # Other (file descriptors...) sniff_sockets[PcapReader( offline if flt is None else tcpdump(offline, diff --git a/scapy/supersocket.py b/scapy/supersocket.py index d9632f2be6a..75eb1022fcb 100644 --- a/scapy/supersocket.py +++ b/scapy/supersocket.py @@ -24,17 +24,19 @@ import scapy.modules.six as six from scapy.packet import Packet import scapy.packet -from scapy.plist import PacketList +from scapy.plist import _PacketList, PacketList, SndRcvList from scapy.utils import PcapReader, tcpdump # Typing imports from scapy.compat import ( Any, + Iterator, List, Optional, Tuple, Type, Union, + cast, ) # Utils @@ -475,3 +477,51 @@ def select(sockets, remain=None): if (WINDOWS or DARWIN): return sockets, None return SuperSocket.select(sockets, remain=remain) + + +# More abstract objects + +class IterSocket(SuperSocket): + desc = "wrapper around an iterable" + nonblocking_socket = True + + def __init__(self, obj): + # type: (Union[Packet, List[Packet], _PacketList[Packet]]) -> None + if not obj: + self.iter = iter([]) # type: Iterator[Packet] + elif isinstance(obj, IterSocket): + self.iter = obj.iter + elif isinstance(obj, SndRcvList): + def _iter(obj=cast(SndRcvList, obj)): + # type: (SndRcvList) -> Iterator[Packet] + for s, r in obj: + if s.sent_time: + s.time = s.sent_time + yield s + yield r + self.iter = _iter() + elif isinstance(obj, (list, PacketList)): + if isinstance(obj[0], bytes): # type: ignore + # Deprecated + self.iter = (conf.raw_layer(x) for x in obj) + else: + self.iter = (y for x in obj for y in x) + else: + self.iter = obj.__iter__() + + @staticmethod + def select(sockets, remain=None): + # type: (List[SuperSocket], Any) -> Tuple[List[SuperSocket], None] + return sockets, None + + def recv(self, *args): + # type: (*Any) -> Optional[Packet] + try: + pkt = next(self.iter) + return pkt.__class__(bytes(pkt)) + except StopIteration: + raise EOFError + + def close(self): + # type: () -> None + pass diff --git a/scapy/utils.py b/scapy/utils.py index b864bc1a342..a1a3fc51e99 100644 --- a/scapy/utils.py +++ b/scapy/utils.py @@ -1635,24 +1635,14 @@ def write(self, pkt): self.write_header(pkt) self.write_packet(pkt) else: - # Import here to avoid a circular dependency - from scapy.plist import SndRcvList - if isinstance(pkt, SndRcvList): - def _iter(pkt=cast(SndRcvList, pkt)): - # type: (SndRcvList) -> Iterator[Packet] - for s, r in pkt: - if s.sent_time: - s.time = s.sent_time - yield s - yield r - pkt_iter = _iter() - else: - pkt_iter = pkt.__iter__() - for p in pkt_iter: + # Import here to avoid circular dependency + from scapy.supersocket import IterSocket + for p in IterSocket(pkt).iter: if not self.header_present: self.write_header(p) - if self.linktype != conf.l2types.get(type(p), None): + if not isinstance(p, bytes) and \ + self.linktype != conf.l2types.get(type(p), None): warning("Inconsistent linktypes detected!" " The resulting PCAP file might contain" " invalid packets." @@ -2148,9 +2138,11 @@ def tcpdump( # An error has occurred return if dump: - return b"".join( + data = b"".join( iter(lambda: proc.stdout.read(1048576), b"") # type: ignore ) + proc.terminate() + return data if getproc: return proc if getfd: