/* Copyright (C) 2020-2022 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.
 */

// Author: Frank Honza <frank.honza@dcso.de>

use nom7::bytes::streaming::tag;
use nom7::bytes::streaming::take;
use nom7::combinator::map_res;
use nom7::number::streaming::*;
use nom7::*;
use std::fmt;
use std::str;

#[derive(Debug, PartialEq)]
pub enum RFBGlobalState {
    TCServerProtocolVersion,
    TCSupportedSecurityTypes,
    TCVncChallenge,
    TCServerInit,
    TCFailureReason,
    TSClientProtocolVersion,
    TCServerSecurityType,
    TSSecurityTypeSelection,
    TSVncResponse,
    TCSecurityResult,
    TSClientInit,
    Skip,
}

impl fmt::Display for RFBGlobalState {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            RFBGlobalState::TCServerProtocolVersion => write!(f, "TCServerProtocolVersion"),
            RFBGlobalState::TCSupportedSecurityTypes => write!(f, "TCSupportedSecurityTypes"),
            RFBGlobalState::TCVncChallenge => write!(f, "TCVncChallenge"),
            RFBGlobalState::TCServerInit => write!(f, "TCServerInit"),
            RFBGlobalState::TCFailureReason => write!(f, "TCFailureReason"),
            RFBGlobalState::TSClientProtocolVersion => write!(f, "TSClientProtocolVersion"),
            RFBGlobalState::TSSecurityTypeSelection => write!(f, "TSSecurityTypeSelection"),
            RFBGlobalState::TSVncResponse => write!(f, "TSVncResponse"),
            RFBGlobalState::TCSecurityResult => write!(f, "TCSecurityResult"),
            RFBGlobalState::TCServerSecurityType => write!(f, "TCServerSecurityType"),
            RFBGlobalState::TSClientInit => write!(f, "TSClientInit"),
            RFBGlobalState::Skip => write!(f, "Skip"),
        }
    }
}

pub struct ProtocolVersion {
    pub major: String,
    pub minor: String,
}

pub struct SupportedSecurityTypes {
    pub number_of_types: u8,
    pub types: Vec<u8>,
}

pub struct SecurityTypeSelection {
    pub security_type: u8,
}

pub struct ServerSecurityType {
    pub security_type: u32,
}

pub struct SecurityResult {
    pub status: u32,
}

pub struct FailureReason {
    pub reason_string: String,
}

pub struct VncAuth {
    pub secret: Vec<u8>,
}

pub struct ClientInit {
    pub shared: u8,
}

pub struct PixelFormat {
    pub bits_per_pixel: u8,
    pub depth: u8,
    pub big_endian_flag: u8,
    pub true_colour_flag: u8,
    pub red_max: u16,
    pub green_max: u16,
    pub blue_max: u16,
    pub red_shift: u8,
    pub green_shift: u8,
    pub blue_shift: u8,
}

pub struct ServerInit {
    pub width: u16,
    pub height: u16,
    pub pixel_format: PixelFormat,
    pub name_length: u32,
    pub name: Vec<u8>,
}

pub fn parse_protocol_version(i: &[u8]) -> IResult<&[u8], ProtocolVersion> {
    let (i, _) = tag("RFB ")(i)?;
    let (i, major) = map_res(take(3_usize), str::from_utf8)(i)?;
    let (i, _) = tag(".")(i)?;
    let (i, minor) = map_res(take(3_usize), str::from_utf8)(i)?;
    let (i, _) = tag("\n")(i)?;
    Ok((
        i,
        ProtocolVersion {
            major: major.to_string(),
            minor: minor.to_string(),
        },
    ))
}

pub fn parse_supported_security_types(i: &[u8]) -> IResult<&[u8], SupportedSecurityTypes> {
    let (i, number_of_types) = be_u8(i)?;
    let (i, types) = take(number_of_types as usize)(i)?;
    Ok((
        i,
        SupportedSecurityTypes {
            number_of_types,
            types: types.to_vec(),
        },
    ))
}

pub fn parse_server_security_type(i: &[u8]) -> IResult<&[u8], ServerSecurityType> {
    let (i, security_type) = be_u32(i)?;
    Ok((i, ServerSecurityType { security_type }))
}

