fix(extraction): extraction with format='pcap' fails (#1127)

pcapkit/foundation/engines/pcap.py

@@ -103,8 +103,11 @@ def run(self) -> 'None':
         link layer protocol type and nanosecond flag are also saved on the
         current :class:`PCAP` engine instance.
 
-        For output, the method will dump the parsed PCAP global header under
-        the name of ``Global Header``.
+        For output, the method will create the output writer from the parsed
+        PCAP global header, via :meth:`Extractor._open_output
+        <pcapkit.foundation.extraction.Extractor._open_output>`, and dump the
+        header under the name of ``Global Header`` -- except to a PCAP writer,
+        which has already written it.
 
         """
         # pylint: disable=attribute-defined-outside-init,protected-access
@@ -121,12 +124,14 @@ def run(self) -> 'None':
         if ext._flag_q:
             return
 
+        ext._open_output(protocol=self._dlink, byteorder=self._gbhdr.byteorder, nanosecond=self._nnsec)
         if ext._flag_f:
             ofile = ext._ofile(f'{ext._ofnm}/Global Header.{ext._fext}')
-            ofile(self._gbhdr.info.to_dict(), name='Global Header')
         else:
-            ext._ofile(self._gbhdr.info.to_dict(), name='Global Header')
             ofile = ext._ofile
+        # NOTE: a PCAP writer has already written the global header on creation.
+        if ofile.kind != 'pcap':
+            ofile(self._gbhdr.info.to_dict(), name='Global Header')
         ext._offmt = ofile.kind
 
     def read_frame(self) -> 'Frame':
@@ -161,11 +166,14 @@ def read_frame(self) -> 'Frame':
         # write plist
         frnum = f'Frame {ext._frnum}'
         if not ext._flag_q:
+            # NOTE: a PCAP writer takes the frame itself, for its octets and
+            # record header; the other writers take its fields.
+            info = frame.info if ext._offmt == 'pcap' else frame.info.to_dict()
             if ext._flag_f:
                 ofile = ext._ofile(f'{ext._ofnm}/{frnum}.{ext._fext}')
-                ofile(frame.info.to_dict(), name=frnum)
+                ofile(info, name=frnum)
             else:
-                ext._ofile(frame.info.to_dict(), name=frnum)
+                ext._ofile(info, name=frnum)
 
         # record fragments
         if ext._flag_r:

pcapkit/foundation/engines/pcapng.py

@@ -12,6 +12,7 @@
 
 from pcapkit.const.pcapng.block_type import BlockType as Enum_BlockType
 from pcapkit.corekit.infoclass import Info, info_final
+from pcapkit.dumpkit.pcap import PCAPIO
 from pcapkit.foundation.engines.engine import EngineBase
 from pcapkit.protocols.misc.pcapng import PCAPNG as P_PCAPNG
 from pcapkit.utilities.exceptions import FormatError, stacklevel
@@ -155,6 +156,12 @@ def run(self) -> 'None':
         shb._ctx = self._ctx
 
         logger.debug('PCAP-NG section %d header read', len(self._ctx_list))
+
+        # NOTE: a PCAP file has one link type and one timestamp resolution,
+        # where a PCAP-NG file declares them per interface, in blocks that follow.
+        if not ext._flag_q and isinstance(ext._ofile, type) and issubclass(ext._ofile, PCAPIO):
+            raise FormatError('PCAP-NG: PCAP output is not supported for a PCAP-NG input')
+        ext._open_output()
         self._write_file(shb.info, name=f'Section Header {len(self._ctx_list)}')
 
     def read_frame(self) -> 'P_PCAPNG':

pcapkit/foundation/extraction.py

@@ -14,6 +14,7 @@
 
 """
 import collections
+import functools
 import importlib
 import io
 import os
@@ -125,6 +126,9 @@ class Extractor(Generic[_P]):
         _flag_f: 'bool'
         #: No output file, i.e., no output file is to be generated.
         _flag_q: 'bool'
+        #: Output opened flag. It indicates if :meth:`~pcapkit.foundation.extraction.Extractor._open_output`
+        #: has replaced the writer class in :attr:`_ofile` with the writer.
+        _flag_o: 'bool'
         #: Reassembly flag. It indicates if datagram reassembly is enabled.
         #: Every engine reads it to decide whether to feed
         #: :attr:`~pcapkit.foundation.extraction.Extractor._reasm`, and the
@@ -693,6 +697,7 @@ def run(self) -> 'None':  # pylint: disable=inconsistent-return-statements
                 logger.debug('using engine %s (%s)', eng.name, eng.module)
                 self._exeng = eng(self)
                 self._exeng.run()
+                self._open_output()
 
                 # start iteration
                 self.record_frames()
@@ -718,6 +723,7 @@ def run(self) -> 'None':  # pylint: disable=inconsistent-return-statements
 
         # start engine
         self._exeng.run()
+        self._open_output()
 
         # start iteration
         self.record_frames()
@@ -884,6 +890,34 @@ def record_header(self) -> 'EngineBase[_P]':
 
         raise FormatError(f'unknown file format: {self._magic!r}')
 
+    def _open_output(self, **kwargs: 'Any') -> 'None':
+        """Create the output writer.
+
+        The engines call this at their global-header step, passing the header's
+        link type, byte order and timestamp resolution as ``protocol``,
+        ``byteorder`` and ``nanosecond``: :class:`~pcapkit.dumpkit.pcap.PCAPIO`
+        writes its own global header from them, and the other writers ignore
+        them. :meth:`run` calls it again, without arguments, after the engine
+        has started, for the engines that have no such step. Only the first
+        call has an effect.
+
+        Under split mode, i.e. :attr:`self._flag_f <Extractor._flag_f>` is
+        :data:`True`, :attr:`self._ofile <Extractor._ofile>` stays a factory,
+        with ``kwargs`` bound, that the engines call once per output file.
+
+        Args:
+            **kwargs: Global header fields for the writer.
+
+        """
+        if self._flag_q or self._flag_o:
+            return
+        self._flag_o = True
+
+        if self._flag_f:
+            self._ofile = functools.partial(self._ofile, **kwargs)  # type: ignore[misc]
+        else:
+            self._ofile = self._ofile(self._ofnm, **kwargs)
+
     def record_frames(self) -> 'None':
         """Read packet frames.
 
@@ -1018,6 +1052,8 @@ def __init__(self,
 
                 * If using PCAP output for TCP flow tracing while the extraction engine is
                   DPKT, Scapy, PyShark or PyPCAPFile.
+                * If using PCAP output while the extraction engine is a third-party
+                  one, i.e. DPKT, Scapy, PyShark, PyPCAP, pcap-ct or PyPCAPFile.
                 * If output file format is not supported.
 
             pcapkit.utilities.warnings.AttributeWarning: If ``trace_analyse`` is
@@ -1031,6 +1067,17 @@ def __init__(self,
         if format is None:
             format = 'tree'
 
+        # NOTE: the third-party engines' adapters hand the output writer each
+        # frame as a plain :obj:`dict` built by ``packet2dict``, which carries
+        # neither the frame's octets nor its record header, so there is nothing
+        # for the PCAP writer to write. The flow tracer has the same limitation,
+        # and is given the same replacement below.
+        if ((engine or 'default').lower() in ('dpkt', 'scapy', 'pyshark', 'pypcap', 'pcap_ct', 'pypcapfile')
+                and format in ('pcap', 'cap')):
+            warn(f"'Extractor(engine={engine})' does not support 'format={format}'; "
+                 "using 'format=\"json\"' instead", FormatWarning, stacklevel=stacklevel())
+            format = 'json'
+
         ifnm, ofnm, fmt, oext, files = self.make_name(fin, fout, format, extension, files=files, nofile=nofile)
 
         self._ifnm = ifnm  # input file name
@@ -1042,6 +1089,7 @@ def __init__(self,
         self._flag_e = False                 # EOF flag
         self._flag_f = files                 # split file flag
         self._flag_q = nofile                # no output flag
+        self._flag_o = False                 # output opened flag
         self._flag_r = reassembly            # reassembly flag
         self._flag_t = trace                 # trace flag
         self._flag_v = False                 # verbose flag
@@ -1199,7 +1247,10 @@ def __init__(self,
             # NOTE: make_dumper() names every subclass it builds 'DictDumper', so the
             # useful name is the output class it wraps.
             logger.debug('dumping %s output to %s via %s', fmt, ofnm, output.__name__)
-            self._ofile = dumper if self._flag_f else dumper(ofnm)  # output file
+            # NOTE: only the writer *class* is kept here. The writer itself is
+            # created by :meth:`_open_output` once the engine has read the global
+            # header, which the PCAP writer needs in order to write its own.
+            self._ofile = dumper  # output file
         else:
             logger.debug('file output disabled')
 

tests/foundation/engines/test_runtime_engines.py

@@ -66,6 +66,8 @@ def make_extractor(**overrides):
         'magic_number': b'\xa1\xb2\xc3\xd4',
         # the third-party engines call this at setup for the global header
         'record_header': mock.Mock(),
+        # the header step creates the writer through this; ``_ofile`` is the writer already
+        '_open_output': mock.Mock(),
     }
     values.update(overrides)
     return types.SimpleNamespace(**values), sink
@@ -94,6 +96,7 @@ def __init__(self, file) -> None:
                 self.file = file
                 self.version = VersionInfo(2, 4)
                 self.protocol = LinkType.ETHERNET
+                self.byteorder = 'little'
                 self.nanosecond = True
                 self.info = FakeInfo(header='global')
 
@@ -109,6 +112,8 @@ def __init__(self, file) -> None:
                 self.assertTrue(engine.nanosecond)
                 self.assertEqual(extractor._offmt, 'unit')
                 self.assertEqual(sink.records[-1][1], 'Global Header')
+                extractor._open_output.assert_called_once_with(
+                    protocol=LinkType.ETHERNET, byteorder='little', nanosecond=True)
 
         quiet, sink = make_extractor(_flag_q=True)
         engine = PCAP(quiet)

tests/foundation/test_extraction.py

@@ -106,6 +106,7 @@ def _bare_extractor(self):
         extractor._ofnm = 'output.txt'
         extractor._offmt = 'tree'
         extractor._flag_q = False
+        extractor._flag_o = True  # no writer class to open
         extractor._flag_d = True
         extractor._flag_r = True
         extractor._flag_t = True

tests/foundation/test_extraction_pcap_output_unit.py

@@ -0,0 +1,108 @@
+# -*- coding: utf-8 -*-
+"""Regression tests for #1127: extraction with ``format='pcap'``.
+
+:class:`~pcapkit.foundation.extraction.Extractor` creates the output writer at
+the engine's global-header step, via
+:meth:`~pcapkit.foundation.extraction.Extractor._open_output`, because
+:class:`~pcapkit.dumpkit.pcap.PCAPIO` needs that header's link type, byte order
+and timestamp resolution to write its own.
+
+"""
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import tempfile
+import unittest
+import warnings
+
+from pcapkit.interface import extract
+from pcapkit.utilities.exceptions import FormatError
+from pcapkit.utilities.warnings import FormatWarning
+from tests._support import sample_path
+
+HAS_DPKT = importlib.util.find_spec('dpkt') is not None
+HAS_SCAPY = importlib.util.find_spec('scapy') is not None
+
+
+def _extract(fin: str, fout: str, **kwargs):
+    with warnings.catch_warnings():
+        warnings.simplefilter('ignore')
+        return extract(fin=fin, fout=fout, nofile=False, **kwargs)
+
+
+def _read(path: str) -> bytes:
+    with open(path, 'rb') as file:
+        return file.read()
+
+
+class PCAPOutputTests(unittest.TestCase):
+
+    def test_pcap_output_reproduces_the_input(self) -> None:
+        for fmt in ('pcap', 'cap'):
+            with self.subTest(format=fmt), tempfile.TemporaryDirectory() as tmp:
+                extractor = _extract(sample_path('in.pcap'), os.path.join(tmp, 'out'), format=fmt)
+                self.assertEqual(extractor.format, 'pcap')
+                self.assertEqual(_read(extractor.output), _read(sample_path('in.pcap')))
+
+    def test_split_pcap_output_writes_one_capture_per_frame(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            extractor = _extract(sample_path('in.pcap'), os.path.join(tmp, 'out'),
+                                 format='pcap', files=True)
+            self.assertEqual(extractor.format, 'pcap')
+            header = _read(sample_path('in.pcap'))[:24]
+            self.assertEqual(_read(os.path.join(extractor.output, 'Global Header.pcap')), header)
+            for num in range(1, extractor.length + 1):
+                path = os.path.join(extractor.output, f'Frame {num}.pcap')
+                self.assertEqual(_read(path)[:24], header)
+                self.assertEqual(_extract(path, os.path.join(tmp, 'again'), format='json').length, 1)
+
+    def test_zero_frame_capture(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            empty = os.path.join(tmp, 'empty.pcap')
+            with open(empty, 'wb') as file:
+                file.write(_read(sample_path('in.pcap'))[:24])
+            extractor = _extract(empty, os.path.join(tmp, 'out'), format='pcap')
+            self.assertEqual(extractor.format, 'pcap')
+            self.assertEqual(_read(extractor.output), _read(empty))
+
+    def test_pcapng_input_cannot_be_written_as_pcap(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            with self.assertRaisesRegex(FormatError, 'PCAP output is not supported'):
+                _extract(sample_path('dhcp.pcapng'), os.path.join(tmp, 'out'), format='pcap')
+            self.assertEqual(os.listdir(tmp), [])
+
+    def test_other_formats_are_unaffected(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            extractor = _extract(sample_path('in.pcap'), os.path.join(tmp, 'out'), format='json')
+            with open(extractor.output, encoding='utf-8') as file:
+                records = json.load(file)
+        self.assertEqual(extractor.format, 'json')
+        self.assertEqual(list(records)[0], 'Global Header')
+        self.assertEqual(len(records), 1 + extractor.length)
+
+
+class ThirdPartyEngineTests(unittest.TestCase):
+    """These engines' frames carry no octets, so PCAP output falls back to JSON."""
+
+    def _check(self, engine: str) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            with self.assertWarnsRegex(FormatWarning, "does not support 'format=pcap'"):
+                extractor = extract(fin=sample_path('in.pcap'), fout=os.path.join(tmp, 'out'),
+                                    format='pcap', engine=engine, nofile=False)
+            self.assertEqual(extractor.format, 'json')
+            with open(extractor.output, encoding='utf-8') as file:
+                self.assertEqual(list(json.load(file))[0], 'Global Header')
+
+    @unittest.skipUnless(HAS_DPKT, 'dpkt not installed')
+    def test_dpkt(self) -> None:
+        self._check('dpkt')
+
+    @unittest.skipUnless(HAS_SCAPY, 'scapy not installed')
+    def test_scapy(self) -> None:
+        self._check('scapy')
+
+
+if __name__ == '__main__':
+    unittest.main()