classdef HelperPCAPNGWriter < handle & HelperPCAPUtils
%HelperPCAPNGWriter Create a PCAPNG file writer object which supports
%writing protocol packets in PCAPNG format
%
%   PCAPNGWRITER = HelperPCAPNGWriter creates a packet capture next
%   generation (PCAPNG) file writer handle object, PCAPNGWRITER, that
%   writes protocol packets into a PCAPNG format file.
%
%   LINKTYPE        - Unique identifier for a protocol
%                     For example: Linktype of WLAN protocol is 127
%                                  Linktype of Bluetooth low energy(BLE)
%                                  protocol is 251
%
%   PCAPNGWRITER = HelperPCAPNGWriter(Name1, Value1, ..., NameN, ValueN)
%   specifies additional name-value pair arguments described below. When a
%   name-value pair is not specified, it's default value is used.
%
%   FileName        - Name of the file. It is specified as a character
%                     vector or a string. The default file name is
%                     'capture'.
%
%   FileComment     - Comment for the file. It is specified as a character
%                     vector or a string. The default value is an empty
%                     character.
%
%   ByteOrder       - Specifies the sequential order in which bytes
%                     are arranged into larger numerical values. It can be
%                     one of the 'little-endian' |'big-endian'. The default
%                     value is 'little-endian'.
%
%   HelperPCAPNGWriter methods:
%
%   writeInterfaceDescriptionBlock - Write interface description block
%                                    into the PCAPNG file and return the
%                                    interface ID. This interface ID must
%                                    be used for logging packets of
%                                    specified interface.
%   write                          - Write enhanced packet block into a
%                                    PCAPNG format file
%   writeCustomBlock               - Write custom block into a PCAPNG
%                                    format file
%
%   % Example 1:
%   %  Write WLAN packets into a PCAPNG format file with
%   %  timestamp and a packet comment
%
%   % Create a PCAP object with file name as sample.pcapng
%   wlanLinkType = 105;
%   pcapngObj = HelperPCAPNGWriter('FileName', 'sample');
%
%   % The methods MUST be called in this sequence:
%   %   1.writeInterfaceDescriptionBlock
%   %   2.write
%
%   % Write interface block
%   interfaceName = 'WLAN';
%   interfaceId = writeInterfaceDescriptionBlock(pcapngObj, wlanLinkType, ...
%                                                interfaceName);
%
%   % Create a WLAN MAC packet of type QoS Data
%   macConfig = wlanMACFrameConfig('FrameType', 'QoS Data');
%   payload = '00576000103afffe80';
%   mpdu = wlanMACFrame(payload, macConfig);
%
%   % Write enhanced packet block
%   timestamp = 0;
%   packetComment = 'This is the first packet';
%   write(pcapngObj, timestamp, interfaceId, mpdu, packetComment);
%
%
%   % Example 2:
%   %  Export the WLAN packet without radiotap, and the BLE packet into
%   %  a PCAPNG format file
%
%   % Create a PCAP object with file name as sample.pcapng
%   wlanLinkType = 105;
%   pcapngObj = HelperPCAPNGWriter('FileName', 'sample');
%
%   % The methods must be called in this sequence:
%   %   1.writeInterfaceDescriptionBlock
%   %   2.write
%
%   % Write interface block
%   interfaceName = 'WLAN';
%   interfaceId = writeInterfaceDescriptionBlock(pcapngObj, wlanLinkType, ...
%                                                interfaceName);
%
%   % Create a WLAN MAC packet of type QoS Data
%   macConfig = wlanMACFrameConfig('FrameType', 'QoS Data');
%   payload = '00576000103afffe80';
%   mpdu = wlanMACFrame(payload, macConfig);
%
%   % Write enhanced packet block
%   timestamp = 124800;
%   packetComment = 'This is WLAN packet';
%   write(pcapngObj, timestamp, interfaceId, mpdu, packetComment);
%
%   % Write interface description block for BLE protocol
%   bleLinkType = 251;
%   interfaceName = 'BLE interface';
%   interfaceId = writeInterfaceDescriptionBlock(pcapngObj, bleLinkType, ...
%                                                interfaceName);
%   % Create a BLE packet
%   llpacket = '42BC13E206120E00050014010A001F0040001700170000007D47C0';
%
%   % Write enhanced packet block
%   timestamp = 129100;
%   packetComment = 'This is BLE packet';
%   write(pcapngObj, timestamp, interfaceId, llpacket, packetComment);
%
%   Reference - https://www.tcpdump.org/linktypes.html

