Example of a Primus plugin that sets memory values before executing a subroutine.

primus_set_mem.ml

(*
MAKEFILE

plugin = primus_set_mem

all: clean build install

build:
	bapbuild -pkg bap-primus -tags 'warn(A)' $(plugin).plugin

install:
	bapbundle install $(plugin).plugin

clean:
	bapbundle remove $(plugin).plugin
	bapbuild -clean
	rm -rf $(plugin).plugin

*)

open Core_kernel
open Bap.Std
open Bap_primus.Std
include Self()

module Param = struct
  open Config;;
  manpage [
    `S "DESCRIPTION";
    `P
      "This Primus component sets memory values before
       executing a function.";
    ]

    let run = flag "run"
      ~doc:"When this flag is present, the component will be executed."

    let subroutine = param (string) "subroutine"
      ~doc:"A subroutine to execute with specified memory values."

    let assignments = param (list ~sep:',' string) "assignments"
      ~doc:"Comma-separated list of ADDR=INT assignments, e.g., '1=467,2=52'."

end

module Assignment = struct

  type t = {addr: int; value: int}
  let addr t = t.addr
  let value t = t.value
  let create addr value = {addr; value}

  exception InvalidAssignment of string

  let invalid s =
    Core_kernel.raise 
      (InvalidAssignment (Printf.sprintf "Error. Invalid assignment: %s. Must be ADDR=INT.\n%!" s))

  let key_val s =
    let parts = String.split s ~on:'=' in
    match List.length parts with
    | 2 -> 
      begin
        match (List.nth parts 0, List.nth parts 1) with
        | Some a, Some v -> create (Int.of_string a) (Int.of_string v)
        | _ -> invalid s     
      end
    | _ -> invalid s

  let rec parse l acc = 
    match l with
    | [] -> acc
    | hd :: tl -> parse tl (List.append acc [key_val hd])

end

module type Printer = sig
  val report : ?newline:string -> string -> unit
end

module type MemAssignments = sig
  val subroutine : string
  val assignments : Assignment.t list
end

module Component (Log : Printer) (MemAssgn : MemAssignments) (Machine : Primus.Machine.S) = struct
  module Memory = Primus.Memory.Make(Machine)
  module Value = Primus.Value.Make(Machine)
  open Machine.Syntax

  let msg = Printf.sprintf "[SetMm] -- %s"

  let word_of_val v = Value.to_word v
  let string_of_word w = Format.asprintf "%a" Word.pp_bin w
  let string_of_val v = string_of_word (word_of_val v)

  let make_addr ?width:(width=64) a = Word.of_int a ~width:width
  let make_word i = Word.of_int i ~width:8

  let set_mem a v =
    let addr = make_addr a in
    let w = make_word v in
    let a_str = string_of_word addr in
    let w_str = string_of_word w in

    Log.report (msg (Printf.sprintf "... At addr [%s] writing [%s]" a_str w_str));

    Value.of_word w >>= fun w' ->
    Memory.set addr w' >>= fun _ ->
    Machine.return ()

  let enter_sub s =
    let name = Sub.name s in

    match String.(=) MemAssgn.subroutine name with
    | false -> Machine.return ()
    | true ->
      begin
        Log.report (msg (Printf.sprintf "::: Setting memory before executing [%s]" name));

        let actions = List.map MemAssgn.assignments
          ~f:(fun a -> set_mem (Assignment.addr a) (Assignment.value a))
          in

        Machine.sequence actions
      end

  let init () = Machine.sequence [
    Primus.Interpreter.enter_sub >>> enter_sub;
    ]

end

let main {Config.get=(!)} =

  let subroutine = !Param.subroutine in
  let assignments = Assignment.parse !Param.assignments [] in

  match !Param.run with
  | true ->
    begin
      let module P : Printer = struct
        let report ?newline:(nl="\n") msg = Format.printf "%s%s%!" msg nl
      end in

      let module A : MemAssignments = struct
        let subroutine = subroutine
        let assignments = assignments
      end in

      let module PluginComponent = Component(P)(A) in
      Primus.Machine.add_component (module PluginComponent)
    end
  | false -> ()

let () = Config.when_ready main