type reg = Int64.t
type id = Int64.t
type addr = Int64.t

exception Invalid_opcode of int
exception Invalid_message
exception Unsupported_message

let regsize = 10 * 8
let opsize = 1
let idsize = 8
let addrsize = 8

type regs = 
    {
      ebx : reg;
      ecx : reg;
      edx : reg;
      esi : reg;
      edi : reg;
      ebp : reg;
      eax : reg;
      eip : reg;
      esp : reg;
      eflags : reg;
    }

type opcode = 
  | RerrorOp   (* 0 *)
  | TlaunchOp
  | RlaunchOp
  | TcontOp
  | RcontOp
  | TstopOp
  | RstopOp
  | TkillOp
  | RkillOp
  | TsetbrkOp
  | RsetbrkOp  (* 10 *)
  | TclrbrkOp
  | RclrbrkOp
  | TreadOp
  | RreadOp
  | AbrkOp
  | AobitOp
  | AackOp
  | TcleaveOp
  | RcleaveOp
  | TwriteOp   (* 20 *)
  | RwriteOp
  | TsetregsOp
  | RsetregsOp
  | TsetsnapOp
  | RsetsnapOp
  | AsnapOp
  | TattachOp
  | RattachOp
  | TtestOp
  | RtestOp    (* 30 *)
  | TstepOp
  | RstepOp

type msg = 
  | Rerror of id * string
  | Tlaunch of string list 
  | Rlaunch of id * regs
  | Tcont of id
  | Rcont of id
  | Tstop of id
  | Rstop of id * regs
  | Tkill of id
  | Rkill of id
  | Tsetbrk of id * addr
  | Rsetbrk of id * addr
  | Tclrbrk of id * addr
  | Rclrbrk of id 
  | Tread of id * addr * Int64.t
  | Rread of id * string
  | Abrk of id * addr * regs
  | Aobit of id
  | Aack of id
  | Tcleave of id
  | Rcleave of id * id
  | Twrite of id * addr * string
  | Rwrite of id
  | Tsetregs of id * regs
  | Rsetregs of id
  | Tsetsnap of id * addr
  | Rsetsnap of id * addr
  | Asnap of id * addr * id
  | Tattach of id
  | Rattach of id
  | Ttest of string
  | Rtest  of Int64.t
  | Tstep of id 
  | Rstep of id

let read_opcode (ic:in_channel) : opcode = 
  let op = input_byte ic in 
    match op with 
      | 0 -> RerrorOp 
      | 1 -> TlaunchOp
      | 2 -> RlaunchOp
      | 3 -> TcontOp
      | 4 -> RcontOp
      | 5 -> TstopOp
      | 6 -> RstopOp
      | 7 -> TkillOp
      | 8 -> RkillOp
      | 9 -> TsetbrkOp
      | 10 -> RsetbrkOp
      | 11 -> TclrbrkOp
      | 12 -> RclrbrkOp
      | 13 -> TreadOp
      | 14 -> RreadOp
      | 15 -> AbrkOp
      | 16 -> AobitOp
      | 17 -> AackOp
      | 18 -> TcleaveOp
      | 19 -> RcleaveOp
      | 20 -> TwriteOp
      | 21 -> RwriteOp
      | 22 -> TsetregsOp
      | 23 -> RsetregsOp
      | 24 -> TsetsnapOp
      | 25 -> RsetsnapOp
      | 26 -> AsnapOp
      | 27 -> TattachOp
      | 28 -> RattachOp
      | 29 -> TtestOp
      | 30 -> RtestOp
      | 31 -> TstepOp
      | 32 -> RstepOp
      | _ -> raise (Invalid_opcode op)

let write_opcode (op:opcode) (oc:out_channel) : unit = 
  match op with 
    | RerrorOp -> output_byte oc 0
    | TlaunchOp -> output_byte oc 1
    | RlaunchOp -> output_byte oc 2
    | TcontOp -> output_byte oc 3
    | RcontOp -> output_byte oc 4
    | TstopOp -> output_byte oc 5
    | RstopOp -> output_byte oc 6
    | TkillOp -> output_byte oc 7
    | RkillOp -> output_byte oc 8
    | TsetbrkOp -> output_byte oc 9
    | RsetbrkOp -> output_byte oc 10
    | TclrbrkOp -> output_byte oc 11
    | RclrbrkOp -> output_byte oc 12
    | TreadOp -> output_byte oc 13
    | RreadOp -> output_byte oc 14
    | AbrkOp -> output_byte oc 15
    | AobitOp -> output_byte oc 16
    | AackOp -> output_byte oc 17
    | TcleaveOp -> output_byte oc 18
    | RcleaveOp -> output_byte oc 19
    | TwriteOp -> output_byte oc 20
    | RwriteOp -> output_byte oc 21
    | TsetregsOp -> output_byte oc 22
    | RsetregsOp -> output_byte oc 23
    | TsetsnapOp -> output_byte oc 24
    | RsetsnapOp -> output_byte oc 25
    | AsnapOp -> output_byte oc 26
    | TattachOp -> output_byte oc 27
    | RattachOp -> output_byte oc 28
    | TtestOp -> output_byte oc 29
    | RtestOp -> output_byte oc 30
    | TstepOp -> output_byte oc 31
    | RstepOp -> output_byte oc 32

let read_u64 (ic:in_channel) : int64 = 
  let u64 = ref Int64.zero in 
    for i = 0 to 7 do 
      u64 := (Int64.logor !u64 
		(Int64.shift_left (Int64.of_int (input_byte ic)) 
		   (8 * i)))
    done;
    !u64

