A simple binary rewriter using BAP

rewriter.ml

open Bap.Std
open Core_kernel
open Bap_main

module Unix = UnixLabels

type chunk = {
  offset : int;
  data : Bigstring.t
}

type patch = {
  chunks : chunk list;
  suffix : Bigstring.t;
}

let mapfile fd size =
  Bigarray.array1_of_genarray @@
  Mmap.V1.map_file fd
    Bigarray.char Bigarray.c_layout true [|size |]

let apply_patch ~output src {suffix; chunks} =
  let fd = Unix.openfile output
      ~mode:Unix.[O_RDWR; O_CREAT]
      ~perm:0o600 in
  (* the size of the new file with suffix included *)
  let size = Bigstring.length src + Bigstring.length suffix in
  (* dst holds the uninitialized contents of the output file  *)
  let dst = mapfile fd size in
  (* first blit the input data *)
  Bigstring.blito ~src ~dst ();
  (* next, append the suffix *)
  Bigstring.blito ~src:suffix ~dst ~dst_pos:(Bigstring.length src) ();
  (* now, for each chunk blit (copy) it to the output *)
  List.iter chunks ~f:(fun {offset; data} ->
      Bigstring.blito ~src:data ~dst ~dst_pos:offset ());
  (* Don't forget to close it *)
  Unix.close fd

let offset mem =
  Bigsubstring.pos (Memory.to_buffer mem)

let find_section image name =
  Image.memory image |>
  Memmap.to_sequence |>
  Seq.find_map ~f:(fun (data,annot) ->
      match Value.get Image.section annot with
      | None -> None
      | Some sec -> Option.some_if (String.equal sec name) data)

let list_of_mem (mem :Memory.t) : (int32 list) =
  List.rev @@
  Memory.fold mem ~word_size:`r32 ~init:[] ~f:(fun w ws ->
      Word.to_int32_exn w :: ws)

let bigstring_of_list ints =
  let size = List.length ints * 4 in
  let dst = Bigstring.create size in
  List.iteri ints ~f:(fun i x ->
      Bigstring.set_int32_t_le dst ~pos:(i*4) x);
  dst

let rewrite (trans : int32 list -> int32 list * int32 list) image output =
  match find_section image ".text" with
  | None -> failwith "there is no .text section in the binary"
  | Some sec ->
    let chunk,suffix = trans (list_of_mem sec) in
    apply_patch ~output (Image.data image) {
      chunks = [{
          offset = offset sec;
          data = bigstring_of_list chunk
        }];
      suffix = bigstring_of_list suffix;
    }

(* our stub transformer, for test *)
let break_binary = function
  | [] -> [],[0xC0FFEEl]
  | x :: xs -> 0xDEADBEEFl :: xs, [0xC0FFEEl]

let output = Extension.Command.argument
    ~doc:"The output file"
    Extension.Type.("FILE" %: string =? "a.out")


let input = Extension.Command.argument
    ~doc:"The output file"
    Extension.Type.("FILE" %: string =? "a.in")

type Extension.Error.t += Fail

let () = Extension.Command.(begin
    declare "rewrite" (args $input $output)
  end) @@ fun input output _ctxt ->
  match Image.create ~backend:"llvm" input with
  | Ok (image,_) ->
    rewrite break_binary image output;
    Ok ()
  | Error err ->
    Format.eprintf "Failed to load the binary: %a@\n"
      Error.pp err;
    Error Fail

Building

bapbuild rewriter.plugin && bapbundle install rewriter.plugin

Using

bap rewrite /bin/true true

Extending

The bogus rewrite binary just ads 0xDEADBEEF to the beginning of the .text section and 0xC0FFEE to the end of the file. Feel free to change it to the function you like. Ideally, choose a better representation than the int32 list. See gitter for the original discussion.