diff --git a/pcapkit/foundation/engines/dpkt.py b/pcapkit/foundation/engines/dpkt.py index e6f791fbb7..9ea63ebf5b 100644 --- a/pcapkit/foundation/engines/dpkt.py +++ b/pcapkit/foundation/engines/dpkt.py @@ -84,6 +84,9 @@ def run(self) -> 'None': as :mod:`dpkt` and :attr:`self._extmp ` as an iterator from :class:`dpkt.pcap.Reader`. + The global header is parsed and dumped first, by + :meth:`self.extractor.record_header `. + Warns: AttributeWarning: If :attr:`self.extractor._exlyr ` and/or :attr:`self.extractor._exptl ` @@ -121,6 +124,9 @@ def run(self) -> 'None': f'Frame {e._frnum:>3d}: {packet2chain(f)}' # pylint: disable=protected-access ) + # extract global header + ext.record_header() + if ext.magic_number in PCAP.MAGIC_NUMBER: logger.debug('dpkt: reading %s as PCAP', ext._ifnm) reader = dpkt.pcap.Reader(ext._ifile) @@ -175,8 +181,6 @@ def read_frame(self) -> 'DPKTPacket': ofile(info, name=frnum) else: ext._ofile(info, name=frnum) - ofile = ext._ofile - ext._offmt = ofile.kind # record fragments if ext._flag_r: diff --git a/pcapkit/foundation/engines/pyshark.py b/pcapkit/foundation/engines/pyshark.py index d48e638288..7628cce4b0 100644 --- a/pcapkit/foundation/engines/pyshark.py +++ b/pcapkit/foundation/engines/pyshark.py @@ -176,6 +176,9 @@ def run(self) -> 'None': as :mod:`pyshark` and :attr:`self._extmp ` as an iterator from :class:`pyshark.FileCapture`. + The global header is parsed and dumped first, by + :meth:`self.extractor.record_header `. + Warns: AttributeWarning: Warns under following circumstances: @@ -217,6 +220,9 @@ def run(self) -> 'None': f'Frame {e._frnum:>3d}: {f.frame_info.protocols}' # pylint: disable=protected-access ) + # extract global header + ext.record_header() + # extract & analyse file logger.debug('pyshark: opening %s', ext._ifnm) self._extmp = self._expkg.FileCapture(ext._ifnm, keep_packets=False) @@ -251,8 +257,6 @@ def read_frame(self) -> 'PySharkPacket': ofile(info, name=frnum) else: ext._ofile(info, name=frnum) - ofile = ext._ofile - ext._offmt = ofile.kind # trace flows if ext._flag_t: diff --git a/pcapkit/foundation/engines/scapy.py b/pcapkit/foundation/engines/scapy.py index 4a6a54ad3a..9a2c47743f 100644 --- a/pcapkit/foundation/engines/scapy.py +++ b/pcapkit/foundation/engines/scapy.py @@ -142,6 +142,9 @@ def run(self) -> 'None': :mod:`scapy.all`, since that is the import that populates the layer registries ``sniff`` needs to dissect anything. + The global header is parsed and dumped first, by + :meth:`self.extractor.record_header `. + Warns: AttributeWarning: If :attr:`self.extractor._exlyr ` and/or :attr:`self.extractor._exptl ` @@ -171,6 +174,9 @@ def run(self) -> 'None': f'Frame {e._frnum:>3d}: {packet2chain(f)}' # pylint: disable=protected-access ) + # extract global header + ext.record_header() + # extract & analyse file logger.debug('scapy: sniffing %s', ext._ifnm) self._extmp = iter(self._expkg.sniff(offline=ext._ifnm)) @@ -206,8 +212,6 @@ def read_frame(self) -> 'ScapyPacket': ofile(info, name=frnum) else: ext._ofile(info, name=frnum) - ofile = ext._ofile - ext._offmt = ofile.kind # record fragments if ext._flag_r: diff --git a/pcapkit/foundation/extraction.py b/pcapkit/foundation/extraction.py index f3a4796813..fc8ccbf30c 100644 --- a/pcapkit/foundation/extraction.py +++ b/pcapkit/foundation/extraction.py @@ -852,7 +852,16 @@ def record_header(self) -> 'EngineBase[_P]': Unless output is disabled, the engine also dumps the parsed header, under the name of ``Global Header`` for PCAP and ``Section Header 1`` for - PCAP-NG. + PCAP-NG, and sets :attr:`self._offmt ` to the output + format. + + This is the shared global-header step of the third-party engines -- + :class:`~pcapkit.foundation.engines.dpkt.DPKT`, + :class:`~pcapkit.foundation.engines.scapy.Scapy` and + :class:`~pcapkit.foundation.engines.pyshark.PyShark` -- which call it at + setup, so their output opens with the same header record as the + built-in engines' and :attr:`format` is set even for a capture with no + frames. Raises: FormatError: If the input is neither a PCAP nor a PCAP-NG file. diff --git a/tests/foundation/engines/test_runtime_engines.py b/tests/foundation/engines/test_runtime_engines.py index 6ee6d5583f..80707a2858 100644 --- a/tests/foundation/engines/test_runtime_engines.py +++ b/tests/foundation/engines/test_runtime_engines.py @@ -64,6 +64,8 @@ def make_extractor(**overrides): '_exctx': ContextRegistry(), '_vfunc': mock.Mock(), 'magic_number': b'\xa1\xb2\xc3\xd4', + # the third-party engines call this at setup for the global header + 'record_header': mock.Mock(), } values.update(overrides) return types.SimpleNamespace(**values), sink @@ -224,6 +226,7 @@ def __next__(self): engine.run() self.assertIs(engine._extmp, pcap_reader) self.assertTrue(callable(extractor._vfunc)) + extractor.record_header.assert_called_once_with() pcapng_reader = FakeReader() extractor.magic_number = PCAPNG.MAGIC_NUMBER[0] @@ -259,7 +262,8 @@ def __init__(self, data: bytes) -> None: packet = engine.read_frame() self.assertIsInstance(packet, ParsedPacket) - self.assertEqual(extractor._offmt, 'unit') + # the output format is set by ``record_header``, not per frame + self.assertIsNone(extractor._offmt) extractor._reasm.ipv4.assert_called_once_with('ipv4') extractor._reasm.ipv6.assert_called_once_with('ipv6') extractor._reasm.tcp.assert_called_once_with('tcp') @@ -341,6 +345,7 @@ def test_scapy_run_and_read_frame(self) -> None: engine.run() sniff.assert_called_once_with(offline='capture.pcap') self.assertTrue(callable(extractor._vfunc)) + extractor.record_header.assert_called_once_with() clean_extractor, _ = make_extractor(_flag_v=False) clean_engine = Scapy(clean_extractor) @@ -421,6 +426,7 @@ def close(self) -> None: with mock.patch.object(engine._expkg, 'FileCapture', return_value=capture) as file_capture: engine.run() file_capture.assert_called_once_with('capture.pcap', keep_packets=False) + extractor.record_header.assert_called_once_with() self.assertFalse(extractor._flag_r) self.assertIsNone(extractor._reasm.ipv4) self.assertTrue(callable(extractor._vfunc)) diff --git a/tests/foundation/test_extraction_record_header_unit.py b/tests/foundation/test_extraction_record_header_unit.py new file mode 100644 index 0000000000..6b27a6e38d --- /dev/null +++ b/tests/foundation/test_extraction_record_header_unit.py @@ -0,0 +1,150 @@ +# -*- coding: utf-8 -*- +"""Regression tests for #1104: the third-party engines share :meth:`Extractor.record_header`. + +`DPKT`_, `Scapy`_ and `PyShark`_ each call +:meth:`~pcapkit.foundation.extraction.Extractor.record_header` at setup, so their +output opens with the header record the built-in engines write -- ``Global +Header`` for PCAP, ``Section Header 1`` for PCAP-NG -- and +:attr:`~pcapkit.foundation.extraction.Extractor.format` is set at header time, +which is what makes it work on a capture with no frames. + +`PyShark`_ cannot run on the 3.14 interpreter these tests use, so it is driven +through a stand-in ``pyshark`` module. + +.. _DPKT: https://dpkt.readthedocs.io +.. _Scapy: https://scapy.net +.. _PyShark: https://kiminewt.github.io/pyshark + +""" +from __future__ import annotations + +import importlib +import importlib.util +import json +import os +import tempfile +import types +import unittest +import warnings +from unittest import mock + +from pcapkit.interface import extract +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 _run(fin: str, engine: str, tmp: str, **kwargs): + """Extract ``fin`` to JSON under ``tmp``, returning the extractor and its records.""" + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + extractor = extract(fin=fin, fout=os.path.join(tmp, 'out'), engine=engine, + format='json', nofile=False, **kwargs) + with open(extractor.output, encoding='utf-8') as file: + return extractor, json.load(file) + + +def _empty_pcap(tmp: str) -> str: + """A PCAP file holding only the global header of ``in.pcap``.""" + path = os.path.join(tmp, 'empty.pcap') + with open(sample_path('in.pcap'), 'rb') as src, open(path, 'wb') as dst: + dst.write(src.read(24)) + return path + + +class _HeaderStepCases: + """Shared cases, run once per real third-party engine.""" + + engine: str + + def test_pcap_output_opens_with_the_global_header(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + _, default = _run(sample_path('in.pcap'), 'default', tmp) + with tempfile.TemporaryDirectory() as tmp: + extractor, records = _run(sample_path('in.pcap'), self.engine, tmp) + self.assertEqual(type(extractor.engine).__name__, self.engine_class) + self.assertEqual(list(records)[0], 'Global Header') + self.assertEqual(records['Global Header'], default['Global Header']) + self.assertEqual(len(records), 1 + extractor.length) + + def test_pcapng_output_opens_with_the_section_header(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + _, records = _run(sample_path('dhcp.pcapng'), self.engine, tmp) + self.assertEqual(list(records)[0], 'Section Header 1') + + def test_zero_frame_capture_has_a_format(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + extractor, records = _run(_empty_pcap(tmp), self.engine, tmp) + self.assertEqual(extractor.format, 'json') + self.assertEqual(list(records), ['Global Header']) + + def test_split_output_writes_a_global_header_file(self) -> None: + with tempfile.TemporaryDirectory() as tmp, warnings.catch_warnings(): + warnings.simplefilter('ignore') + extractor = extract(fin=sample_path('in.pcap'), fout=os.path.join(tmp, 'out'), + engine=self.engine, format='json', files=True, nofile=False) + self.assertEqual(extractor.format, 'json') + self.assertIn('Global Header.json', os.listdir(os.path.join(tmp, 'out'))) + + +@unittest.skipUnless(HAS_DPKT, 'dpkt not installed') +class DPKTHeaderStepTests(_HeaderStepCases, unittest.TestCase): + engine = 'dpkt' + engine_class = 'DPKT' + + +@unittest.skipUnless(HAS_SCAPY, 'scapy not installed') +class ScapyHeaderStepTests(_HeaderStepCases, unittest.TestCase): + engine = 'scapy' + engine_class = 'Scapy' + + +class PySharkHeaderStepTests(unittest.TestCase): + """`PyShark`_ through a stand-in module, so neither it nor :program:`tshark` is needed.""" + + def _run(self, fin: str, packets: list): + from pcapkit.foundation.engines.pyshark import PyShark + + # The dotted-string patch re-reads ``pyshark`` as an attribute of ``pcapkit.toolkit``, + # which breaks on 3.10 when that attribute of ``pcapkit`` is not the object in + # ``sys.modules`` (what a module-table restore can leave). Patch the submodule + # object from ``sys.modules``, which is also where the engine imports it from. + toolkit = importlib.import_module('pcapkit.toolkit.pyshark') + + class FakeCapture: + def __init__(self, *args, **kwargs) -> None: + self.packets = iter(packets) + + def next(self): + return next(self.packets) + + def close(self) -> None: + pass + + fake = types.ModuleType('pyshark') + fake.FileCapture = FakeCapture # type: ignore[attr-defined] + with tempfile.TemporaryDirectory() as tmp, \ + mock.patch.dict('sys.modules', {'pyshark': fake}), \ + mock.patch.object(PyShark, 'unsupported_reason', return_value=None), \ + mock.patch.object(toolkit, 'packet2dict', return_value={'packet': True}): + extractor, records = _run(fin, 'pyshark', tmp) + fmt = extractor.format + self.assertIsInstance(extractor.engine, PyShark) + return fmt, records + + def test_pcap_output_opens_with_the_global_header(self) -> None: + packet = types.SimpleNamespace(number='1') + fmt, records = self._run(sample_path('in.pcap'), [packet]) + self.assertEqual(fmt, 'json') + self.assertEqual(list(records), ['Global Header', 'Frame 1']) + + def test_zero_frame_capture_has_a_format(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + fmt, records = self._run(_empty_pcap(tmp), []) + self.assertEqual(fmt, 'json') + self.assertEqual(list(records), ['Global Header']) + + +if __name__ == '__main__': + unittest.main()