pub fn parse_vnc_auth(i: &[u8]) -> IResult<&[u8], VncAuth> {
    let (i, secret) = take(16_usize)(i)?;
    Ok((
        i,
        VncAuth {
            secret: secret.to_vec(),
        },
    ))
}

pub fn parse_security_type_selection(i: &[u8]) -> IResult<&[u8], SecurityTypeSelection> {
    let (i, security_type) = be_u8(i)?;
    Ok((i, SecurityTypeSelection { security_type }))
}

pub fn parse_security_result(i: &[u8]) -> IResult<&[u8], SecurityResult> {
    let (i, status) = be_u32(i)?;
    Ok((i, SecurityResult { status }))
}

pub fn parse_failure_reason(i: &[u8]) -> IResult<&[u8], FailureReason> {
    let (i, reason_length) = be_u32(i)?;
    let (i, reason_string) = map_res(take(reason_length as usize), str::from_utf8)(i)?;
    Ok((
        i,
        FailureReason {
            reason_string: reason_string.to_string(),
        },
    ))
}

pub fn parse_client_init(i: &[u8]) -> IResult<&[u8], ClientInit> {
    let (i, shared) = be_u8(i)?;
    Ok((i, ClientInit { shared }))
}

pub fn parse_pixel_format(i: &[u8]) -> IResult<&[u8], PixelFormat> {
    let (i, bits_per_pixel) = be_u8(i)?;
    let (i, depth) = be_u8(i)?;
    let (i, big_endian_flag) = be_u8(i)?;
    let (i, true_colour_flag) = be_u8(i)?;
    let (i, red_max) = be_u16(i)?;
    let (i, green_max) = be_u16(i)?;
    let (i, blue_max) = be_u16(i)?;
    let (i, red_shift) = be_u8(i)?;
    let (i, green_shift) = be_u8(i)?;
    let (i, blue_shift) = be_u8(i)?;
    let (i, _) = take(3_usize)(i)?;
    let format = PixelFormat {
        bits_per_pixel,
        depth,
        big_endian_flag,
        true_colour_flag,
        red_max,
        green_max,
        blue_max,
        red_shift,
        green_shift,
        blue_shift,
    };
    Ok((i, format))
}

