#!/usr/bin/env ocaml

#load "unix.cma";;
#directory "+guestfs";;
#load "mlguestfs.cma";;

open Printf
module G = Guestfs

let (//) = Filename.concat
let quote = Filename.quote
let (+^) = Int64.add
let (-^) = Int64.sub
let ( *^ ) = Int64.mul
let (/^) = Int64.div

type ('a,'b) alternate = Either of 'a | Or of 'b
type visit_r = OK | Prune

let rec main () =
  if Array.length Sys.argv < 2 then (
    eprintf "usage: %s disk [output.disk]\n" Sys.executable_name;
    exit 1
  );
  let disk = Sys.argv.(1) in

  (* Make an overlay over the input disk. *)
  let overlay = Filename.temp_file "ovl" ".qcow2" in
  let cmd =
    sprintf "qemu-img create -f qcow2 -b %s %s" (quote disk) (quote overlay) in
  prerr_endline cmd;
  let i = Sys.command cmd in
  if i <> 0 then exit i;

  let g = new G.guestfs () in
  g#add_drive_opts ~format:"qcow2" overlay;
  g#launch ();

  (* Estimate the size of the original disk. *)
  let orig_size = estimate_size g in
  printf "original size [uncompressed, estimate]: %s\n%!"
    (human_size orig_size);

  (* Mount up the filesystems. *)
  inspect_and_mount g;

  (* Get "inspection signature" before removing anything. *)
  let orig_sig = inspection_signature g in

  (* Start deleting things. *)
  remove g "/home";
  remove g "/opt";
  remove g "/usr/etc";
  remove g "/usr/games";
  remove g "/usr/include";
  remove g "/usr/lib";
  remove g "/usr/lib64";
  remove g "/usr/libexec";
  remove g "/usr/local";
  remove g "/usr/share" ~except:[
    "/usr/share/cirros";
    "/usr/share/icons";
    "/usr/share/pixmaps";
  ];
  remove g "/usr/src";
  remove g "/var" ~except:[
    "/var/lib"                          (* /var/lib for package dbs *)
  ];

  g#umount_all ();

  (* Get "inspection signature" after. *)
  inspect_and_mount g;
  let final_sig = inspection_signature g in
  g#umount_all ();

  if orig_sig <> final_sig then (
    eprintf "oops, that changed the results of inspection\n";
    exit 1
  );

  (* Estimate the size of the final disk. *)
  let final_size = estimate_size g in
  printf "final size [uncompressed, estimate]: %s\n%!" (human_size final_size);

  (* If there's an output disk, then run virt-sparsify and xz to
   * generate the final image.
   *)
  if Array.length Sys.argv >= 3 then (
    let output = Sys.argv.(2) in

    let cmd =
      sprintf "virt-sparsify --quiet --format qcow2 %s --convert raw %s"
        (quote overlay) (quote output) in
    prerr_endline cmd;
    let i = Sys.command cmd in
    if i <> 0 then exit i;

    let cmd = sprintf "du -h %s" (quote output) in
    prerr_endline cmd;
    let i = Sys.command cmd in
    if i <> 0 then exit i;

    let cmd = sprintf "rm -f %s.xz; xz --best --block-size=%Ld %s"
      (quote output) (16L *^ 1024L *^ 1024L) (quote output) in
    prerr_endline cmd;
    let i = Sys.command cmd in
    if i <> 0 then exit i;

    let cmd = sprintf "ls -lh %s.xz" (quote output) in
    prerr_endline cmd;
    let i = Sys.command cmd in
    if i <> 0 then exit i
  );

  Unix.unlink overlay

(* Estimate the size of a disk, uncompressed bytes, after using
 * virt-sparsify.
 *)
and estimate_size g =
  let size = ref 0L in

  (* Virt-sparsify is usually able to recover any unused space in VGs,
   * so we don't have to examine that.  Just look at filesystems.
   *)
  let fses = g#list_filesystems () in
  List.iter (
    fun (dev, typ) ->
      (* Assume swap contains random data and virt-sparsify would
       * be able to replace it with a blank swap partition.
       *)
      if typ <> "swap" then (
        let devsize = g#blockdev_getsize64 dev in
        let mounted = try g#mount_ro dev "/"; true with G.Error _ -> false in

        (* Assume virt-sparsify wouldn't be able to touch an unmountable dev. *)
        if not mounted then
          size := !size +^ devsize
        else (
          (* For a mountable filesystem, assume virt-sparsify would be
           * able to recover all unused space.
           *)
          let statvfs = g#statvfs "/" in
          size := !size +^ (devsize -^ (statvfs.G.bsize *^ statvfs.G.bfree));

          g#umount_all ()
        )
      )
  ) fses;

  !size

and human_size i =
  let sign, i = if i < 0L then "-", Int64.neg i else "", i in

  if i < 1024L then
    sprintf "%s%Ld" sign i
  else (
    let f = Int64.to_float i /. 1024. in
    let i = i /^ 1024L in
    if i < 1024L then
      sprintf "%s%.1fK" sign f
    else (
      let f = Int64.to_float i /. 1024. in
      let i = i /^ 1024L in
      if i < 1024L then
        sprintf "%s%.1fM" sign f
      else (
        let f = Int64.to_float i /. 1024. in
        (*let i = i /^ 1024L in*)
        sprintf "%s%.1fG" sign f
      )
    )
  )

(* This is like rm -rf except:
 * - It doesn't delete directories or symlinks.
 * - The optional ?except parameter can be used to "save" a prefix from
 *   being deleted.
 *)
and remove g ?(except = []) path =
  if g#exists path then (
    if g#is_dir ~followsymlinks:false path then (
      let rec f pathname _ _ { G.mode = mode } =
        if is_except pathname then
          Prune                         (* prune ?except subtrees *)
        else if is_dir mode then OK     (* don't delete directories *)
        else if is_lnk mode then OK     (* don't delete symlinks *)
        else (
          g#rm pathname;                (* otherwise delete it *)
          OK
        )
      and is_except pathname =
        List.exists (
          fun prefix -> starts_with pathname prefix
        ) except
      and starts_with pathname prefix =
        let pathnamelen = String.length pathname in
        let prefixlen = String.length prefix in
        prefixlen <= pathnamelen &&
          String.sub pathname 0 prefixlen = prefix
      in

      visit g ~f path
    ) else
      g#rm path
  )

(* See cat/ls.c: visit() *)
and visit g ~f ?(depth = 0) dir =
  (* Call the callback function on the top level directory. *)
  let r =
    if depth = 0 then (
      let stat = g#lstat dir in
      f dir dir None stat
    )
    else OK in

  match r with
  | Prune -> ()
  | OK ->
    let names = g#ls dir in
    let stats = g#lstatlist dir names in

    assert (Array.length names = Array.length stats);

    for i = 0 to Array.length names - 1 do
      let path =
        if dir = "/" then "/" ^ names.(i)
        else dir ^ "/" ^ names.(i) in
      let r = f path dir (Some names.(i)) stats.(i) in

      if r = OK && is_dir stats.(i).G.mode then
        visit g ~f ~depth:(depth+1) path
    done

and is_dir mode = Int64.logand mode 0o170_000L = 0o040_000L

and is_lnk mode = Int64.logand mode 0o170_000L = 0o120_000L

and inspect_and_mount g =
  let roots = g#inspect_os () in
  if Array.length roots <> 1 then
    failwith "inspect_vm: no operating systems or multi-boot guest";

  let root = roots.(0) in

  let mps = g#inspect_get_mountpoints root in
  let cmp (a,_) (b,_) = compare (String.length a) (String.length b) in
  let mps = List.sort cmp mps in
  List.iter (
    fun (mp, dev) ->
      try g#mount dev mp with G.Error msg -> eprintf "%s (ignored)\n%!" msg
  ) mps

(* This lets us check that no parts of inspection change before
 * and after removing bits from the guest.
 *)
and inspection_signature g =
  let root = (g#inspect_get_roots ()).(0) in
  (g#inspect_get_arch root,
   g#inspect_get_distro root,
   g#inspect_get_drive_mappings root,
   g#inspect_get_filesystems root,
   g#inspect_get_format root,
   g#inspect_get_hostname root,
   g#inspect_get_icon root,
   g#inspect_get_major_version root,
   g#inspect_get_minor_version root,
   g#inspect_get_mountpoints root,
   g#inspect_get_package_format root,
   g#inspect_get_package_management root,
   g#inspect_get_product_name root,
   g#inspect_get_product_variant root,
   g#inspect_get_type root,
   (try Either (g#inspect_get_windows_current_control_set root)
    with G.Error msg -> Or msg),
   (try Either (g#inspect_get_windows_systemroot root)
    with G.Error msg -> Or msg),
   g#inspect_is_live root,
   g#inspect_is_multipart root,
   g#inspect_is_netinst root,
   g#inspect_list_applications2 root)

let () = main ()