(** ocaml-exploit

    @author Sang Kil Cha <sangkil.cha\@gmail.com>

*)
(*
    Copyright (c) 2014, Sang Kil Cha
    All rights reserved.
    This software is free software; you can redistribute it and/or
    modify it under the terms of the GNU Library General Public
    License version 2, with the special exception on linking
    described in file LICENSE.

    This software is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
*)

open Bitstring

type sock =
  | TCP of Unix.file_descr
  | UDP of Unix.file_descr * Unix.sockaddr

let (|>) a b = b a

let create_socket domain socktype proto addr port =
  let sock = Unix.socket domain socktype proto in
  let hostinfo = Unix.gethostbyname addr in
  let host = hostinfo.Unix.h_addr_list.(0) in
  sock, Unix.ADDR_INET (host, port)

let connect_tcp addr port =
  let connect domain socktype addr port =
    let sock, addr  = create_socket domain socktype 0 addr port in
    Unix.connect sock addr; sock
  in
  TCP (connect Unix.PF_INET Unix.SOCK_STREAM addr port)

let connect_udp addr port =
  let proto = (Unix.getprotobyname "udp").Unix.p_proto in
  let sock, addr =
    create_socket Unix.PF_INET Unix.SOCK_DGRAM proto addr port
  in
  let sock =
    Unix.socket
      Unix.PF_INET Unix.SOCK_DGRAM
      (Unix.getprotobyname "udp").Unix.p_proto
  in
  UDP (sock, addr)

let disconnect = function
  | TCP sock -> Unix.close sock
  | UDP (sock, _) -> Unix.close sock

let send_msg sock msg =
  let len = String.length msg in
  match sock with
    | TCP sock ->
        let nbytes = Unix.send sock msg 0 len [] in
        assert( nbytes = len )
    | UDP (sock, addr) ->
        let nbytes = Unix.sendto sock msg 0 len [] addr in
        assert( nbytes = len )

let recv_msg sock maxlen =
  let str = String.make maxlen '\x00' in
  match sock with
    | TCP sock ->
        let nbytes = Unix.recv sock str 0 maxlen [] in
        String.sub str 0 nbytes
    | UDP (sock, addr) ->
        let nbytes, _ = Unix.recvfrom sock str 0 maxlen [] in
        String.sub str 0 nbytes

let wait () = ignore (read_line ())

let ( * ) str num =
  if num <= 0 then failwith "wrong num";
  let rec loop acc num =
    if num = 0 then acc
    else loop (str::acc) (num-1)
  in
  let str_list = loop [] num in
  String.concat "" str_list

let int32_to_hexstring i =
  Printf.sprintf "%08lx" i

let hexstring_to_string hexstr =
  let hexchar prev curr =
    int_of_string (Printf.sprintf "0x%c%c" prev curr) |> char_of_int
  in
  let len = String.length hexstr in
  assert (len mod 2 = 0);
  let str = String.make (len / 2) '\x00' in
  let _ =
    BatString.fold_left (fun (pos, prev) ch ->
      match prev with
        | None -> pos, Some ch
        | Some prev -> str.[pos] <- (hexchar prev ch); pos+1, None
    ) (0, None) hexstr
  in
  str

let print str = print_endline str; flush stdout

let explode s =
  let rec exp i l =
    if i < 0 then l else exp (i - 1) (s.[i] :: l)
  in
  exp (String.length s - 1) []

let string_to_hexstring little str =
  let accumulate acc chunk =
    Printf.sprintf "%s%08lx" acc chunk
  in
  let rec little_loop acc bs =
    bitmatch bs with
      | {chunk : 32 : littleendian; rest  : -1 : bitstring }
        when bitstring_length rest = 0 -> accumulate acc chunk
      | {chunk : 32 : littleendian;
         chunk2 : 24 : littleendian; rest : -1 : bitstring }
        when bitstring_length rest = 0 ->
          let acc = accumulate acc chunk in
          Printf.sprintf "%s%06x\n" acc chunk2
      | {chunk : 32 : littleendian; rest  : -1 : bitstring } ->
          let acc = accumulate acc chunk in
          little_loop acc rest
      | {chunk : 8 : littleendian; rest : -1 : bitstring }
        when bitstring_length rest = 0 ->
          Printf.sprintf "%s%02x\n" acc chunk
      | {chunk : 16 : littleendian; rest : -1 : bitstring }
        when bitstring_length rest = 0 ->
          Printf.sprintf "%s%04x\n" acc chunk
  in
  let rec big_loop acc bs =
    bitmatch bs with
      | {chunk : 32 : bigendian; rest : -1 : bitstring }
        when bitstring_length rest = 0 -> accumulate acc chunk
      | {chunk : 32 : bigendian;
         chunk2 : 24 : bigendian; rest : -1 : bitstring }
        when bitstring_length rest = 0 ->
          let acc = accumulate acc chunk in
          Printf.sprintf "%s%06x\n" acc chunk2
      | {chunk : 32 : bigendian; rest : -1 : bitstring } ->
          let acc = accumulate acc chunk in
          big_loop acc rest
      | {chunk : 8 : bigendian; rest : -1 : bitstring }
        when bitstring_length rest = 0 ->
          Printf.sprintf "%s%02x\n" acc chunk
      | {chunk : 16 : bigendian; rest : -1 : bitstring }
        when bitstring_length rest = 0 ->
          Printf.sprintf "%s%04x\n" acc chunk
  in
  let bs = bitstring_of_string str in
  if little then little_loop "" bs
  else big_loop "" bs

let hexdump str =
  bitstring_of_string str |> hexdump_bitstring stdout

let new_char_list f t =
  let rec loop acc cur =
    if cur < t then loop (cur::acc) (char_of_int ((int_of_char cur)+1))
    else List.rev (cur::acc)
  in
  loop [] f

let pattern_create num =
  let lower_alpha = new_char_list 'a' 'z' in
  let upper_alpha = new_char_list 'A' 'Z' in
  let numeric = new_char_list '0' '9' in
  let get_next l u n =
    match l, u, n with
      | lhd::[] , uhd::[], nhd::[] -> failwith "too large pattern"
      | lhd::ltl, uhd::[], nhd::[] -> ltl, upper_alpha, numeric
      | lhd::ltl, uhd::utl, nhd::[] -> l, utl, numeric
      | lhd::ltl, uhd::utl, nhd::ntl -> l, u, ntl
      | _ -> failwith "wrong pattern"
  in
  let rec create_loop acc (l,u,n) (cnt:int) =
    if cnt >= num then
      String.concat ""
        (List.rev (List.map (fun c -> Printf.sprintf "%c" c) acc))
    else
      let nl, nu, nn = get_next l u n in
      let acc =
        if cnt + 2 < num then (List.hd n)::(List.hd u)::(List.hd l)::acc
        else if cnt + 1 < num then (List.hd u)::(List.hd l)::acc
        else if cnt < num then (List.hd l)::acc
        else acc
      in
      create_loop acc (nl,nu,nn) (cnt+3)
  in
  create_loop [] (lower_alpha,upper_alpha,numeric) 0

let pattern_pos p num =
  let pattern = pattern_create num in
  let regexp = Str.regexp (Str.quote p) in
  (Str.search_forward regexp pattern 0) + 1

