// genfuzz
// A performant general purpose fuzzer
// © Jean Pereira <counterswarm.de>

use base64::{engine::general_purpose, Engine as _};
use clipboard::ClipboardContext;
use clipboard::ClipboardProvider;
use rand::seq::SliceRandom;
use rand::{thread_rng, Rng};
use std::fs;
use std::fs::File;
use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use std::path::PathBuf;
use std::thread;

const MAX_MUTATIONS: usize = 15;

const MUT_MIN_BYTE_RANGE: usize = 8;
const MUT_MAX_BYTE_RANGE: usize = 64;

const SAMPLE_DIRECTORY: &str = "./fonts";
const PAYLOAD_ENCODER: bool = true;
const TCP_SERVER: bool = true;
const TCP_SERVER_DEFAULT_PORT: usize = 8080;
const TCP_SERVER_RANDOM_PORT: bool = false;
const TEMPLATE_FILE: bool = true;
const TEMPLATE_FILE_PATH: &str = "./templates/template.html";

static mut TEMPLATE_FILE_CONTENT: Option<String> = None;

fn secure_random(range_max: usize) -> usize {
  let mut rng = thread_rng();
  rng.gen_range(0..=range_max)
}

fn mut_flip_bit(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let i = rng.gen_range(0..buffer.len());
  let bit = rng.gen_range(0..8);
  buffer[i] ^= 1 << bit;
}

fn mut_flip_byte(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let i = rng.gen_range(0..buffer.len());
  let j = rng.gen_range(0..buffer.len());
  let temp = buffer[i];
  buffer[i] = buffer[j];
  buffer[j] = temp;
}

fn mut_replace_byte(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let i = rng.gen_range(0..buffer.len());
  buffer[i] = rng.gen();
}

fn mut_remove_byte(buffer: &mut Vec<u8>) {
  let mut rng = rand::thread_rng();
  let i = rng.gen_range(0..buffer.len());
  buffer.remove(i);
}

fn mut_add_byte(buffer: &mut Vec<u8>) {
  let mut rng = rand::thread_rng();
  let i = rng.gen_range(0..buffer.len());
  let b = rng.gen();
  buffer.insert(i, b);
}

fn mut_repeat_byte(buffer: &mut Vec<u8>) {
  let mut rng = rand::thread_rng();
  let i = rng.gen_range(0..buffer.len());
  let repeat_times = rng.gen_range(1..11);
  let b = buffer[i];
  for _ in 0..repeat_times {
  buffer.insert(i, b);
  }
}

fn mut_reg_flip_bits(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let region_size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len());
  let start = rng.gen_range(0..=buffer.len() - region_size).max(MUT_MIN_BYTE_RANGE);
  let end = start + region_size.min(MUT_MAX_BYTE_RANGE);

  for i in start..end {
    let bit = rng.gen_range(0..8);
    buffer[i] ^= 1 << bit;
  }
}

fn mut_reg_flip_bytes(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let region_size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len());
  let start = rng.gen_range(0..=buffer.len() - region_size).max(MUT_MIN_BYTE_RANGE);
  let end = start + region_size.min(MUT_MAX_BYTE_RANGE);

  for _ in start..end {
    let i = rng.gen_range(start..end);
    let j = rng.gen_range(start..end);
    let temp = buffer[i];
    buffer[i] = buffer[j];
    buffer[j] = temp;
  }
}

fn mut_reg_replace_bytes(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let region_size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len());
  let start = rng.gen_range(0..=buffer.len() - region_size).max(MUT_MIN_BYTE_RANGE);
  let end = start + region_size.min(MUT_MAX_BYTE_RANGE);

  for i in start..end {
    buffer[i] = rng.gen();
  }
}

fn mut_reg_remove_bytes(buffer: &mut Vec<u8>) {
  let mut rng = rand::thread_rng();
  let region_size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len());
  let start = rng.gen_range(0..=buffer.len() - region_size).max(MUT_MIN_BYTE_RANGE);
  let end = start + region_size.min(MUT_MAX_BYTE_RANGE);

  let i = rng.gen_range(start..end);
  buffer.remove(i);
}

