(**************************************************************************)
(*  This file is part of BINSEC.                                          *)
(*                                                                        *)
(*  Copyright (C) 2016-2019                                               *)
(*    CEA (Commissariat à l'énergie atomique et aux énergies              *)
(*         alternatives)                                                  *)
(*                                                                        *)
(*  you can redistribute it and/or modify it under the terms of the GNU   *)
(*  Lesser General Public License as published by the Free Software       *)
(*  Foundation, version 2.1.                                              *)
(*                                                                        *)
(*  It 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.  See the         *)
(*  GNU Lesser General Public License for more details.                   *)
(*                                                                        *)
(*  See the GNU Lesser General Public License version 2.1                 *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).            *)
(*                                                                        *)
(**************************************************************************)

open Ida_options
open Graph
open Dot_ast
open Format

module IF = Ida_cfg.Function
module IU = Ida_utils

module Node = struct
  module T = struct
    (* based on section containing funcion *)
    type t =
      | Entrypoint
      | Text
      | Plt

    (* convert type from fillcolor value *)
    let of_string = function
      | "#000000" -> Text
      | "#ff00ff" -> Plt
      | _ -> Entrypoint (* #008000 *)
    ;;

    let to_string = function
      | Entrypoint -> "EP"
      | Text -> "TEXT"
      | Plt -> "PLT"
    ;;

    let pp ppf = function
      | Entrypoint -> fprintf ppf "E"
      | Text -> fprintf ppf "T"
      | Plt -> fprintf ppf "P"
    ;;
  end

  type t = {
    nid : Dot_ast.id;
    func : IF.t;
    typ : T.t;
  }

  let nid t = t.nid ;;
  let func t = t.func ;;
  let typ t = t.typ ;;

  let create ?(nid=(String "")) func typ =
    { nid = nid; func = func; typ = typ; }
  ;;

  let pp ppf n =
    Format.fprintf ppf "(%a):%a:%a"
      IU.Dot.pp_id n.nid IF.pp n.func T.pp n.typ
  ;;

  let pp_short ppf n =
    Format.fprintf ppf "[%a]%a"
      T.pp n.typ IF.pp n.func
  ;;

  let pp_list ppf nodes =
    Format.fprintf ppf "%a"
      (Print_utils.pp_list ~sep:", " pp_short) nodes
  ;;

  let equal n1 n2 =
    n1.nid == n2.nid &&
    IF.same n1.func n2.func &&
    n1.typ == n2.typ
  ;;
end

module Edge = struct
  type t = {
    src : Node.t;
    dst : Node.t;
  }

  let src t = t.src ;;
  let dst t = t.dst ;;

  let create s d = { src = s; dst = d; } ;;

  let pp ppf t =
    Format.fprintf ppf "%a -> %a"
      Node.pp t.src Node.pp t.dst
  ;;

  let pp_list ppf edges =
    Format.fprintf ppf "%a"
      (Print_utils.pp_list ~sep:", " pp) edges
  ;;
end

module A = struct
  type t = Node.t
  let compare = compare
  let equal = (=)
  let hash = Hashtbl.hash
end

module G = Cfg.Make (A)(A)(A)
include G

(* Parse a call graph generated by IDA Pro *)
module Parse = struct
  let get_label lb =
    let len = String.length lb - 2 in
    String.sub lb 0 len
  ;;

  let get_attrs nid attrs =
    let loop acc = function
      | [] -> acc
      | _ :: (_, Some (String fillcolor)) ::
        _ :: (_, Some (String label)) :: _ ->
        let label' = get_label label in
        Logger.debug "nid: %a, label: %s, fillcolor: %s"
          IU.Dot.pp_id nid label' fillcolor;
        let fname = IF.name label' in
        let typ = Node.T.of_string fillcolor in
        let node = Node.create ~nid fname typ in
        node :: acc
      | _ :: (_, Some (String fillcolor)) ::
        (_, Some (String label)) :: _ ->
        let label' = get_label label in
        Logger.debug "nid: %a, label: %s, fillcolor: %s"
          IU.Dot.pp_id nid label' fillcolor;
        let fname = IF.name label' in
        let typ = Node.T.of_string fillcolor in
        let node = Node.create ~nid fname typ in
        node :: acc
      | _ -> acc
    in loop [] (List.flatten attrs)
  ;;

  let get_node id nodes =
    try
      List.find (fun n -> Node.nid n = id) nodes
    with Not_found ->
      Logger.error "Can't find node given id: %a" IU.Dot.pp_id id;
      raise (Failure "Can't find node given id")
  ;;

  let build_cg ~cg_file =
    Logger.debug "Parsing callgraph dot file %s" cg_file;
    if not @@ Sys.file_exists cg_file then
      Logger.fatal "Could not find file %s" cg_file;
    let dot_p = Dot.parse_dot_ast cg_file in

    (* Extract nodes in callgraph *)
    let snodes = List.fold_left (fun acc stmt ->
        match stmt with
        | Node_stmt ((nid,_), attrs) ->
          let node_attrs = get_attrs nid attrs in
          node_attrs @ acc
        | _ -> acc
      ) [] dot_p.stmts in

    (* Extract edges in callgraph *)
    let sedges = List.fold_left (fun acc stmt ->
        match stmt with
        | Edge_stmt (node, nodes, _attr) -> begin
            match node with
            | NodeId (node_id,_) -> begin
                try
                  let node = get_node node_id snodes in
                  begin
                    match List.hd nodes with
                    | NodeId (succ_id,_) -> begin
                        try
                          let succ_node = get_node succ_id snodes in
                          let edge = Edge.create node succ_node in
                          Logger.debug "%a" Edge.pp edge;
                          edge :: acc
                        with Not_found -> acc
                      end
                    | NodeSub _ -> acc
                  end
                with Not_found -> acc
              end
            | NodeSub _ -> acc
          end
        | _-> acc
      ) [] dot_p.stmts in
    let cg = create 1000 in
    List.iter (fun e ->
        add_edge_a cg (Edge.src e) (Edge.dst e)) sedges;
    cg
  ;;
end
