similarity.ml

   1 (* virt-similarity
   2  * Copyright (C) 2013 Red Hat Inc.
   3  *
   4  * This program is free software; you can redistribute it and/or modify
   5  * it under the terms of the GNU General Public License as published by
   6  * the Free Software Foundation; either version 2 of the License, or
   7  * (at your option) any later version.
   8  *
   9  * This program is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12  * GNU General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU General Public License along
  15  * with this program; if not, write to the Free Software Foundation, Inc.,
  16  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  17  *)
  18
  19 open Printf
  20
  21 open Cladogram
  22 open Utils
  23
  24 let blocksize = 65536
  25 module HashAlgorithm = Hash.MD5
  26 let hash_bytes_per_block = HashAlgorithm.hash_bytes_per_block
  27
  28 (* Int64 operators. *)
  29 let (+^)   = Int64.add
  30 let (-^)   = Int64.sub
  31 let ( *^ ) = Int64.mul
  32 let (/^)   = Int64.div
  33
  34 (* Read command line parameters. *)
  35 let n, filenames, debug =
  36   let display_version () =
  37     printf "%s %s (%s)\n"
  38       Config.package_name Config.package_version
  39       (if Config.ocamlopt then "native" else "bytecode");
  40     exit 0
  41   in
  42
  43   let debug = ref false in
  44
  45   let argspec = Arg.align [
  46     "--debug", Arg.Set debug, " Enable debugging";
  47     "-d", Arg.Set debug, " Enable debugging";
  48     "--version", Arg.Unit display_version, " Display version number and exit";
  49     "-V", Arg.Unit display_version, " Display version number and exit";
  50   ] in
  51
  52   let filenames = ref [] in
  53   let collect_filenames str = filenames := str :: !filenames in
  54
  55   let usage_msg = "
  56 virt-similarity: Find clusters of similar/cloned virtual machines
  57 Copyright (C) 2013 Red Hat Inc.
  58
  59 For full documentation see the virt-similarity(1) man page.
  60
  61 Usage:
  62
  63   virt-similarity [options] disk.img disk.img [disk.img ...]
  64
  65 You must supply at least one disk image.  You can supply disk
  66 images in most common formats (raw, qcow2, vmdk, etc.)
  67
  68 Options:
  69 " in
  70
  71   (* Do the argument parsing. *)
  72   Arg.parse argspec collect_filenames usage_msg;
  73
  74   (* Check the arguments. *)
  75   let filenames = Array.of_list (List.rev !filenames) in
  76   let n = Array.length filenames in
  77   if n < 2 then (
  78     eprintf "virt-similarity: At least two disk images must be specified.\n";
  79     exit 1
  80   );
  81
  82   let debug = !debug in
  83
  84   n, filenames, debug
  85
  86 (* Read in the cache file. *)
  87 let cache = Cache.read_cache_file ()
  88
  89 (* Read the disk images, hash them, and update the cache. *)
  90 let read_disk_image filename =
  91   let g = new Guestfs.guestfs () in
  92   g#add_drive_ro filename;
  93   g#launch ();
  94
  95   let dev = "/dev/sda" in
  96   let size = g#blockdev_getsize64 dev in
  97   let rec loop offset last_percent accum =
  98     if offset <= size -^ Int64.of_int blocksize then (
  99       let percent = 100L *^ offset /^ size in
 100       if percent <> last_percent then printf "%Ld%%\r%!" percent;
 101
 102       let block = g#pread_device dev blocksize offset in
 103       let hash = HashAlgorithm.digest block in
 104       loop (offset +^ Int64.of_int blocksize) percent (hash :: accum)
 105     )
 106     else accum
 107   in
 108   let hashes_reversed = loop 0L (-1L) [] in
 109   g#close ();
 110
 111   Array.of_list (List.rev hashes_reversed)
 112
 113 let cache =
 114   List.fold_left (
 115     fun cache filename ->
 116       let cache, hashes =
 117         match Cache.get_hash cache filename with
 118         | Some hashes ->
 119           if debug then
 120             printf "%s: disk image is already in the cache\n%!" filename;
 121           cache, hashes
 122         | None ->
 123           printf "%s: reading disk image ...\n%!" filename;
 124           let hashes = read_disk_image filename in
 125           Cache.update_hash cache filename hashes, hashes in
 126       if debug then
 127         printf "%s: number of blocks = %d\n" filename (Array.length hashes);
 128       cache
 129   ) cache (Array.to_list filenames)
 130
 131 (* Write the updated cache file. *)
 132 let () = Cache.write_cache_file cache
 133
 134 (* Work out the similarity for each pair of guests and store it in a
 135  * matrix, where matrix.(i).(j) is the distance between filenames.(i)
 136  * and filenames.(j).
 137  *)
 138 let hash_of_zero =
 139   let zero_string = String.make blocksize (Char.chr 0) in
 140   HashAlgorithm.digest zero_string
 141
 142 let calculate_distance hash1 hash2 =
 143   let hash1 = Array.to_list hash1 in
 144   let hash2 = Array.to_list hash2 in
 145   let rec loop = function
 146     | [], [] -> 0
 147     | (x :: xs), [] when x = hash_of_zero -> loop (xs, [])
 148     | (x :: xs), [] -> 1 + loop (xs, [])
 149     | [], (y :: ys) when y = hash_of_zero -> loop ([], ys)
 150     | [], (y :: ys) -> 1 + loop ([], ys)
 151     | (x :: xs), (y :: ys) when x = y -> loop (xs, ys)
 152     | (x :: xs), (y :: ys) -> 1 + loop (xs, ys)
 153   in
 154   loop (hash1, hash2)
 155
 156 let matrix =
 157   let matrix = Array.make_matrix n n 0 in
 158   List.iter (
 159     fun (i, j) ->
 160       let hi = Cache.get_hash cache filenames.(i) in
 161       let hj = Cache.get_hash cache filenames.(j) in
 162       match hi, hj with
 163       | Some hi, Some hj ->
 164         let d = calculate_distance hi hj in
 165         if debug then
 166           printf "distance from %s to %s = %d\n" filenames.(i) filenames.(j) d;
 167         matrix.(i).(j) <- d;
 168         matrix.(j).(i) <- d
 169       | _ -> assert false
 170   ) (pairs_of_ints n);
 171   matrix
 172
 173 (* Construct the tree (cladogram). *)
 174 let cladogram = construct_cladogram matrix n
 175
 176 let () =
 177   let format_leaf i = Filename.basename filenames.(i) in
 178   let lines = format_cladogram ~format_leaf cladogram in
 179   List.iter print_endline lines