let write_u64 (u64:int64) (oc:out_channel) : unit = 
  for i = 0 to 7 do 
    (output_byte oc (Int64.to_int 
		       (Int64.logand 
			  (Int64.shift_right_logical u64 (8 * i))
			  (Int64.of_int 255))))
  done
      

let read_regs (ic:in_channel) : regs = 
  let ebx = read_u64 ic in 
  let ecx = read_u64 ic in 
  let edx = read_u64 ic in 
  let esi = read_u64 ic in 
  let edi = read_u64 ic in 
  let ebp = read_u64 ic in 
  let eax = read_u64 ic in 
  let eip = read_u64 ic in 
  let esp = read_u64 ic in 
  let eflags = read_u64 ic in 
    { 
      ebx = ebx;
      ecx = ecx;
      edx = edx;
      esi = esi;
      edi = edi;
      ebp = ebp;
      eax = eax;
      eip = eip;
      esp = esp;
      eflags = eflags;
    }
      

let writemsg (m:msg) (oc:out_channel) : unit  = 
  (match m with 
     | Tlaunch slist -> (* Tlaunch[1] narg[8] arg[s] ... *)
	 let slen = List.length slist in 
	 let out_len = 
	   (List.fold_left (fun x y -> x + 1 + String.length y) 
	      (opsize + addrsize) slist) in
	   write_u64 (Int64.of_int out_len) oc;
	   write_opcode TlaunchOp oc;
	   write_u64 (Int64.of_int slen) oc;
	   List.iter (fun x -> output_string oc x; output_byte oc 0) slist

     | Tstop (id) -> (* op[1] id[8] *)
	 let out_len = opsize + idsize in 
	   write_u64 (Int64.of_int out_len) oc;
	   write_opcode TstopOp oc;
	   write_u64 id oc

     | Tsetbrk (pid,addr) -> (* op[1] id[8] addr[8] *)
	 let out_len = opsize + idsize + addrsize in 
	   write_u64 (Int64.of_int out_len) oc;
	   write_opcode TsetbrkOp oc;
	   write_u64 pid oc;
	   write_u64 addr oc
	 
     | Tclrbrk (pid,bpid) -> (* op[1] id[8] bp[8] *)
	 let out_len = opsize + idsize + idsize in
	   write_u64 (Int64.of_int out_len) oc;
	   write_opcode TclrbrkOp oc;
	   write_u64 pid oc;
	   write_u64 bpid oc

     | Tcont (pid) -> (* op[1] id[8] *)
	 let out_len = opsize + idsize in 
	   write_u64 (Int64.of_int out_len) oc;
	   write_opcode TcontOp oc;
	   write_u64 pid oc
    
     | Aack (pid) -> (* op[1] id[8] *)
	 let out_len = opsize + idsize in 
	   write_u64 (Int64.of_int out_len) oc;
	   write_opcode AackOp oc;
	   write_u64 pid oc

     | Ttest s -> (* Ttest[1] cnt[8] bytes[cnt]... *)
	 let strlen = String.length s in 
	 let out_len = (opsize + addrsize + strlen) in 
	   write_u64 (Int64.of_int out_len) oc;
	   write_opcode TtestOp oc;
	   write_u64 (Int64.of_int strlen) oc;
	   if strlen > 0 then 
	     output oc s 0 strlen

     | Tstep (pid) -> (* Tstep[1] id[8] *)
	 let out_len = opsize + idsize in 
	   write_u64 (Int64.of_int out_len) oc;
	   write_opcode TstepOp oc;
	   write_u64 pid oc

     | _ -> raise Unsupported_message
  );
  flush oc

let readmsg (ic:in_channel) : msg = 
  let mlen = Int64.to_int (read_u64 ic) in 
  let op = read_opcode ic in 
  let m = 
    match op with 
      | RlaunchOp -> (* op[1] id[8] reg[sizeof(Reg)] *)
	  if mlen <> (opsize + idsize + regsize) then 
	    raise Invalid_message;
	  let id = read_u64 ic in 
	  let regs = read_regs ic in 
	    Rlaunch(id,regs)

      | RsetbrkOp -> (* op[1] id[8] bp[8] *)
	  if mlen <> (opsize + idsize + idsize) then
	    raise Invalid_message;
	  let pid = read_u64 ic in
	  let bpid = read_u64 ic in 
	    Rsetbrk(pid,bpid)

      | RclrbrkOp -> (* op[1] id[8] *)
	  if mlen <> (opsize + idsize) then
	    raise Invalid_message;
	  let id = read_u64 ic in 
	    Rclrbrk(id)
	  
      | RcontOp -> (* op[1] id[8] *)
	  if mlen <> (opsize + idsize) then
	    raise Invalid_message;
	  let id = read_u64 ic in 
	    Rcont(id)
	  
      | AbrkOp -> (* Abrk[1] id[8] bp[8] reg[sizeof(Reg)] *)
	  if mlen <> (opsize + idsize + idsize + regsize) then
	    raise Invalid_message;
	  let id = read_u64 ic in
	  let bpid = read_u64 ic in 
	  let regs = read_regs ic in 
	    Abrk(id,bpid,regs)

      | AobitOp -> (* op[1] id[8] *)
	  if mlen <> (opsize + idsize) then
	    raise Invalid_message;
	  let id = read_u64 ic in 
	    Aobit(id)

      | RtestOp -> (* Rtest[1] cnt[8] *)
	  if mlen <> (opsize + addrsize)  then
	    raise Invalid_message;
	  let count = read_u64 ic in 
	    Rtest(count)

      | RstepOp -> (* Rstep[1] id[8] *)
	  if mlen <> (opsize + idsize)  then
	    raise Invalid_message;
	  let id = read_u64 ic in 
	    Rstep(id)

      | _ -> raise Unsupported_message
  in 
    m