Source code for hermespy.modem.frame_generator.scapy

from .frame_generator import FrameGenerator
from ..bits_source import BitsSource

import numpy as np
from typing import Type

from scapy.packet import Packet, raw  # type: ignore

__author__ = "Egor Achkasov"
__copyright__ = "Copyright 2024, Barkhausen Institut gGmbH"
__credits__ = ["Egor Achkasov", "Jan Adler"]
__license__ = "AGPLv3"
__version__ = "1.5.0"
__maintainer__ = "Jan Adler"
__email__ = "jan.adler@barkhauseninstitut.org"
__status__ = "Prototype"




class FrameGeneratorScapy(FrameGenerator):
    """Scapy wrapper frame generator.

    Attrs:
       packet: Scapy packet header to which a payload would be attached.
       packet_type: Type of the first layer of the packet header.
    """

    packet: Packet
    packet_type: Type[Packet]

    def __init__(self, packet: Packet) -> None:
        """
        Args:
            packet: Packet to which a payload will be attached.
        """
        self.packet = packet
        self.packet_num_bits = len(packet) * 8
        self.packet_type = packet.layers()[0]



    def pack_frame(self, source: BitsSource, num_bits: int) -> np.ndarray:
        """Generate a frame of num_bits bits from the given bitsource.
        Note that the payload size is num_bits minus number of bits in the packet header.
        Note that payload can be of size 0, in which case no data would be sent (except for the packet header).

        Args:
            source: Payload source.
            num_bits: Number of bits in the whole resulting frame.

        Raises:
            ValueError: If num_bits is not enough to fit the packet.
        """

        payload_num_bits = num_bits - self.packet_num_bits
        if payload_num_bits < 0:
            raise ValueError(
                f"Packet header is bigger then the requested amount of bits ({len(self.packet)*8} > {num_bits})."
            )
        packet_new = self.packet_type()
        packet_new.add_payload(np.packbits(source.generate_bits(payload_num_bits)).tobytes())
        return np.unpackbits(np.frombuffer(raw(packet_new), np.uint8))




    def unpack_frame(self, frame: np.ndarray) -> np.ndarray:
        if frame.size < self.packet_num_bits:
            raise ValueError(
                f"The frame contains less bits then the header ({frame.size} < {self.packet_num_bits})."
            )
        return frame[self.packet_num_bits :]