% Copyright 2019 The MathWorks, Inc.

%#codegen

    properties (Access = private)
        %FID File identifier
        %   Scalar MATLAB integer-valued double, called a file identifier
        FID = -1

        %IsLittleEndian Byte order flag
        %   Set this property to true to indicate that the byte order is
        %   'little-endian'. The default value is true.
        IsLittleEndian (1, 1) logical = true

        %IsHeader Header flag
        %   Set this property to true to indicate that the section header
        %   block is present. The default value is false.
        IsHeader (1, 1) logical = false
    end

    properties (SetAccess = private, Hidden)
        %InterfacesCount Number of interface description blocks written
        %into PCAPNG format file
        %   Represents the count of interface description blocks written.
        %   It increments with the addition of new interface blocks.
        InterfacesCount = 0
    end

    properties (GetAccess = public, SetAccess = private)
        %FileName Name of the file
        %   It is specified as a character vector or a string. The default
        %   file name is 'capture'.
        FileName = 'capture'

        %FileComment Comment for the file
        %   It is specified as a character vector or a string. The default
        %   value is an empty character.
        FileComment = blanks(0)

        %ByteOrder Order of bytes
        %   Specifies the sequential order in which bytes are arranged into
        %   larger numerical values. The default value is 'little-endian'.
        ByteOrder = 'little-endian';
    end

    methods (Access = private)
        function sectionHeaderBlock = createSectionHeaderBlock(obj, varargin)
            %createSectionHeaderBlock Create a section header block and
            %return it as a byte array

            % Used to detect the block format
            blockType = obj.formByteArray([10; 13; 13; 10]);
            % Used to detect the byte ordering. Write in big-endian format
            byteOrderMagic = obj.formByteArray([26; 43; 60; 77]);
            % Major Version Number
            majorVersion = obj.formByteArray([0; 1]);
            % Minor Version Number
            minorVersion = [0; 0];
            % Section length in octets of this section, excluding the
            % section header block itself
            sectionLength = obj.formByteArray([255; 255; 255; 255; 255; 255; 255; 255]);
            % Application Code
            userApplicationCode = [obj.formByteArray([0; 4]); obj.formByteArray([0; 6])];
            % Application used to log packets into a file. The user
            % application is "MATLAB"
            userApplication = [77; 65; 84; 76; 65; 66; 0; 0];
            % Option delimits the end of the optional fields
            endOfOptions = [0; 0; 0; 0];
            % Default block length
            defLength = 44;

            % Check if comment is provided as an input
            if nargin > 1
                comment = varargin{1};
                % Add Comments, if any, for the section header block
                commentLength = numel(comment);
                comment = obj.getByteArray(comment, commentLength);
                % The value [0; 1] represents that the field is comment
                commentCode = obj.formByteArray([0; 1]);
                % Calculate the length of comment
                commentLengthDec = obj.getByteArray(commentLength, 2);
                commentCode = [commentCode; commentLengthDec];
                paddingComment = obj.getPadBytes(comment);
                paddingLength = length(paddingComment);
                commentBytes = [commentCode; comment; paddingComment];
                % Get total length of section header block
                defLength = defLength+4;
                totalLength = obj.getByteArray(commentLength+paddingLength+defLength, 4);
            else
                commentBytes = [];
                % Get total length of section header block
                totalLength = obj.getByteArray(defLength, 4);
            end

            % Section header block
            sectionHeaderBlock  = [blockType; ...
                totalLength; ...
                byteOrderMagic; ...
                majorVersion; ...
                minorVersion; ...
                sectionLength; ...
                commentBytes; ...
                userApplicationCode; ...
                userApplication; ...
                endOfOptions; ...
                totalLength];
        end

        function [interfaceDescriptionBlock, interfaceID] = createInterfaceDescriptionBlock(obj, linkType, interface)
            %createInterfaceDescriptionBlock Create an interface
            %description block as a byte array. Return the interface block
            %and interface ID

            % Used to detect the block format
            blockType = obj.formByteArray([0; 0; 0; 1]);
            % Value that defines the link layer type of this interface
            linkTypeBytes = obj.getByteArray(linkType, 2);
            % Must be filled with 0, and ignored by PCAPNG file readers
            reserved = [0; 0];
            % Maximum number of octets captured from each packet
            snapLen = obj.formByteArray([0; 0; 255; 255]);
            % Interface code
            interfaceCode = [obj.formByteArray([0; 2]); obj.getByteArray(numel(interface), 2)];
            % Interface name in bytes
            interfaceLength = numel(interface);
            interfaceName = obj.getByteArray(interface, interfaceLength);
            % Padding for interface
            interfacePadding  = obj.getPadBytes(interfaceName);
            % Option delimits the end of the optional fields
            endOfOptions = [0; 0; 0; 0];
            % Default block length
            defLength = 28;

            % Get total length of interface description block
            totalLength = obj.getByteArray(defLength+numel(interfaceName)+numel(interfacePadding), 4);

            % Interface block and the corresponding interface ID
            interfaceDescriptionBlock = [blockType; ...
                totalLength; ...
                linkTypeBytes; ...
                reserved; ...
                snapLen; ...
                interfaceCode; ...
                interfaceName; ...
                interfacePadding; ...
                endOfOptions; ...
                totalLength];
            interfaceID = obj.InterfacesCount;

            % Increment the interface count if a new interface is added
            obj.InterfacesCount = obj.InterfacesCount+1;
        end

        function enhancedPacketBlock = createEnhancedPacketBlock(obj, interfaceID, packetData, timeStamp, varargin)
            %createEnhancedPacketBlock Create enhanced packet block and
            %return it as a byte array

            % Block type of enhanced packet block
            blockType = obj.formByteArray([0; 0; 0; 6]);
            % Interface ID specifies the interface this packet comes from
            interfaceIDBytes = obj.getByteArray(interfaceID, 4);
            % Total time in microseconds
            totalTime = obj.getByteArray((fix(timeStamp/1000000))*1000000+fix(rem(timeStamp, 1000000)), 8);
            if obj.IsLittleEndian
                % Timestamp high takes first four bytes of total time
                timestampHigh = totalTime(5:8);
                % Timestamp low takes last four bytes of total time
                timestampLow = totalTime(1:4);
            else
                % Timestamp high takes first four bytes of total time
                timestampHigh = totalTime(1:4);
                % Timestamp low takes last four bytes of total time
                timestampLow = totalTime(5:8);
            end
            % Maximum length of packet that can be captured
            maxCaptureLen = 65535;
            % Number of octets captured, it is the minimum value of packet
            % data and the snapshot length
            maxPacketLength = min(numel(packetData), maxCaptureLen);
            packetData = packetData(1:maxPacketLength, 1);
            capturedLength = obj.getByteArray(maxPacketLength, 4);
            % Actual length of the packet when it was transmitted on the
            % network
            packetLength = obj.getByteArray(numel(packetData), 4);
            % End of options
            endOfOptions = [0; 0; 0; 0];
            % Default block length
            defLength = 32;

            % Get the data to be padded for packet data
            if isempty(packetData)
                dataPadding = [];
            else
                dataPadding = obj.getPadBytes(packetData);
            end

            % Check if comment is provided as an input
            if nargin > 4
                % Add Comments, if any, for the enhanced packet block
                comment = varargin{1};
                commentLength = numel(comment);
                comment = obj.getByteArray(comment, commentLength);
                % The value [0; 1] represents that the field is comment
                commentCode = obj.formByteArray([0; 1]);
                % Calculate the length of comment
                commentLengthDec = obj.getByteArray(commentLength, 2);
                commentCode = [commentCode; commentLengthDec];
                % Get the data to be padded for comment
                commentPading = obj.getPadBytes(comment);
                paddingLength = length(commentPading);
                commentBytes = [commentCode; comment; commentPading; endOfOptions];
                % Get total length of the enhanced packet block
                defLength = defLength+8;
                totalLength = obj.getByteArray(defLength+numel(dataPadding)+numel(packetData(:, 1))+...
                    paddingLength+commentLength, 4);
            else
                commentBytes = [];
                totalLength = obj.getByteArray(defLength+numel(dataPadding)+numel(packetData(:, 1)), 4);
            end

            % Enhanced packet block
            enhancedPacketBlock = [blockType; ...
                totalLength; ...
                interfaceIDBytes; ...
                timestampHigh; ...
                timestampLow; ...
                capturedLength; ...
                packetLength; ...
                packetData;...
                dataPadding; ...
                commentBytes; ...
                totalLength];
        end

        function customBlock = createCustomBlock(obj, customData)
            %createCustomBlock Create a custom block with custom data

            % Block type of custom block
            blockTypeCB = obj.formByteArray([0; 0; 173; 11]);
            % Private enterprise number (to be obtained from IANA)
            PEN = [0; 0; 0; 0];
            % Default block length
            defLength = 16;
            % Custom data length
            customDataLength = numel(customData);
            if isempty(customData)
                customData = [];
            else
                customData = obj.getByteArray(customData, customDataLength);
            end
            % Get the data to be padded for packet data
            if isempty(customData)
                paddingData = [];
            else
                paddingData = obj.getPadBytes(customData);
            end

            % Total length of the block
            totalLength = obj.getByteArray(customDataLength+numel(paddingData)+defLength, 4);

            % Custom block
            customBlock = [blockTypeCB; ...
                totalLength; ...
                PEN; ...
                customData; ...
                paddingData; ...
                totalLength];
        end

        function openFile(obj, fileName)
            %openFile Open a file
            %
            %   FILENAME  - Name of the file

            % Open and clear the existing content in the file
            obj.FID = fopen(fileName, 'w');

            % Check if file ID is negative
            if obj.FID < 0
                error('Unable to open file %s because value of its file ID is less than zero.', fileName);
            end
        end

        function closeFile(obj)
            %closeFile Close the file

            if obj.FID > 0
                fclose(obj.FID);
                obj.FID = -1;
            end
        end

        function delete(obj)
            %delete Function acts as a destructor

            obj.closeFile;
        end
    end

    methods
        function obj = HelperPCAPNGWriter(varargin)
            %HelperPCAPNGWriter Create a PCAPNG packet writer configuration
            %object.

            % Name-value pair check
            if mod(nargin,2)
                error('Incorrect number of input arguments. Number of input arguments must be even.');
            end

            % Apply name-value pairs
            for i = 1:2:length(varargin)
                if strcmp('FileName', varargin{i}) || strcmp('FileComment', varargin{i}) || strcmp('ByteOrder', varargin{i})
                    obj.(varargin{i}) = varargin{i+1};
                end
            end

            obj.ByteOrder = char(obj.ByteOrder);
            if ~strcmp(obj.ByteOrder, 'little-endian')
                obj.IsLittleEndian = false;
                obj.IsLittleEndianUtil = false;
            end

            % Initialize interface count, before writing any blocks
            obj.InterfacesCount = 0;
        end

        function set.FileName(obj, value)
            validateattributes(value, {'char' ,'string'}, {'row', 'nonempty'}, ...
                    mfilename, 'FileName');
            obj.FileName = value;
        end

        function set.FileComment(obj, value)
            validateattributes(value, {'char' ,'string'}, {'row'}, ...
                    mfilename, 'FileComment');
            obj.FileComment = value;
        end

        function set.ByteOrder(obj, value)
            validatestring(value, {'little-endian', 'big-endian'}, ...
                    mfilename, 'ByteOrder');
            obj.ByteOrder = value;
        end

        function interfaceID = writeInterfaceDescriptionBlock(obj, linkType, interface)
            %writeInterfaceDescriptionBlock Write interface description
            %block into the PCAPNG file and return the interface ID. This
            %interface ID has to be used for logging packets of specified
            %interface.
            %
            %   LINKTYPE  - Protocol number
            %   INTERFACE - Interface on which the packet is captured

            % Write section header block
            if ~obj.IsHeader
                % Form file name
                fileName = [char(obj.FileName), '.',  'pcapng'];
                % Open file for writing packets
                openFile(obj, fileName);
                if isempty(char(obj.FileComment))
                    sectionHeaderBlockBytes = obj.createSectionHeaderBlock;
                else
                    % Write section header block at the start of the file
                    sectionHeaderBlockBytes = obj.createSectionHeaderBlock(char(obj.FileComment));
                end
                fwrite(obj.FID, sectionHeaderBlockBytes);
                % Set the header flag
                obj.IsHeader = true;
            end
            [InterfaceDescriptionBlock, interfaceID] = obj.createInterfaceDescriptionBlock(linkType, interface);
            fwrite(obj.FID, InterfaceDescriptionBlock);
        end

        function write(obj, timestamp, interfaceID, packet, varargin)
            %write Write enhanced packet block into PCAPNG format file
            %
            %   write(OBJ, TIMESTAMP, INTERFACEID, PACKET) writes packet
            %   into PCAPNG format file
            %
            %   TIMESTAMP   - Packet arrival time in microseconds. It is
            %                 specified as a double value. The default
            %                 value is zero.
            %   INTERFACEID - Unique identifier for an interface
            %   PACKET      - Packet is specified as one of the following
            %                 types:
            %                   - A character vector representing octets in
            %                     hexadecimal format.
            %                   - A string scalar representing octets in
            %                     hexadecimal format.
            %                   - A numeric vector, where each element is
            %                     in the range of [0 - 255] inclusive,
            %                     representing octets in decimal format.
            %                   - An n-by-2 character array, where each row
            %                     represents an octet in hexadecimal
            %                     format.
            %
            %   write(OBJ, TIMESTAMP, INTERFACEID, PACKET, COMMENT) writes
            %   packet along with comment into PCAPNG file
            %
            %   COMMENT     - Packet comment in PCAPNG format. It is
            %                 specified as a character vector or a string.
            %                 The default value is an empty character.

            % Validate packet and return octets in decimal format
            packetData = obj.validatePayloadFormat(packet);

            if interfaceID < 0 || interfaceID > obj.InterfacesCount
                error('Unrecognized interface ID %d. Interface ID must be between 0 and total number of interfaces added.', interfaceID);
            end

            if obj.FID < 0
                error('Unable to write packet. To write a packet in PCAPNG file, add interface description block.');
            end

            if nargin > 4
                enhancedPacketBlock = obj.createEnhancedPacketBlock(interfaceID, packetData, timestamp, char(varargin{1}));
            else
                enhancedPacketBlock = obj.createEnhancedPacketBlock(interfaceID, packetData, timestamp);
            end

            fwrite(obj.FID, enhancedPacketBlock);
        end

        function writeCustomBlock(obj, customData)
            %writeCustomBlock Write custom block into a PCAPNG  format file
            %
            %   writeCustomBlock(OBJ, CUSTOMDATA) writes custom block into
            %   PCAPNG format file
            %   CUSTOMDATA  - Any user defined data

            customBlock = obj.createCustomBlock(char(customData));
            fwrite(obj.FID, customBlock);
        end
    end
end
