addpath("from_matlab_example");

rxFrame1 = trace1.iq(:);
rxFrame2 = trace2.iq(:);
rxFrame3 = trace3.iq(:);
rxFrame4 = trace4.iq(:);

beaconFrames1 = extractBeaconFrames(rxFrame1);
beaconFrames2 = extractBeaconFrames(rxFrame2);
beaconFrames3 = extractBeaconFrames(rxFrame3);
beaconFrames4 = extractBeaconFrames(rxFrame4);

% 1.
disp("Beacon Frame 1")
printBeaconTable(beaconFrames1);
disp("Beacon Frame 2")
printBeaconTable(beaconFrames2);
disp("Beacon Frame 3")
printBeaconTable(beaconFrames3);
disp("Beacon Frame 4")
printBeaconTable(beaconFrames4);

% 2.
%{
Address 1: used to announce the destination station, in a beacon frame its 
           FF:FF:FF:FF:FF:FF (meaning broadcast to all stations)
Address 2: used to remark the source device address of the frame, 
           the ap mac who sent the beacon is inserted
Address 3: used for the bssid of the device sending the beacon, means
           usually address 2 equals address 3, when no virtual ap is used
Address 4: this field does not exists, after the seq ctrl the frame body
           follows containing beacon data instead of an address, like
           timestamp, beacon interval, capability info, ssid, potentially
           the channel number, see the extracts from 2.1
%}

% 3.
%{
Downlink:
Address 1: MAC of the destination station
Address 2: MAC of the AP which further routes the frame
Address 3: MAC of the initial sender which has started the process of 
           the frame to AP's
Address 4: not used
Uplink:
Address 1: MAC of the AP the WiFi frame should route through 
           (the connected AP)
Address 2: MAC of the sending station
Address 3: used to store the final destination MAC address

Address 4: not used

4 Address case:
Address 4: is only used in wireless bridge solutions, environments where
           multiple AP's connect different LAN's to one big one
%}

function beaconFrames = extractBeaconFrames(rxFrame)
    beaconFrames = struct( ...
                           "SSID", {}, ...
                           "BSSID", {}, ...
                           "Offset", {}, ...
                           "MAC_Config", {}, ...
                           "Bits", {}, ...
                           "SNR_dB", {} ...
                         );

    searchOffset = 0;
    index = 1;

    while searchOffset < length(rxFrame)
        oldOffset = searchOffset;
        [bitsData, decParams, searchOffset, res] = recoverOFDMBits(rxFrame, searchOffset);

        if searchOffset <= oldOffset
            searchOffset = oldOffset + 1;
        end

        if isempty(bitsData)
            continue;
        end
    
        [cfgMAC, ~, decodeStatus] = wlanMPDUDecode(bitsData, SuppressWarnings=true);
    
        if ~decodeStatus && matches(cfgMAC.FrameType, "Beacon")
            if isempty(cfgMAC.ManagementConfig.SSID)
                ssid = "Hidden";
            else
                ssid = string(cfgMAC.ManagementConfig.SSID);
            end

            if isfield(res, "PacketOffset")
                beaconFrames(index).Offset = res.PacketOffset;
            else
                beaconFrames(index).Offset = oldOffset;
            end

            if isfield(res, "LLTFSNR")
                beaconFrames(index).SNR_dB = res.LLTFSNR;
            else
                beaconFrames(index).SNR_dB = NaN;
            end

            beaconFrames(index).SSID = ssid;
            beaconFrames(index).BSSID = string(cfgMAC.Address3);
            beaconFrames(index).MAC_Config = cfgMAC;
            beaconFrames(index).Bits = bitsData;

            index = index + 1;
        end
    end
end

function printBeaconTable(beaconFrames)
    beaconCount = numel(beaconFrames);

    tableFormat = "%-20s %-18s %-10s %-8s %-8s %-8s %-12s %-8s %-80s\n";

    if(beaconCount < 1)
        disp("No valid beacons could be extracted")
        return;
    end

    fprintf("\n");
    fprintf(tableFormat, ...
        "SSID", "BSSID", "Interval", "B. ch.", "Op. ch.", "SNR", "Vendor", "Std.", "Rates");

    for i = 1 : beaconCount
        mgmt = beaconFrames(i).MAC_Config.ManagementConfig;

        ssid = string(beaconFrames(i).SSID);
        bssid = toMACAddress(beaconFrames(i).BSSID);
        interval = string(mgmt.BeaconInterval);
        snr = sprintf("%.1f", beaconFrames(i).SNR_dB);
    
        [bch, opch] = extractChannels(mgmt);
        vendor = lookupVendor(beaconFrames(i).BSSID);
        rates = getRates(mgmt);
        std = getWifiStandard(mgmt);
    
        fprintf(tableFormat, ...
            ssid, bssid, interval, bch, opch, snr, vendor, std, rates);
    end
    
    fprintf("\n");
end

function mac = toMACAddress(mac)
    mac = upper(string(mac));
    
    if strlength(mac) ~= 12
        return;
    end
    
    macChars = char(mac);

    mac = string(sprintf("%s:%s:%s:%s:%s:%s", ...
        macChars(1:2), ...
        macChars(3:4), ...
        macChars(5:6), ...
        macChars(7:8), ...
        macChars(9:10), ...
        macChars(11:12)));
end

function rates = getRates(mgmt)
    basicRates = string(mgmt.BasicRates);
    additionalRates = string(mgmt.AdditionalRates);
    
    if isempty(basicRates) && isempty(additionalRates)
        rates = "-";
        return;
    end
    
    basicRates = "[" + strjoin(basicRates, ", ") + "]";    
    rates = strjoin([basicRates, additionalRates], ", ");
end

function vendor = lookupVendor(bssid)
    bssid = upper(string(bssid));
    
    if strlength(bssid) < 6
        vendor = "-";
        return;
    end
    
    oui = extractBefore(bssid, 7);
    
    switch oui
        case { "C884A1", "5CE176" }
            vendor = "Cisco";
        case { "6466B3", "647002" }
            vendor = "TP-Link";
        otherwise
            vendor = "Unknown";
    end
end

function std = getWifiStandard(mgmt)
    rates = [string(mgmt.BasicRates), string(mgmt.AdditionalRates)];
    
    if any(contains(rates, ["6 Mbps", "9 Mbps", "12 Mbps", "18 Mbps", "24 Mbps", "36 Mbps", "48 Mbps", "54 Mbps"]))
        std = "802.11g";
    elseif any(contains(rates, ["1 Mbps", "2 Mbps", "5.5 Mbps", "11 Mbps"]))
        std = "802.11b";
    else
        std = "Unknown";
    end
end

function [b_ch, op_ch] = extractChannels(mgmt)
    b_ch = "-";
    op_ch = "-";

    ies = mgmt.InformationElements;

    if(isempty(ies)) 
        return; 
    end

    for i = 1 : size(ies, 1)
        id = ies{i, 1};
        value = ies{i, 2};

        if isempty(id) || isempty(value)
            continue;
        end

        switch id(1)
            case 3
                b_ch = string(double(value(1)));
            case 61
                op_ch = string(double(value(1)));
        end

        if ~strcmp(b_ch, "-") && ~strcmp(op_ch, "-")
            return;
        end
    end
end