fn mut_reg_add_bytes(buffer: &mut Vec<u8>) {
  let mut rng = rand::thread_rng();
  let region_size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len());
  let start = rng.gen_range(0..=buffer.len() - region_size).max(MUT_MIN_BYTE_RANGE);
  let end = start + region_size.min(MUT_MAX_BYTE_RANGE);

  let i = rng.gen_range(start..end);
  let b = rng.gen();
  buffer.insert(i, b);
}

fn mut_reg_repeat_bytes(buffer: &mut Vec<u8>) {
  let mut rng = rand::thread_rng();
  let region_size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len());
  let start = rng.gen_range(0..=buffer.len() - region_size);
  let end = start + region_size;

  let i = rng.gen_range(start..end);
  let repeat_times = rng.gen_range(1..11);
  let b = buffer[i];
  for _ in 0..repeat_times {
    buffer.insert(i, b);
  }
}

fn mut_reg_block_swap_bytes(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len() - 1) & !1;
  let start = rng.gen_range(0..=buffer.len() - size - 1) & !1;

  for i in (start..start + size).step_by(2) {
    buffer.swap(i, i + 1);
  }
}

fn mut_reg_set_bytes(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len() - 1) & !1;
  let start = rng.gen_range(0..=buffer.len() - size - 1) & !1;
  let byte = rng.gen();

  for i in start..start + size {
    buffer[i] = byte;
  }
}

fn mut_reg_shift_bits(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len() - 1) & !1;
  let start = rng.gen_range(0..=buffer.len() - size - 1) & !1;

  for i in start..start + size {
    buffer[i] ^= 0xff;
  }
}

fn mut_reg_overwrite_with_offset(buffer: &mut [u8]) {
  let mut rng = rand::thread_rng();
  let size = rng.gen_range(MUT_MIN_BYTE_RANGE..=MUT_MAX_BYTE_RANGE).min(buffer.len() - 1);
  let src_start = rng.gen_range(0..=buffer.len() - size - 1);
  let dst_start = rng.gen_range(0..=buffer.len() - size - 1);

  for i in 0..size {
    let src_index = (src_start + i) % buffer.len();
    let dst_index = (dst_start + i) % buffer.len();
    buffer[dst_index] = buffer[src_index];
  }
}

fn mut_reg_bit_shift(buffer: &mut [u8]) {
  let shift = rand::thread_rng().gen_range(1..8);
  for i in MUT_MIN_BYTE_RANGE..MUT_MAX_BYTE_RANGE {
    buffer[i] = buffer[i].rotate_left(shift);
  }
}

fn get_random_file(path: &str) -> Option<String> {
  let mut file_paths: Vec<PathBuf> = fs::read_dir(path)
  .ok()?
  .filter_map(|entry| entry.ok().map(|e| e.path()))
  .filter(|path| path.is_file())
  .collect();

  if file_paths.is_empty() {
    return None;
  }

  file_paths.shuffle(&mut thread_rng());

  Some(file_paths[0].to_string_lossy().into_owned())
}

fn handle_read(mut stream: &TcpStream) -> std::io::Result<()> {
  let mut buf = [0u8; 4096];
  let mut has_error = false;

  match stream.read(&mut buf) {
    Ok(_) => {
      let req_str = String::from_utf8_lossy(&buf);
      if req_str.contains("favicon.ico") {
        has_error = true;
      }
    }
    Err(e) => return Err(e),
  }

  if has_error {
    return Err(std::io::Error::new(std::io::ErrorKind::Other, "favicon.ico requested"));
  }

  Ok(())
}

