diff --git a/dpkt/pcap.py b/dpkt/pcap.py index aa783f77..a558fea3 100644 --- a/dpkt/pcap.py +++ b/dpkt/pcap.py @@ -254,29 +254,36 @@ def __init__(self, fileobj, snaplen=1500, linktype=DLT_EN10MB, nano=False): self._pack_hdr = self._PktHdr._pack_hdr self.__f.write(bytes(fh)) - def writepkt(self, pkt, ts=None): + def writepkt(self, pkt, ts=None, pktlen=None): """Write single packet and optional timestamp to file. Args: pkt: `bytes` will be called on this and written to file. - ts (float): Timestamp in seconds. Defaults to current time. + ts (float, optional): Timestamp in seconds. Defaults to current time. + pktlen (int, optional): Packet length in original transmission (might + be more than the number of bytes available from + the capture). Defaults to length of pkt. """ if ts is None: ts = time.time() + self.writepkt_time(bytes(pkt), ts, pktlen) - self.writepkt_time(bytes(pkt), ts) - - def writepkt_time(self, pkt, ts): + def writepkt_time(self, pkt, ts, pktlen=None): """Write single packet and its timestamp to file. Args: pkt (bytes): Some `bytes` to write to the file ts (float): Timestamp in seconds + pktlen (int, optional): Packet length in original transmission (might + be more than the number of bytes available from + the capture). Defaults to length of pkt. """ n = len(pkt) + if pktlen is None: + pktlen = n sec = int(ts) usec = intround(ts % 1 * self._precision_multiplier) - ph = self._pack_hdr(sec, usec, n, n) + ph = self._pack_hdr(sec, usec, n, pktlen) self.__f.write(ph + pkt) def writepkts(self, pkts): @@ -286,17 +293,23 @@ def writepkts(self, pkts): Packets must be of type `bytes` as they will not be cast. Args: - pkts: iterable containing (ts, pkt) + pkts: iterable containing (ts, pkt) or (ts, pktlen, pkt) """ fd = self.__f pack_hdr = self._pack_hdr precision_multiplier = self._precision_multiplier - for ts, pkt in pkts: + for packet_data in pkts: + if len(packet_data) == 3: + ts, pktlen, pkt = packet_data + else: + ts, pkt = packet_data + pktlen = len(pkt) + n = len(pkt) sec = int(ts) usec = intround(ts % 1 * precision_multiplier) - ph = pack_hdr(sec, usec, n, n) + ph = pack_hdr(sec, usec, n, pktlen) fd.write(ph + pkt) def close(self): @@ -372,20 +385,22 @@ def dispatch(self, cnt, callback, *args): *args -- optional arguments passed to callback on execution """ processed = 0 - if cnt > 0: - for _ in range(cnt): - try: - ts, pkt = next(iter(self)) - except StopIteration: - break - callback(ts, pkt, *args) - processed += 1 - else: - for ts, pkt in self: - callback(ts, pkt, *args) - processed += 1 + loop_indefinitely = cnt == 0 + + while loop_indefinitely or processed < cnt: + try: + pktiter = next(iter(self)) + except StopIteration: + break + self._invoke_callback(pktiter, callback, *args) + processed += 1 + return processed + def _invoke_callback(self, pktiter, callback, *args): + ts, pkt = pktiter + callback(ts, pkt, *args) + def loop(self, callback, *args): self.dispatch(0, callback, *args) @@ -399,6 +414,44 @@ def __iter__(self): yield (hdr.tv_sec + (hdr.tv_usec / self._divisor), buf) +class PktlenReader(Reader): + """ + Extended pcap reader exposing the length in original transmission (might + be more than the number of bytes available from the capture). + Iterator returns (timestamp, pktlen, buf) tuple, dispatch accepts + callbacks with more arguments.. + """ + + def dispatch(self, cnt, callback, *args): + """Collect and process packets with a user callback. + + Return the number of packets processed, or 0 for a savefile. + + Arguments: + + cnt -- number of packets to process; + or 0 to process all packets until EOF + callback -- function with (timestamp, pktlen, pkt, *args) prototype + *args -- optional arguments passed to callback on execution + """ + return super().dispatch(cnt, callback, *args) + + def _invoke_callback(self, pktiter, callback, *args): + ts, pktlen, pkt = pktiter + callback(ts, pktlen, pkt, *args) + + def __iter__(self): + fd = self._Reader__f + ph = self._Reader__ph + while 1: + buf = fd.read(ph.__hdr_len__) + if not buf: + break + hdr = ph(buf) + buf = fd.read(hdr.caplen) + yield (hdr.tv_sec + (hdr.tv_usec / self._divisor), hdr.len, buf) + + class UniversalReader(object): """ Universal pcap reader for the libpcap and pcapng file formats