de6db3ee42cbc2f269c857f5545e3087988bffee
[virt-similarity.git] / 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