fn fuzz_file(random_file_path: Option<String>) -> String {
  if let Some(path) = random_file_path {
    let mut buffer = std::fs::read(path.clone()).unwrap();

    let rand_mutations = secure_random(MAX_MUTATIONS) + 1;

    println!("Mutation count: {}", rand_mutations);
    println!("File: {}", path);

    for _ in 0..rand_mutations {
      match secure_random(17) {
      0 => mut_flip_byte(&mut buffer),
      1 => mut_replace_byte(&mut buffer),
      2 => mut_remove_byte(&mut buffer),
      3 => mut_add_byte(&mut buffer),
      4 => mut_repeat_byte(&mut buffer),
      5 => mut_flip_bit(&mut buffer),
      6 => mut_reg_flip_bits(&mut buffer),
      7 => mut_reg_flip_bytes(&mut buffer),
      8 => mut_reg_replace_bytes(&mut buffer),
      9 => mut_reg_remove_bytes(&mut buffer),
      10 => mut_reg_add_bytes(&mut buffer),
      11 => mut_reg_repeat_bytes(&mut buffer),
      12 => mut_reg_block_swap_bytes(&mut buffer),
      13 => mut_reg_set_bytes(&mut buffer),
      14 => mut_reg_shift_bits(&mut buffer),
      15 => mut_reg_overwrite_with_offset(&mut buffer),
      16 => mut_reg_bit_shift(&mut buffer),
      _ => (),
      }
    }

    let mut payload = match PAYLOAD_ENCODER {
      true => general_purpose::STANDARD.encode(&buffer),
      false => String::from_utf8_lossy(&buffer).to_string(),
    };

    if TEMPLATE_FILE {
      unsafe {
        payload = match &TEMPLATE_FILE_CONTENT {
          Some(template) => template.replace("{{payload}}", &payload),
          None => payload
        }
      }
    }

    return payload;
  }

  panic!("No files found in directory");
}

fn handle_write(mut stream: TcpStream) {
  match get_random_file(SAMPLE_DIRECTORY) {
    Some(random_file_path) => {
      let response = fuzz_file(Some(random_file_path));
      let mut response_file = File::create("./templates/response.html").expect("Failed to create file");
      response_file
      .write_all(response.as_bytes())
      .expect("Failed to create file");
      match stream.write(response.as_bytes()) {
        Ok(_) => println!("Response sent"),
        Err(e) => println!("Failed sending response: {}", e),
      }
    }
    None => println!("No files found in directory"),
  }
}

fn handle_client(stream: TcpStream) {
  if let Ok(_) = handle_read(&stream) {
    handle_write(stream);
  }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {

  unsafe {
    if TEMPLATE_FILE {
      TEMPLATE_FILE_CONTENT = Some(fs::read_to_string(TEMPLATE_FILE_PATH).expect("Failed to read template file"));
    }
  }

  if TCP_SERVER {
    let port; // secure_random(8000) + 1000;

    if TCP_SERVER_RANDOM_PORT {
      port = secure_random(8000) + 1000
    } else {
      port = TCP_SERVER_DEFAULT_PORT
    }

    let bind = format!("0.0.0.0:{}", port);
    let listener = TcpListener::bind(&bind).unwrap();

    println!("Server running on: 0.0.0.0:{}\nServer address copied to clipboard", port);

    let mut ctx: ClipboardContext = ClipboardProvider::new().unwrap();
    ctx.set_contents(bind).unwrap();

    for stream in listener.incoming() {
      match stream {
        Ok(stream) => {
          thread::spawn(|| handle_client(stream));
        }
        Err(e) => {
          println!("Unable to connect: {}", e);
        }
      }
    }
  } else {
      match get_random_file(SAMPLE_DIRECTORY) {
        Some(random_file_path) => {
          let response = fuzz_file(Some(random_file_path));
          let mut response_file = File::create("./templates/response.html").expect("Failed to create file");
          response_file
          .write_all(response.as_bytes())
          .expect("Failed to create file");
        }
        None => println!("No files found in directory"),
      }
  }

  return Ok(());
}