/* Copyright (c) 2018 Debily
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in the
 * Software without restriction, including without limitation the rights to use, copy,
 * modify, merge, publish, distribute, sublicense, and/or sell copies of the
 * Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
 * KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
 * WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
 * PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS
 * OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
 * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

extern crate rand;
use self::rand::{Rng, OsRng};
use std;

/// The `MutatorEngine` struct used to mutate fuzzing data
pub struct MutatorEngine {
    rng: OsRng,
    mutationgrade: f64,
    truncationgrade: f64,
}

impl MutatorEngine {
    /// Creates a new `MutatorEngine` struct
    pub fn new() -> Result<MutatorEngine, std::io::Error> {
        Ok(MutatorEngine {
            rng: OsRng::new()?,
            mutationgrade: 0.01,
            truncationgrade: 0.999,
        })
    }

    /// If you don't need specialized mutation you can use
    /// this function to choose the kind of
    /// mutation at random.
    #[inline]
    pub fn mutate(&mut self, data: &mut Vec<u8>) {
        match rand_usize(0, 4, &mut self.rng) {
            0 => self.special_ints(data),
            1 => self.bitflipping(data),
            2 => self.byteflipping(data),
            3 => self.add_substract_binary(data),
            4 => self.chunk_spew(data),
            _ => unreachable!(),
        }
    }

    /// Flips bits at random in `data`.
    #[inline]
    pub fn bitflipping(&mut self, data: &mut Vec<u8>) {
        let mut count = ((data.len() as f64 * 8.0) * self.mutationgrade) as i64;
        let mut bit: usize;
        let mut idx_bit: usize;
        let mut idx_byte: usize;
        let mut data_len: usize;

        if count == 0 {
            count = 1;
        }


        for _ in 0..count {
            data_len = data.len();

            if data_len <= 1 {
                self.append_data(data);
                continue;
            }

            bit = rand_usize(0, data_len * 8 - 1, &mut self.rng);
            idx_bit = bit % 8;
            idx_byte = bit / 8;

            if let Some(idx) = data.get_mut(idx_byte) {
                *idx ^= 1 << idx_bit;
            }
        }
    }

    /// This function flips bytes in `data`
    #[inline]
    pub fn byteflipping(&mut self, data: &mut Vec<u8>) {
        let mut count = (data.len() as f64 * self.mutationgrade) as usize;
        let mut index: usize;

        if count == 0 {
            count = 1;
        }
        for _ in 0..count {
            index = rand_usize(
                0,
                if data.len() <= 1 {
                    self.append_data(data);
                    continue;
                } else {
                    data.len() - 1
                },
                &mut self.rng,
            );
            if index < data.len() {
                data[index] = rand_u8(&mut self.rng);
            }
        }
    }

    /// In this function some special integers will be
    /// placed in `data`.
    #[inline]
    pub fn special_ints(&mut self, data: &mut Vec<u8>) {
        let byte: [u8; 3] = [0xff, 0x7f, 0];
        let short: [[u8; 2]; 2] = [[0xff, 0xff], [0x0, 0x0]];
        let int32: [[u8; 4]; 5] = [
            [0xff, 0xff, 0xff, 0xff],
            [0, 0, 0, 0],
            [0x80, 0, 0, 0],
            [0x40, 0, 0, 0],
            [0x7f, 0xff, 0xff, 0xff],
        ];
        let mut count = (data.len() as f64 * self.mutationgrade) as usize;
        let mut index: usize;
        let mut n_size: u8;
        let mut sz: i64;
        let mut result: usize;

        if count == 0 {
            count = 1;
        }
        for _ in 0..count {
            result = rand_usize(0, 2, &mut self.rng);
            if result == 0 {
                /* u8 */
                index = rand_usize(
                    0,
                    if data.len() == 0 {
                        self.append_data(data);
                        continue;
                    } else {
                        if data.len() > 0 {
                            data.len() - 1
                        } else {
                            continue;
                        }
                    },
                    &mut self.rng,
                );
                data[index] = byte[rand_usize(0, 2, &mut self.rng)];
            } else if result == 1 {
                /* u16 */
                n_size = 2;
                sz = data.len() as i64 - n_size as i64;
                if sz <= 0 {
                    continue;
                }
                index = rand_usize(0, sz as usize, &mut self.rng);
                for i in short[rand_usize(0, 1, &mut self.rng)].iter() {
                    if index < data.len() {
                        data[index] = *i;
                        index += 1;
                    }
                }
            } else {
                /* u32 */
                n_size = 4;
                sz = data.len() as i64 - n_size as i64;
                if sz <= 0 {
                    continue;
                }
                index = rand_usize(0, sz as usize, &mut self.rng);
                for i in int32[rand_usize(0, 1, &mut self.rng)].iter() {
                    if index < data.len() {
                        data[index] = *i;
                        index += 1;
                    }
                }
            }
        }
    }

    /// In this function values in `data` will be decremented
    /// or incremented a random number of times.
    #[inline]
    pub fn add_substract_binary(&mut self, data: &mut Vec<u8>) {
        let mut count = (data.len() as f64 * self.mutationgrade) as usize;
        let mut substract: usize;
        let mut index: usize;
        let mut un8: u8;

        if count == 0 {
            count = 1;
        }
        for _ in 0..count {
            substract = rand_usize(0, 1, &mut self.rng);
            index = rand_usize(
                0,
                if data.len() == 0 {
                    self.append_data(data);
                    continue;
                } else {
                    if data.len() > 0 {
                        data.len() - 1
                    } else {
                        continue;
                    }
                },
                &mut self.rng,
            );
            un8 = rand_u8(&mut self.rng);
            if substract == 1 {
                if data[index] >= un8 {
                    data[index] = data[index] - un8;
                }
            } else {
                if (data[index] as u32 + un8 as u32) < (u8::max_value() as u32) {
                    data[index] = data[index] + un8;
                }
            }
        }
    }

    /// This function will replace a value in `data` with
    /// another one in `data`.<br>
    /// Here is an example:<br>
    /// [1,2,3,4,5,6,7,8,9,0]<br>
    ///  <- - - - - - -<br>
    /// [6,2,3,4,5,6,7,8,9,0]<br>
    /// In the above example 6 has been copied where 1 has
    /// previously been.
    #[inline]
    pub fn chunk_spew(&mut self, data: &mut Vec<u8>) {
        let mut count = (data.len() as f64 * self.mutationgrade) as usize;
        let mut offset_sz_src: usize;
        let mut offsetsrc: usize;
        let mut offsetdest: usize;

        if count == 0 {
            count = 1;
        }
        for _ in 0..count {
            offset_sz_src = rand_usize(
                0,
                if data.len() / 2 == 0 {
                    self.append_data(data);
                    continue;
                } else {
                    data.len() / 2
                },
                &mut self.rng,
            );
            offsetsrc = rand_usize(
                0,
                if (data.len() as i64 - offset_sz_src as i64) < 0 {
                    self.append_data(data);
                    continue;
                } else {
                    data.len() - offset_sz_src
                },
                &mut self.rng,
            );
            offsetdest = rand_usize(
                0,
                if (data.len() as i64 - offset_sz_src as i64) < 0 {
                    self.append_data(data);
                    continue;
                } else {
                    data.len() - offset_sz_src
                },
                &mut self.rng,
            );
            for _ in 0..offset_sz_src {
                if offsetsrc < data.len() && offsetdest < data.len() {
                    /* Using the offsets like raw pointers */
                    data[offsetdest] = data[offsetsrc];
                    offsetdest += 1;
                    offsetsrc += 1;
                } else {
                    /* Errors do not matter, we will just save time and break to continue */
                    break;
                }
            }
        }
    }

    /// Appends random bytes to `data`.
    /// # Remarks
    /// Do not use this function too often,
    /// if `data` becomes too long it will take too long to
    /// process it with the other functions since it takes
    /// exponentially longer to process.
    /// This function will be called automagically if a function
    /// detects that `data` is too short.
    #[inline]
    pub fn append_data(&mut self, data: &mut Vec<u8>) {
        let mut count = (data.len() as f64 * self.mutationgrade) as usize;

        if count == 0 {
            count = 1;
        }
        for _ in 0..count {
            data.push(rand_u8(&mut self.rng));
        }
    }

    /// Converts `data` to valid UTF-8.
    /// This can be used to make fuzzing parsers that only
    /// accept UTF-8 encoding more effective.
    #[inline]
    pub fn to_utf8(&mut self, data: &Vec<u8>) -> Vec<u8> {
        /*
         * Using the std library implementation because it gets much more maintenance
         * and bug fixes.
         */
        Vec::from(String::from_utf8_lossy(&data[..]).as_bytes())
    }

    /// Shrinks `data` to a smaller size.
    /// # Remarks
    /// Do not use this function too often, in a loop it
    /// will keep the length of `data` at 0.
    #[allow(dead_code)]
    /*
     * Since this function might not
     * be called at all
     */
    #[inline]
    pub fn truncate(&mut self, data: &mut Vec<u8>) {
        let mut count: usize;

        if data.len() == 0 {
            return;
        }
        count = (data.len() as f64 * self.truncationgrade) as usize;
        if count == 0 {
            count = 1;
        }
        data.truncate(count);
    }

    /// Creates a String based on the data provided in `data`
    #[inline]
    pub fn to_string(&mut self, data: &Vec<u8>) -> Result<String, ()> {
        if let Ok(ret) = std::str::from_utf8(&self.to_utf8(data)[..]) {
            Ok(String::from(ret))
        } else {
            Err(())
        }
    }

    /// Borrows the engine's RNG
    #[inline]
    pub fn borrow_rng(&mut self) -> &mut OsRng {
        &mut self.rng
    }
}

#[inline]
fn rand_usize(min: usize, max: usize, rng: &mut OsRng) -> usize {
    rng.gen_range(min, max + 1)
}


#[inline]
fn rand_u8(rng: &mut OsRng) -> u8 {
    rng.gen_range(0u16, u8::max_value() as u16 + 1) as u8
    /* The cast to u16 is needed because it would not include the highest u8 otherwise.
     * The cast itself is save from overflowing because it can only generate a number between
     * the highest and lowest u8
     */
}

/// The mutator will fuzz itself in the hopes of identifing bugs early.
/// To try it just type `cargo test` in the crates root.
#[test]
fn mutate_tester() {
    let mut data = vec![0xde, 0xad, 0xbe, 0xef];
    let testnum = 500000;

    if let Ok(mut mutatoren) = MutatorEngine::new() {
        mutatoren.mutation_grade(0.199);
        for _ in 0..testnum {
            mutatoren.mutate(&mut data);
            let utf8_vec = mutatoren.to_utf8(&data);
            mutatoren.truncate(&mut data);
            match std::str::from_utf8(&utf8_vec[..]) {
                Ok(_) => {}
                Err(_) => panic!("to_utf8() failed."),
            };
        }
    }
}