/* Copyright (C) 2023 Open Information Security Foundation

 *

 * You can copy, redistribute or modify this Program under the terms of

 * the GNU General Public License version 2 as published by the Free

 * Software Foundation.

 *

 * This program is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU General Public License for more details.

 *

 * You should have received a copy of the GNU General Public License

 * version 2 along with this program; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA

 * 02110-1301, USA.

 */

use nom7::bytes::streaming::take;

use nom7::combinator::cond;

use nom7::number::streaming::{be_u16, be_u32, be_u64, be_u8};

use nom7::IResult;

use suricata_derive::EnumStringU8;

#[derive(Clone, Debug, Default, EnumStringU8)]

#[repr(u8)]

pub enum WebSocketOpcode {

    #[default]

    Continuation = 0,

    Text = 1,

    Binary = 2,

    Close = 8,

    Ping = 9,

    Pong = 10,

}

#[derive(Clone, Debug, Default)]

pub struct WebSocketPdu {

    pub flags: u8,

    pub fin: bool,

    pub compress: bool,

    pub opcode: u8,

    pub mask: Option<u32>,

    pub payload: Vec<u8>,

    pub to_skip: u64,

}

// cf rfc6455#section-5.2

pub fn parse_message(i: &[u8], max_pl_size: u32) -> IResult<&[u8], WebSocketPdu> {

    let (i, flags_op) = be_u8(i)?;

    let fin = (flags_op & 0x80) != 0;

    let compress = (flags_op & 0x40) != 0;

    let flags = flags_op & 0xF0;

    let opcode = flags_op & 0xF;

    let (i, mask_plen) = be_u8(i)?;

    let mask_flag = (mask_plen & 0x80) != 0;

    let (i, payload_len) = match mask_plen & 0x7F {

        126 => {

            let (i, val) = be_u16(i)?;

            Ok((i, val.into()))

        }

        127 => be_u64(i),

        _ => Ok((i, (mask_plen & 0x7F).into())),

    }?;

    let (i, xormask) = cond(mask_flag, take(4usize))(i)?;

    let mask = if mask_flag {

        let (_, m) = be_u32(xormask.unwrap())?;

        Some(m)

    } else {

        None

    };

    // we limit payload_len to u32, so as to build on 32-bit system

    // where we cannot take(usize) with a u64

    let (to_skip, payload_len) = if payload_len < u64::from(max_pl_size) {

        (0, payload_len as u32)

    } else {

        (payload_len - (max_pl_size as u64), max_pl_size)

    };

    let (i, payload_raw) = take(payload_len)(i)?;

    let mut payload = payload_raw.to_vec();

    if let Some(xorkey) = xormask {

        for i in 0..payload.len() {

            payload[i] ^= xorkey[i % 4];

        }

    }

    Ok((

        i,

        WebSocketPdu {

            flags,

            fin,

            compress,

            opcode,

            mask,

            payload,

            to_skip,

        },

    ))

}