pub fn parse_server_init(i: &[u8]) -> IResult<&[u8], ServerInit> {
    let (i, width) = be_u16(i)?;
    let (i, height) = be_u16(i)?;
    let (i, pixel_format) = parse_pixel_format(i)?;
    let (i, name_length) = be_u32(i)?;
    let (i, name) = take(name_length as usize)(i)?;
    let init = ServerInit {
        width,
        height,
        pixel_format,
        name_length,
        name: name.to_vec(),
    };
    Ok((i, init))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Simple test of some valid data.
    #[test]
    fn test_parse_version() {
        let buf = b"RFB 003.008\n";

        let result = parse_protocol_version(buf);
        match result {
            Ok((remainder, message)) => {
                // Check the first message.
                assert_eq!(message.major, "003");

                // And we should have 0 bytes left.
                assert_eq!(remainder.len(), 0);
            }
            Err(Err::Incomplete(_)) => {
                panic!("Result should not have been incomplete.");
            }
            Err(Err::Error(err)) | Err(Err::Failure(err)) => {
                panic!("Result should not be an error: {:?}.", err);
            }
        }
    }

    #[test]
    fn test_parse_server_init() {
        let buf = [
            0x05, 0x00, 0x03, 0x20, 0x20, 0x18, 0x00, 0x01, 0x00, 0xff, 0x00, 0xff, 0x00, 0xff,
            0x10, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x1e, 0x61, 0x6e, 0x65, 0x61,
            0x67, 0x6c, 0x65, 0x73, 0x40, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x68, 0x6f, 0x73, 0x74,
            0x2e, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e,
        ];

        let result = parse_server_init(&buf);
        match result {
            Ok((remainder, message)) => {
                // Check the first message.
                assert_eq!(message.width, 1280);
                assert_eq!(message.height, 800);
                assert_eq!(message.pixel_format.bits_per_pixel, 32);

                // And we should have 0 bytes left.
                assert_eq!(remainder.len(), 0);
            }
            Err(Err::Incomplete(_)) => {
                panic!("Result should not have been incomplete.");
            }
            Err(Err::Error(err)) | Err(Err::Failure(err)) => {
                panic!("Result should not be an error: {:?}.", err);
            }
        }
    }

    #[test]
    fn test_parse_pixel_format() {
        let buf = [
            0x20, /* Bits per pixel: 32 */
            0x18, /* Depth: 24 */
            0x00, /* Big endian flag: False */
            0x01, /* True color flag: True */
            0x00, 0xff, /* Red maximum: 255 */
            0x00, 0xff, /* Green maximum: 255 */
            0x00, 0xff, /* Blue maximum: 255 */
            0x10, /* Red shift: 16 */
            0x08, /* Green shift: 8 */
            0x00, /* Blue shift: 0 */
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0xa0,
        ];

        let result = parse_pixel_format(&buf);
        match result {
            Ok((remainder, message)) => {
                assert_eq!(message.bits_per_pixel, 32);
                assert_eq!(message.depth, 24);
                assert_eq!(message.big_endian_flag, 0);
                assert_eq!(message.true_colour_flag, 1);
                assert_eq!(message.red_max, 255);
                assert_eq!(message.green_max, 255);
                assert_eq!(message.blue_max, 255);
                assert_eq!(message.red_shift, 16);
                assert_eq!(message.green_shift, 8);
                assert_eq!(message.blue_shift, 0);
                assert_eq!(remainder.len(), 5);
            }
            Err(Err::Incomplete(_)) => {
                panic!("Result should not have been incomplete.");
            }
            Err(Err::Error(err)) | Err(Err::Failure(err)) => {
                panic!("Result should not be an error: {:?}.", err);
            }
        }
    }

    #[test]
    fn test_parse_supported_security_types() {
        let buf = [
            0x01, /* Number of security types: 1 */
            0x02, /* Security type: VNC (2) */
            0x00, 0x01, 0x00, 0xff, 0x00, 0xff, 0x00, 0xff, 0x10, 0x08, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x01, 0xa0,
        ];

        let result = parse_supported_security_types(&buf);
        match result {
            Ok((remainder, message)) => {
                assert_eq!(message.number_of_types, 1);
                assert_eq!(message.types[0], 2);
                assert_eq!(remainder.len(), 19);
            }
            Err(Err::Incomplete(_)) => {
                panic!("Result should not have been incomplete.");
            }
            Err(Err::Error(err)) | Err(Err::Failure(err)) => {
                panic!("Result should not be an error: {:?}.", err);
            }
        }
    }

    #[test]
    fn test_parse_vnc_auth() {
        let buf = [
            0x54, 0x7b, 0x7a, 0x6f, 0x36, 0xa1, 0x54, 0xdb, 0x03, 0xa2, 0x57, 0x5c, 0x6f, 0x2a,
            0x4e, 0xc5, /* Authentication challenge: 547b7a6f36a154db03a2575c6f2a4ec5 */
            0x00, 0x00, 0x00, 0x01, 0xa0,
        ];

        let result = parse_vnc_auth(&buf);
        match result {
            Ok((remainder, message)) => {
                assert_eq!(
                    hex::encode(message.secret),
                    "547b7a6f36a154db03a2575c6f2a4ec5"
                );
                assert_eq!(remainder.len(), 5);
            }
            Err(Err::Incomplete(_)) => {
                panic!("Result should not have been incomplete.");
            }
            Err(Err::Error(err)) | Err(Err::Failure(err)) => {
                panic!("Result should not be an error: {:?}.", err);
            }
        }
    }

    #[test]
    fn test_parse_client_init() {
        let buf = [
            0x00, /*Share desktop flag: False*/
            0x7b, 0x7a, 0x6f, 0x36, 0xa1, 0x54, 0xdb, 0x03, 0xa2, 0x57, 0x5c, 0x6f, 0x2a, 0x4e,
            0xc5, 0x00, 0x00, 0x00, 0x01, 0xa0,
        ];

        let result = parse_client_init(&buf);
        match result {
            Ok((remainder, message)) => {
                assert_eq!(message.shared, 0);
                assert_eq!(remainder.len(), 20);
            }
            Err(Err::Incomplete(_)) => {
                panic!("Result should not have been incomplete.");
            }
            Err(Err::Error(err)) | Err(Err::Failure(err)) => {
                panic!("Result should not be an error: {:?}.", err);
            }
        }
    }
}