# -*- coding: utf-8 -*-
"""DTLS Handshake-Message Reassembly
=======================================
.. module:: pcapkit.foundation.reassembly.dtls
:mod:`pcapkit.foundation.reassembly.dtls` contains
:class:`~pcapkit.foundation.reassembly.dtls.DTLS` only, which reconstructs a
DTLS handshake message split across several records back to origin
(:rfc:`6347#section-4.2.3`). A byte range is punched out of the message's
hole list as each fragment arrives, and the message is whole once no hole
remains.
.. versionadded:: 1.6.0
"""
from typing import TYPE_CHECKING, Generic
from pcapkit.foundation.reassembly.data.data import Completion, Deferred
from pcapkit.foundation.reassembly.data.dtls import (_AT, Buffer, BufferID, Datagram, DatagramID,
Fragment, HoleDescriptor, Packet)
from pcapkit.foundation.reassembly.reassembly import ReassemblyBase
from pcapkit.protocols.application.dtls.protocol import DTLS as DTLS_Protocol
if TYPE_CHECKING:
from typing import Optional, Type
from pcapkit.const.tls.handshake_type import HandshakeType as Enum_HandshakeType
__all__ = ['DTLS']
[docs]
class DTLS(ReassemblyBase[Packet[_AT], Datagram[_AT], BufferID, Buffer], Generic[_AT]): # pylint: disable=abstract-method
"""Reassembly for DTLS handshake messages.
Args:
strict: if :data:`True`, report a message that is not completely
reassembled as the tuple of its received runs; otherwise as one
contiguous payload, its holes zero-filled
store: if store reassembled datagram in memory, i.e.,
:attr:`self._dtgram <pcapkit.foundation.reassembly.reassembly.Reassembly._dtgram>`
(if not, datagram will be discarded after callback)
Note:
The buffer ID is the UDP 4-tuple plus the DTLS **epoch**: the record
sequence number resets on each cipher-state change
(:rfc:`6347#section-4.1`), so a new epoch opens a fresh buffer.
Overlapping fragments are accepted
(:rfc:`6347#section-4.2.3`, :rfc:`9147#section-5.5`). The
:class:`~pcapkit.foundation.extraction.Extractor` feeds it through
:func:`pcapkit.toolkit.pcap.dtls_reassembly`, which reads the 4-tuple
off the frame, and :func:`~pcapkit.foundation.analysis.dtls.dtls_handshake`
decodes what it reassembles. A complete message's
:attr:`~pcapkit.foundation.reassembly.data.dtls.Datagram.packet` holds
the same decode, run on first read. Both call the protocol's
:meth:`~pcapkit.protocols.application.dtls.protocol.DTLS.analyze_handshake`; the
``packet`` calls that of :attr:`__protocol_type__` (:issue:`2622`).
Example:
>>> from pcapkit.foundation.reassembly.dtls import DTLS
# Initialise instance:
>>> dtls_reassembly = DTLS()
# Call reassembly:
>>> dtls_reassembly(packet_dict)
# Fetch result:
>>> result = dtls_reassembly.datagram
"""
if TYPE_CHECKING:
protocol: 'Type[DTLS_Protocol]'
##########################################################################
# Defaults.
##########################################################################
#: Protocol name of current reassembly object.
__protocol_name__ = 'DTLS'
#: Protocol of current reassembly object. An extension point -- a protocol
#: named here need only define ``analyze_handshake``, with the signature of
#: :meth:`DTLS.analyze_handshake
#: <pcapkit.protocols.application.dtls.protocol.DTLS.analyze_handshake>`, and is read
#: for it when a message completes, so one without it raises
#: :exc:`AttributeError` there.
__protocol_type__ = DTLS_Protocol
# No ``__timeout__``: RFC 6347 section 4.2.3 and RFC 9147 section 5.5 give
# DTLS handshake reassembly no timeout, so it inherits the base default,
# :data:`math.inf`, and a buffer is never expired by time. It is held until
# its messages complete or the capture ends.
##########################################################################
# Methods.
##########################################################################
def __init__(self, *, strict: 'bool' = True, store: 'bool' = True) -> 'None':
"""Initialise packet reassembly.
Args:
strict: if :data:`True`, report a message that is not completely
reassembled as the tuple of its received runs; otherwise as
one contiguous payload, its holes zero-filled
store: if store reassembled datagram in memory
"""
super().__init__(strict=strict, store=store)
#: Optional[Datagram[_AT]]: The message, if any, that the most recent
#: call to :meth:`reassembly` completed; see :attr:`last`.
self._last: 'Optional[Datagram[_AT]]' = None
##########################################################################
# Properties.
##########################################################################
@property
def last(self) -> 'Optional[Datagram[_AT]]':
"""The message reassembled by the most recent call to
:meth:`reassembly`, or :data:`None` if that fragment did not complete
one.
This is a synchronous alternative to :attr:`datagram`/:meth:`fetch`
for a caller that wants to know, immediately after feeding one
fragment, whether *that* fragment was the one that completed its
message, rather than waiting for the buffer (which may hold several
other messages still in flight under the same buffer ID) to be
flushed.
"""
return self._last
##########################################################################
# Methods.
##########################################################################
[docs]
def reassembly(self, info: 'Packet[_AT]') -> 'None':
"""Reassembly procedure.
Arguments:
info: :term:`info` record of packets to be reassembled
"""
# clear cache
self._flag_n = False
self.__cached__.clear()
self._last = None
BUFID = info.bufid # Buffer Identifier
MSN = info.message_seq # Message Sequence Number
FO = info.fragment_offset # Fragment Offset
FL = info.fragment_length # Fragment Length
# No time-based expiry: a buffer is held until its messages complete,
# so a fragment is never abandoned for arriving late.
buf = self._buffer.get(BUFID)
if buf is None:
buf = Buffer(msg={}, timestamp=info.timestamp)
self._buffer[BUFID] = buf
frag = buf.msg.get(MSN)
# A fresh buffer for this MSN, or a disagreement with the one held
# there: a different ``length`` or ``msg_type`` means this fragment
# cannot belong to the message under construction. A reused 4-tuple
# restarts ``message_seq`` at 0 (:rfc:`6347#section-4.2.2`). Restarting
# prevents merging disagreeing fragments. Known limitation: a late
# fragment of the displaced message whose length and msg_type match
# the new one can still splice into it.
if frag is None or frag.length != info.length or frag.msg_type != info.msg_type:
if frag is not None:
# The displaced partial is reported as PARTIAL rather than dropped.
self._dtgram.extend(self.submit(Buffer(msg={MSN: frag}, timestamp=buf.timestamp),
bufid=BUFID))
length = info.length
frag = Fragment(
ind=[],
msg_type=info.msg_type,
length=length,
raw=bytearray(length),
hdl=([HoleDescriptor(first=0, last=length - 1)] if length else []),
)
buf.msg[MSN] = frag
frag.ind.append(info.num)
# The fragment's data length is what its header declares, but only
# the octets the capture actually holds can be written, and a write
# past the message's own declared length is clipped.
held = max(min(len(info.payload), FL, len(frag.raw) - FO), 0)
if held:
stop = FO + held
frag.raw[FO:stop] = info.payload[:held]
self._update_hole_descriptors(frag.hdl, FO, stop - 1)
# the message is whole once no hole remains
if not frag.hdl:
del buf.msg[MSN]
if not buf.msg:
buf = self._buffer.pop(BUFID)
datagram = self._datagram_from_fragment(frag, BUFID, MSN, Completion.COMPLETE)
self._dtgram.append(datagram)
self._last = datagram
for callback in self.__callback_fn__:
callback([datagram])
[docs]
def submit(self, buf: 'Buffer', *, bufid: 'BufferID', checked: 'bool' = False, timeout: 'bool' = False) -> 'list[Datagram[_AT]]': # type: ignore[override] # pylint: disable=arguments-differ,unused-argument,line-too-long
"""Submit reassembled payload.
Arguments:
buf: buffer dict of reassembled packets
bufid: buffer identifier
checked: unused; every message's completeness is decided from its
own hole list, so there is nothing for a caller to assert
ahead of time
timeout: whether this buffer is being submitted because
:meth:`~pcapkit.foundation.reassembly.reassembly.Reassembly.expire`
abandoned it under the reassembly timeout
Returns:
Every message still held in ``buf``, each reported
:attr:`~pcapkit.foundation.reassembly.data.data.Completion.COMPLETE`
if its hole list is empty, otherwise
:attr:`~pcapkit.foundation.reassembly.data.data.Completion.TIMEOUT`
or :attr:`~pcapkit.foundation.reassembly.data.data.Completion.PARTIAL`.
"""
ret = [] # type: list[Datagram[_AT]]
for (msn, frag) in buf.msg.items():
if not frag.hdl:
completion = Completion.COMPLETE
else:
completion = Completion.TIMEOUT if timeout else Completion.PARTIAL
ret.append(self._datagram_from_fragment(frag, bufid, msn, completion))
for callback in self.__callback_fn__:
callback(ret)
return ret
[docs]
def _datagram_from_fragment(self, frag: 'Fragment', bufid: 'BufferID', message_seq: 'int',
completion: 'Completion') -> 'Datagram[_AT]':
"""Build a :class:`~pcapkit.foundation.reassembly.data.dtls.Datagram`
from a message's buffered :class:`~pcapkit.foundation.reassembly.data.dtls.Fragment`.
Arguments:
frag: message-under-construction to report
bufid: buffer identifier it was held under
message_seq: its message sequence number
completion: how completely it was reassembled
Returns:
The reassembled (or partially reassembled) message. Its ``packet``
is the decode :func:`~pcapkit.foundation.analysis.dtls.dtls_handshake`
runs, by this reassembly's own
:attr:`protocol <pcapkit.foundation.reassembly.reassembly.Reassembly.protocol>`,
deferred to its first read, if the message is complete, and
:data:`None` otherwise, as that function returns.
"""
packet = None # type: Optional[Deferred[tuple[Enum_HandshakeType, int, int]]]
if completion:
body = bytes(frag.raw)
payload = body # type: 'bytes | tuple[bytes, ...]'
packet = Deferred(self.protocol.analyze_handshake,
(frag.msg_type, frag.length, message_seq), body)
elif self._flag_s:
payload = self._received_runs(frag)
else:
payload = bytes(frag.raw)
return Datagram(
completed=completion,
id=DatagramID(
src=bufid[0],
sport=bufid[1],
dst=bufid[2],
dport=bufid[3],
epoch=bufid[4],
),
index=tuple(frag.ind),
message_seq=message_seq,
msg_type=frag.msg_type,
length=frag.length,
body=payload,
packet=packet,
)
[docs]
@staticmethod
def _received_runs(frag: 'Fragment') -> 'tuple[bytes, ...]':
"""Extract the octet runs actually received, for ``strict=True``'s
report of an incomplete message.
Arguments:
frag: message-under-construction to extract runs from
Returns:
Received octet runs, in offset order; empty runs (a hole butting
straight up against the start, the end, or another hole) are left
out.
"""
runs = [] # type: list[bytes]
pos = 0
for hole in sorted(frag.hdl, key=lambda hole: hole.first):
if hole.first > pos:
runs.append(bytes(frag.raw[pos:hole.first]))
pos = max(pos, hole.last + 1)
if pos < frag.length:
runs.append(bytes(frag.raw[pos:frag.length]))
return tuple(run for run in runs if run)
[docs]
@staticmethod
def _update_hole_descriptors(hdl: 'list[HoleDescriptor]', first: 'int', last: 'int') -> 'None':
"""Punch a received byte range out of a message's hole list.
Offsets are message-relative and inclusive. They never wrap and need
no open upper bound, since every fragment of a message declares the
same total ``length`` (:rfc:`6347#section-4.2.3`).
Arguments:
hdl: the message's hole descriptor list, mutated in place
first: offset of the first octet the arriving fragment actually
held
last: offset of the last octet the arriving fragment actually
held, inclusive
"""
new_hdl = [] # type: list[HoleDescriptor]
for hole in hdl: # step one
if first > hole.last or last < hole.first: # steps two and three
new_hdl.append(hole)
continue
# step four: this hole is (partly) filled and replaced by what of
# it survives, on either side of the arriving range
if first > hole.first: # step five
new_hdl.append(HoleDescriptor(first=hole.first, last=first - 1))
if last < hole.last: # step six
new_hdl.append(HoleDescriptor(first=last + 1, last=hole.last))
# step seven: go on to the next hole, since one fragment may span
# several of them
hdl[:] = new_hdl