/* Copyright (C) 2019 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: Zach Kelly <zach.kelly@lmco.com>

use crate::rdp::error::RdpError;

use byteorder::ReadBytesExt;

use memchr::memchr;

use nom7::{Err, IResult, Needed};

use std::io::Cursor;

use widestring::U16CString;

/// converts a raw u8 slice of little-endian wide chars into a String

pub fn le_slice_to_string(input: &[u8]) -> Result<String, Box<dyn std::error::Error>> {

    let mut vec = Vec::new();

    let mut cursor = Cursor::new(input);

    while let Ok(x) = cursor.read_u16::<byteorder::LittleEndian>() {

        if x == 0 {

            break;

        }

        vec.push(x);

    }

    match U16CString::new(vec) {

        Ok(x) => match x.to_string() {

            Ok(x) => Ok(x),

            Err(e) => Err(e.into()),

        },

        Err(e) => Err(e.into()),

    }

}

/// converts a raw u8 slice of null-padded utf7 chars into a String, dropping the nulls

pub fn utf7_slice_to_string(input: &[u8]) -> Result<String, Box<dyn std::error::Error>> {

    let s = match memchr(b'\0', input) {

        Some(end) => &input[..end],

        None => input,

    };

    match std::str::from_utf8(s) {

        Ok(s) => Ok(String::from(s)),

        Err(e) => Err(e.into()),

    }

}

/// parses a PER length determinant, to determine the length of the data following

/// x.691-spec: section 10.9

pub fn parse_per_length_determinant(input: &[u8]) -> IResult<&[u8], u32, RdpError> {

    if input.is_empty() {

        // need a single byte to begin length determination

        Err(Err::Incomplete(Needed::new(1)))

    } else {

        let bit7 = input[0] >> 7;

        match bit7 {

            0b0 => {

                // byte starts with 0b0.  Length stored in the lower 7 bits of the current byte

                let length = input[0] as u32 & 0x7f;

                Ok((&input[1..], length))

            }

            _ => {

                let bit6 = (input[0] >> 6) & 0x1;

                match bit6 {

                    0b0 => {

                        // byte starts with 0b10.  Length stored in the remaining 6 bits and the next byte

                        if input.len() < 2 {

                            Err(Err::Incomplete(Needed::new(2)))

                        } else {

                            let length = ((input[0] as u32 & 0x3f) << 8) | input[1] as u32;

                            Ok((&input[2..], length))

                        }

                    }

                    _ => {

                        // byte starts with 0b11.  Without an example to confirm 16K+ lengths are properly

                        // handled, leaving this branch unimplemented

                        Err(Err::Error(RdpError::UnimplementedLengthDeterminant))

                    }

                }

            }

        }

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    use crate::rdp::error::RdpError;

    use nom7::Needed;

    #[test]

    fn test_le_string_abc() {

        let abc = &[0x41, 0x00, 0x42, 0x00, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00];

        assert_eq!(String::from("ABC"), le_slice_to_string(abc).unwrap());

    }

    #[test]

    fn test_le_string_empty() {

        let empty = &[];

        assert_eq!(String::from(""), le_slice_to_string(empty).unwrap());

    }

    #[test]

    fn test_le_string_invalid() {

        let not_utf16le = &[0x00, 0xd8, 0x01, 0x00];

        assert!(le_slice_to_string(not_utf16le).is_err());

    }

    #[test]

    fn test_utf7_string_abc() {

        let abc = &[0x41, 0x42, 0x43, 0x00, 0x00];

        assert_eq!(String::from("ABC"), utf7_slice_to_string(abc).unwrap());

    }

    #[test]

    fn test_utf7_string_empty() {

        let empty = &[];

        assert_eq!(String::from(""), utf7_slice_to_string(empty).unwrap());

    }

    #[test]

    fn test_utf7_string_invalid() {

        let not_utf7 = &[0x80];

        assert!(utf7_slice_to_string(not_utf7).is_err());

    }

    #[test]

    fn test_length_single_length() {

        let bytes = &[0x28];

        assert_eq!(Ok((&[][..], 0x28)), parse_per_length_determinant(bytes));

    }

    #[test]

    fn test_length_double_length() {

        let bytes = &[0x81, 0x28];

        assert_eq!(Ok((&[][..], 0x128)), parse_per_length_determinant(bytes));

    }

    #[test]

    fn test_length_single_length_incomplete() {

        let bytes = &[];

        assert_eq!(

            Err(Err::Incomplete(Needed::new(1))),

            parse_per_length_determinant(bytes)

        )

    }

    #[test]

    fn test_length_16k_unimplemented() {

        let bytes = &[0xc0];

        assert_eq!(

            Err(Err::Error(RdpError::UnimplementedLengthDeterminant)),

            parse_per_length_determinant(bytes)

        )

    }

    #[test]

    fn test_length_double_length_incomplete() {

        let bytes = &[0x81];

        assert_eq!(

            Err(Err::Incomplete(Needed::new(2))),

            parse_per_length_determinant(bytes)

        )

    }

}