Example of a Primus component that sets register values before running a function.

primus_set_regs.ml

(*
MAKEFILE:

plugin = primus_set_regs

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 register 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 register values."

    let assignments = param (list ~sep:',' string) "assignments"
      ~doc:"Comma-separated list of REG=INT assignments, e.g., 'R9=42,R10=23'."

end

module Assignment = struct

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

  exception InvalidAssignment of string

  let invalid s =
    Core_kernel.raise 
      (InvalidAssignment (Printf.sprintf "Error. Invalid assignment: %s. Must be REG=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 k, Some v -> create k (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 RegAssignments = sig
  val subroutine : string
  val assignments : Assignment.t list
end

module Component (Log : Printer) (RegAssgn : RegAssignments) (Machine : Primus.Machine.S) = struct
  module Env = Primus.Env.Make(Machine)
  module Value = Primus.Value.Make(Machine)
  open Machine.Syntax

  let msg = Printf.sprintf "[SetRg] -- %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_var ?width:(width=64) name = Var.create name (Type.imm width)
  let make_word ?width:(width=64) i = Word.of_int i ~width:width

  let set_var k v =
    let var = make_var k in
    let w = make_word v in
    let w_str = string_of_word w in

    Log.report (msg (Printf.sprintf "... Setting [%s] to [%s]" k w_str));

    Value.of_word w >>= fun w' ->
    Env.set var w' >>= fun _ ->
    Machine.return ()

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

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

        let actions = List.map RegAssgn.assignments
          ~f:(fun a -> set_var (Assignment.reg 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 : RegAssignments = 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