(* NOTE: unsure what happens when file is larger than a max int32
 but if a file of this size exists then that would be very weird
 unless a block device is used
 In short, code could still improve by handling that condition*)

let usage_msg = "bininject -i inbut.bin -t target.bin -s offset_num";;

let binfile = ref "";;
let targetfile = ref "";;
let offset = ref "";; (* offset of target to inject into*)
let help = ref false;;
let findhole = ref false;;
let targetbyte = ref "0";;

let brief_help_str = "Basic usage: " ^ usage_msg ^ "\nSee -h or -help for more information";; 

let print_brief_help () : unit =
	print_endline brief_help_str; ();;

let params_other = ref [];;

let args_rest filename = 
	params_other := filename :: !params_other;;

let speclist =
[
	("-i", Arg.Set_string binfile, "Binary input file to inject");
	("-t", Arg.Set_string targetfile, "Target binary to inject into");
	("-s", Arg.Set_string offset, "Offset to inject at. Either decimal or prefix with 0x for hex");
	("-h", Arg.Set help, "Display this help message");
	("-0", Arg.Set findhole, "Find a hole in the binary and insert there (-s is ignored)");
	("-b", Arg.Set_string targetbyte, "Look for this byte instead of 0. Accepts 0x prefix for hex (Valid only for -0 mode. Ignored for -s)");
];;


let open_source_file sf =
	let ic = In_channel.open_bin sf in
	let b = Bytes.of_string (In_channel.input_all ic) in
	In_channel.close ic;
	b;;

let write_to_target tf buff (offset:int option) =
	Out_channel.with_open_gen [Out_channel.Open_wronly; Out_channel.Open_creat] 0o644 tf (fun oc ->
        Out_channel.set_binary_mode oc true;
		Out_channel.seek oc (Int64.of_int 0);
		match offset with
			None -> print_endline "Couldn't find valid offset for injection, EOF reached. Exiting."
			| Some off -> begin 
				Out_channel.seek oc (Int64.of_int off);
				Out_channel.output_bytes oc buff
			end
	);;

let hex_to_int str =
	Scanf.sscanf str "0x%x" (fun x -> x);;

let dec_to_int str =
	Scanf.sscanf str "%d" (fun x -> x);;


let parse_num str =
	if String.length str > 2 && (String.starts_with ~prefix:"0x" str) then
		hex_to_int str
	else
		dec_to_int str;;

let rec r_find_first_zero_offset seq_length target_byte file_chan byte_cnt_acc =
	match In_channel.input_byte file_chan with
		None -> None
		| Some byte -> begin
			if byte = target_byte then
				if byte_cnt_acc >= seq_length then Some (Int64.to_int (In_channel.pos file_chan) - byte_cnt_acc - 1)
				(*Found byte, increment count and recurse*)
				else (r_find_first_zero_offset [@tailcall]) seq_length target_byte file_chan (1 + byte_cnt_acc)
			(*Byte not found, reset count and try again*)
			else (r_find_first_zero_offset [@tailcall]) seq_length target_byte file_chan 0
		end


let input_file_get_size infile =
	In_channel.with_open_bin infile In_channel.length;;


let main =
	Arg.parse speclist args_rest usage_msg;
	if !help then 
		(Arg.usage speclist usage_msg; exit 1);;
	if (String.length !binfile = 0 || String.length !targetfile = 0) then 
		(print_brief_help (); exit 1);;
	if !findhole then
		write_to_target 
			!targetfile
			(open_source_file !binfile) 
			(r_find_first_zero_offset 
				(Int64.to_int (input_file_get_size !binfile)) 
				(parse_num !targetbyte)
				(In_channel.open_bin !targetfile)
				0)
	else
		if !offset = "" then begin
				print_brief_help ();
				print_endline "-s expected";
				exit 1;
			end
		else
			write_to_target !targetfile (open_source_file !binfile) (Some (parse_num !offset));;


let () = main;;

.