(* Guestfs Browser.
 * Copyright (C) 2010 Red Hat Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *)

open ExtString
open ExtList
open Unix
open Printf

open Utils
open DeviceSet
open Slave_types

open Filetree_type
open Filetree_markup
open Filetree_ops

module G = Guestfs
module UTF8 = CamomileLibraryDefault.Camomile.UTF8

type t = Filetree_type.t

(* Temporary directory for shared use by any function in this file.
 * It is cleaned up when the program exits.
 *)
let tmpdir = tmpdir ()

let rec create ~packing () =
  let view = GTree.view ~packing () in
  (*view#set_rules_hint true;*)
  (*view#selection#set_mode `MULTIPLE; -- add this later *)

  (* Hash of index numbers -> hdata.  We do this because it's more
   * efficient for the GC compared to storing OCaml objects directly in
   * the rows.
   *)
  let hash = Hashtbl.create 1023 in

  (* The columns stored in each row.  The hidden [index_col] column is
   * an index into the hash table that records everything else about
   * this row (see hdata above).  The other display columns, eg.
   * [name_col] contain Pango markup and thus have to be escaped.
   *)
  let cols = new GTree.column_list in
  (* Hidden: *)
  let index_col = cols#add Gobject.Data.int in
  (* Displayed: *)
  let mode_col = cols#add Gobject.Data.string in
  let name_col = cols#add Gobject.Data.string in
  let size_col = cols#add Gobject.Data.string in
  let date_col = cols#add Gobject.Data.string in

  (* Create the model. *)
  let model = GTree.tree_store cols in

  (* Create the view. *)
  view#set_model (Some (model :> GTree.model));

  let renderer = GTree.cell_renderer_text [], ["markup", mode_col] in
  let mode_view = GTree.view_column ~title:"Permissions" ~renderer () in
  mode_view#set_resizable true;
  ignore (view#append_column mode_view);

  let renderer = GTree.cell_renderer_text [], ["markup", name_col] in
  let name_view = GTree.view_column ~title:"Filename" ~renderer () in
  name_view#set_resizable true;
  name_view#set_sizing `AUTOSIZE;
  ignore (view#append_column name_view);

  let renderer = GTree.cell_renderer_text [`XALIGN 1.], ["markup", size_col] in
  let size_view = GTree.view_column ~title:"Size" ~renderer () in
  size_view#set_resizable true;
  ignore (view#append_column size_view);

  let renderer = GTree.cell_renderer_text [`XALIGN 1.], ["markup", date_col] in
  let date_view = GTree.view_column ~title:"Date" ~renderer () in
  date_view#set_resizable true;
  ignore (view#append_column date_view);

  let t = {
    view = view; model = model; hash = hash;
    index_col = index_col;
    mode_col = mode_col; name_col = name_col; size_col = size_col;
    date_col = date_col;
  } in

  (* Open a context menu when a button is pressed. *)
  ignore (view#event#connect#button_press ~callback:(button_press t));

  t

(* Handle mouse button press on the selected row.  This opens the
 * pop-up context menu.
 * http://scentric.net/tutorial/sec-selections-context-menus.html
 *)
and button_press ({ model = model; view = view } as t) ev =
  let button = GdkEvent.Button.button ev in
  let x = int_of_float (GdkEvent.Button.x ev) in
  let y = int_of_float (GdkEvent.Button.y ev) in
  let time = GdkEvent.Button.time ev in

  (* Right button for opening the context menu. *)
  if button = 3 then (
(*
    (* If no row is selected, select the row under the mouse. *)
    let paths =
      let sel = view#selection in
      if sel#count_selected_rows < 1 then (
        match view#get_path_at_pos ~x ~y with
        | None -> []
        | Some (path, _, _, _) ->
            sel#unselect_all ();
            sel#select_path path;
            [path]
      ) else
        sel#get_selected_rows (* actually returns paths *) in
*)
    (* Select the row under the mouse. *)
    let paths =
      let sel = view#selection in
      match view#get_path_at_pos ~x ~y with
      | None -> []
      | Some (path, _, _, _) ->
          sel#unselect_all ();
          sel#select_path path;
          [path] in

    (* Get the hdata for all the paths.  Filter out rows that it doesn't
     * make sense to select.
     *)
    let paths =
      List.filter_map (
        fun path ->
          let row = model#get_iter path in
          let hdata = get_hdata t row in
          match hdata with
          | { content=(Loading | ErrorMessage _ | Info _) } -> None
          | { content=(Top _ | Directory _ | File _ |
                           TopWinReg _ | RegKey _ | RegValue _ ) } ->
              Some (path, hdata)
      ) paths in

    (* Based on number of selected rows and what is selected, construct
     * the context menu.
     *)
    (match make_context_menu t paths with
     | Some menu -> menu#popup ~button ~time
     | None -> ()
    );

    (* Return true so no other handler will run. *)
    true
  )
  (* We didn't handle this, defer to other handlers. *)
  else false

and make_context_menu t paths =
  let menu = GMenu.menu () in
  let factory = new GMenu.factory menu in

  let rec add_file_items path =
    let item = factory#add_item "View ..." in
    (match Config.opener with
     | Some opener ->
         ignore (item#connect#activate ~callback:(view_file t path opener));
     | None ->
         item#misc#set_sensitive false
    );
    let item = factory#add_item "File information" in
    ignore (item#connect#activate ~callback:(file_information t path));
    let item = factory#add_item "MD5 checksum" in
    ignore (item#connect#activate ~callback:(checksum_file t path "md5"));
    let item = factory#add_item "SHA1 checksum" in
    ignore (item#connect#activate ~callback:(checksum_file t path "sha1"));
    ignore (factory#add_separator ());
    let item = factory#add_item "Download ..." in
    ignore (item#connect#activate ~callback:(download_file t path));

  and add_directory_items path =
    let item = factory#add_item "Directory information" in
    item#misc#set_sensitive false;
    let item = factory#add_item "Calculate disk usage" in
    ignore (item#connect#activate ~callback:(disk_usage t path));
    ignore (factory#add_separator ());
    let item = factory#add_item "Download ..." in
    item#misc#set_sensitive false;
    let item = factory#add_item "Download as .tar ..." in
    ignore (item#connect#activate
              ~callback:(download_dir_tarball t Tar path));
    let item = factory#add_item "Download as .tar.gz ..." in
    ignore (item#connect#activate
              ~callback:(download_dir_tarball t TGZ path));
    let item = factory#add_item "Download as .tar.xz ..." in
    ignore (item#connect#activate
              ~callback:(download_dir_tarball t TXZ path));
    let item = factory#add_item "Download list of filenames ..." in
    ignore (item#connect#activate ~callback:(download_dir_find0 t path));

  and add_top_os_items path =
    let item = factory#add_item "Operating system information" in
    ignore (item#connect#activate ~callback:(display_inspection_data t path));
    ignore (factory#add_separator ());
    add_top_volume_items path

  and add_top_volume_items path =
    let item = factory#add_item "Filesystem used & free" in
    item#misc#set_sensitive false;
    let item = factory#add_item "Block device information" in
    item#misc#set_sensitive false;
    ignore (factory#add_separator ());
    add_directory_items path

  and add_topwinreg_items path =
    let item = factory#add_item "Download hive file ..." in
    item#misc#set_sensitive false;
    ignore (factory#add_separator ());
    add_regkey_items path

  and add_regkey_items path =
    let item = factory#add_item "Download as .reg file ..." in
    (match Config.hivexregedit with
     | Some hivexregedit ->
         ignore (item#connect#activate
                   ~callback:(download_as_reg t path hivexregedit));
     | None ->
         item#misc#set_sensitive false
    )

  and add_regvalue_items path =
    let item = factory#add_item "Copy value to clipboard" in
    ignore (item#connect#activate ~callback:(copy_regvalue t path));

  in

  let has_menu =
    match paths with
    | [] -> false

    (* single selection *)
    | [path, { content=Top (OS os)} ] ->  (* top level operating system *)
        add_top_os_items path; true

    | [path, { content=Top (Volume dev) }] -> (* top level volume *)
        add_top_volume_items path; true

    | [path, { content=Directory _ }] -> (* directory *)
        add_directory_items path; true

    | [path, { content=File _ }] ->      (* file *)
        add_file_items path; true

    | [path, { content=TopWinReg _ }] -> (* top level registry node *)
        add_topwinreg_items path; true

    | [path, { content=RegKey _ }] ->    (* registry node *)
        add_regkey_items path; true

    | [path, { content=RegValue _ }] ->  (* registry key/value pair *)
        add_regvalue_items path; true

    | [_, { content=(Loading|ErrorMessage _|Info _) }] -> false

    | _::_::_ ->
        (* At the moment multiple selection is disabled.  When/if we
         * enable it we should do something intelligent here. XXX
         *)
        false in
  if has_menu then Some menu else None

let clear { model = model; hash = hash } =
  model#clear ();
  Hashtbl.clear hash

let rec add ({ model = model } as t) name data =
  clear t;

  (* Populate the top level of the filetree.  If there are operating
   * systems from inspection, these have their own top level entries
   * followed by only unreferenced filesystems.  If we didn't get
   * anything from inspection, then at the top level we just show
   * filesystems.
   *)
  let other_filesystems =
    DeviceSet.of_list (List.map fst data.insp_all_filesystems) in
  let other_filesystems =
    List.fold_left (fun set { insp_filesystems = fses } ->
                      DeviceSet.subtract set (DeviceSet.of_array fses))
      other_filesystems data.insp_oses in

  (* Add top level operating systems. *)
  List.iter (add_top_level_os t name) data.insp_oses;

  (* Add top level left-over filesystems. *)
  DeviceSet.iter (add_top_level_vol t name) other_filesystems;

  (* If it's Windows and registry files exist, create a node for
   * each file.
   *)
  List.iter (
    fun os ->
      (match os.insp_winreg_SAM with
       | Some filename ->
           add_top_level_winreg t name os "HKEY_LOCAL_MACHINE\\SAM" filename
       | None -> ()
      );
      (match os.insp_winreg_SECURITY with
       | Some filename ->
           add_top_level_winreg t name os "HKEY_LOCAL_MACHINE\\SECURITY"
             filename
       | None -> ()
      );
      (match os.insp_winreg_SOFTWARE with
       | Some filename ->
           add_top_level_winreg t name os "HKEY_LOCAL_MACHINE\\SOFTWARE"
             filename
       | None -> ()
      );
      (match os.insp_winreg_SYSTEM with
       | Some filename ->
           add_top_level_winreg t name os "HKEY_LOCAL_MACHINE\\SYSTEM"
             filename
       | None -> ()
      );
      (match os.insp_winreg_DEFAULT with
       | Some filename ->
           add_top_level_winreg t name os "HKEY_USERS\\.DEFAULT" filename
       | None -> ()
      );
  ) data.insp_oses;

  (* Expand the first top level node. *)
  match model#get_iter_first with
  | None -> ()
  | Some row ->
      t.view#expand_row (model#get_path row)

(* Add a top level operating system node. *)
and add_top_level_os ({ model = model } as t) name os =
  let markup =
    sprintf "<b>%s</b>\n<small>%s</small>\n<small>%s</small>"
      (markup_escape name) (markup_escape os.insp_hostname)
      (markup_escape os.insp_product_name) in

  let row = model#append () in
  make_node t row (Top (OS os)) None;
  model#set ~row ~column:t.name_col markup

(* Add a top level volume (left over filesystem) node. *)
and add_top_level_vol ({ model = model } as t) name dev =
  let markup =
    sprintf "<b>%s</b>\n<small>from %s</small>"
      (markup_escape dev) (markup_escape name) in

  let row = model#append () in
  make_node t row (Top (Volume dev)) None;
  model#set ~row ~column:t.name_col markup

(* Add a top level Windows Registry node. *)
and add_top_level_winreg ({ model = model } as t) name os rootkey
    remotefile =
  let cachefile = tmpdir // string_of_int (unique ()) ^ ".hive" in

  let markup =
    sprintf "<b>%s</b>\n<small>from %s</small>"
      (markup_escape rootkey) (markup_escape name) in

  let row = model#append () in
  make_node t row
    (TopWinReg (OS os, rootkey, remotefile, cachefile)) None;
  model#set ~row ~column:t.name_col markup

(* Generic function to make an openable node to the tree. *)
and make_node ({ model = model } as t) row content hiveh =
  let hdata =
    { state=NodeNotStarted; content=content; visited=false; hiveh=hiveh } in
  store_hdata t row hdata;

  (* Create a placeholder "loading ..." row underneath this node so
   * the user has something to expand.
   *)
  let placeholder = model#append ~parent:row () in
  let hdata = { state=IsLeaf; content=Loading; visited=false; hiveh=None } in
  store_hdata t placeholder hdata;
  model#set ~row:placeholder ~column:t.name_col "<i>Loading ...</i>";
  ignore (t.view#connect#row_expanded ~callback:(expand_row t))

and make_leaf ({ model = model } as t) row content hiveh =
  let hdata = { state=IsLeaf; content=content; visited=false; hiveh=hiveh } in
  store_hdata t row hdata

(* This is called when the user expands a row. *)
and expand_row ({ model = model } as t) row _ =
  match get_hdata t row with
  | { state=NodeNotStarted; content=Top src } as hdata ->
      (* User has opened a top level node that was not previously opened. *)

      (* Mark this row as loading, so we don't try to open it again. *)
      hdata.state <- NodeLoading;

      (* Get a stable path for this row. *)
      let path = model#get_path row in

      Slave.read_directory ~fail:(when_read_directory_fail t path)
        src "/" (when_read_directory t path)

  | { state=NodeNotStarted; content=Directory direntry } as hdata ->
      (* User has opened a filesystem directory not previously opened. *)

      (* Mark this row as loading. *)
      hdata.state <- NodeLoading;

      (* Get a stable path for this row. *)
      let path = model#get_path row in

      let src, pathname = get_pathname t row in

      Slave.read_directory ~fail:(when_read_directory_fail t path)
        src pathname (when_read_directory t path)

  | { state=NodeNotStarted;
      content=TopWinReg (src, rootkey, remotefile, cachefile) } as hdata ->
      (* User has opened a Windows Registry top level node
       * not previously opened.
       *)

      (* Mark this row as loading. *)
      hdata.state <- NodeLoading;

      (* Get a stable path for this row. *)
      let path = model#get_path row in

      (* Since the user has opened this top level registry node for the
       * first time, we now need to download the hive.
       *)
      cache_registry_file ~fail:(when_downloaded_registry_fail t path)
        t path src remotefile cachefile (when_downloaded_registry t path)

  | { state=NodeNotStarted; content=RegKey node } as hdata ->
      (* User has opened a Windows Registry key node not previously opened. *)

      (* Mark this row as loading. *)
      hdata.state <- NodeLoading;

      expand_hive_node t row node

  (* Ignore when a user opens a node which is loading or has been loaded. *)
  | { state=(NodeLoading|IsNode) } -> ()

  (* In some circumstances these can be nodes, eg. if we have added Info
   * nodes below them.  Just ignore them if opened.
   *)
  | { content=(File _ | RegValue _) } | { state=IsLeaf } -> ()

  (* Node should not exist in the tree. *)
  | { state=NodeNotStarted; content=(Loading | ErrorMessage _ | Info _) } ->
      assert false

(* This is the callback when the slave has read the directory for us. *)
and when_read_directory ({ model = model } as t) path entries =
  debug "when_read_directory";

  let row = model#get_iter path in

  (* Sort the entries by lexicographic ordering. *)
  let cmp { dent_name = n1 } { dent_name = n2 } =
    UTF8.compare n1 n2
  in
  let entries = List.sort ~cmp entries in

  (* Add the entries. *)
  List.iter (
    fun direntry ->
      let { dent_name = name; dent_stat = stat; dent_link = link } =
        direntry in
      let row = model#append ~parent:row () in
      if is_directory stat.G.mode then
        make_node t row (Directory direntry) None
      else
        make_leaf t row (File direntry) None;
      model#set ~row ~column:t.name_col (markup_of_name direntry);
      model#set ~row ~column:t.mode_col (markup_of_mode stat.G.mode);
      model#set ~row ~column:t.size_col (markup_of_size stat.G.size);
      model#set ~row ~column:t.date_col (markup_of_date stat.G.mtime);
  ) entries;

  (* Remove the placeholder "Loading" entry.  NB. Must be done AFTER
   * adding the other entries, or else Gtk will unexpand the row.
   *)
  (try
     let row = find_child_node_by_content t row Loading in
     ignore (model#remove row)
   with Invalid_argument _ | Not_found -> ()
  );

  (* The original directory entry has now been loaded, so
   * update its state.
   *)
  let hdata = get_hdata t row in
  hdata.state <- IsNode;
  set_visited t row

(* This is called instead of when_read_directory when the read directory
 * (or mount etc) failed.  Convert the "Loading" entry into the
 * error message.
 *)
and when_read_directory_fail ({ model = model } as t) path exn =
  debug "when_read_directory_fail: %s" (Printexc.to_string exn);

  match exn with
  | G.Error msg ->
      let row = model#get_iter path in
      let row = model#iter_children ~nth:0 (Some row) in

      let hdata =
        { state=IsLeaf; content=ErrorMessage msg; visited=false; hiveh=None } in
      store_hdata t row hdata;

      model#set ~row ~column:t.name_col (markup_escape msg)

  | exn ->
      (* unexpected exception: re-raise it *)
      raise exn

(* Called when the top level registry node has been opened and the
 * hive file was downloaded to the cache file successfully.
 *)
and when_downloaded_registry ({ model = model } as t) path () =
  debug "when_downloaded_registry";
  let row = model#get_iter path in
  let hdata = get_hdata t row in
  let h = Option.get hdata.hiveh in

  (* Continue as if expanding any other hive node. *)
  let root = Hivex.root h in
  expand_hive_node t row root

(* Called instead of {!when_downloaded_registry} if the download failed. *)
and when_downloaded_registry_fail ({ model = model } as t) path exn =
  debug "when_downloaded_registry_fail: %s" (Printexc.to_string exn);

  match exn with
  | G.Error msg
  | Hivex.Error (_, _, msg) ->
      let row = model#get_iter path in
      let row = model#iter_children ~nth:0 (Some row) in

      let hdata =
        { state=IsLeaf; content=ErrorMessage msg; visited=false; hiveh=None } in
      store_hdata t row hdata;

      model#set ~row ~column:t.name_col (markup_escape msg)

  | exn ->
      (* unexpected exception: re-raise it *)
      raise exn

(* Expand a hive node. *)
and expand_hive_node ({ model = model } as t) row node =
  debug "expand_hive_node";
  let hdata = get_hdata t row in
  let h = Option.get hdata.hiveh in

  (* Read the hive entries (values, subkeys) at this node and add them
   * to the tree.
   *)
  let values = Hivex.node_values h node in
  let cmp v1 v2 = UTF8.compare (Hivex.value_key h v1) (Hivex.value_key h v2) in
  Array.sort cmp values;
  Array.iter (
    fun value ->
      let row = model#append ~parent:row () in
      make_leaf t row (RegValue value) (Some h);
      model#set ~row ~column:t.name_col (markup_of_regvalue h value);
      model#set ~row ~column:t.size_col (markup_of_regvaluesize h value);
      model#set ~row ~column:t.date_col (markup_of_regvaluetype h value);
  ) values;

  let children = Hivex.node_children h node in
  let cmp n1 n2 = UTF8.compare (Hivex.node_name h n1) (Hivex.node_name h n2) in
  Array.sort cmp children;
  Array.iter (
    fun node ->
      let row = model#append ~parent:row () in
      make_node t row (RegKey node) (Some h);
      model#set ~row ~column:t.name_col (markup_of_regkey h node);
  ) children;

  (* Remove the placeholder "Loading" entry.  NB. Must be done AFTER
   * adding the other entries, or else Gtk will unexpand the row.
   *)
  (try
     let row = find_child_node_by_content t row Loading in
     ignore (model#remove row)
   with Invalid_argument _ | Not_found -> ()
  );

  (* The original entry has now been loaded, so update its state. *)
  hdata.state <- IsNode;
  set_visited t row