Fix case where binary name doesn't have any extension.
[virt-df.git] / diskzip / diskzip.ml
1 (* 'diskzip' command for intelligently compressing disk images.
2    (C) Copyright 2007 Richard W.M. Jones, Red Hat Inc.
3    http://libvirt.org/
4
5    This program is free software; you can redistribute it and/or modify
6    it under the terms of the GNU General Public License as published by
7    the Free Software Foundation; either version 2 of the License, or
8    (at your option) any later version.
9
10    This program is distributed in the hope that it will be useful,
11    but WITHOUT ANY WARRANTY; without even the implied warranty of
12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13    GNU General Public License for more details.
14
15    You should have received a copy of the GNU General Public License
16    along with this program; if not, write to the Free Software
17    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
18  *)
19
20 open Unix
21 open Printf
22
23 open Diskzip_gettext.Gettext
24
25 type output = File of string | Dir of string
26 type extcompress = BZip2 | GZip | External of string
27
28 let rec main () =
29   (* Program name changes behaviour. *)
30   let compressing =
31     let name = Sys.argv.(0) in
32     let name = Filename.basename name in (* just the executable name *)
33     let name =                           (* remove .opt or .exe *)
34       try Filename.chop_extension name
35       with Invalid_argument("Filename.chop_extension") -> name in
36     let name = String.lowercase name in
37     match name with
38     | "diskzcat" -> false
39     | "diskzip" -> true
40     | name ->
41         eprintf
42           (f_"diskzip: unknown executable name '%s', assuming 'diskzip'\n")
43           name in
44   let compressing = ref compressing in
45
46   (* Command line argument parsing. *)
47   let version () =
48     printf "diskzip\n"; (* XXX version XXX *)
49     exit 0
50   in
51
52   let output = ref None in
53   let set_output path =
54     if !output <> None then (
55       prerr_endline (s_"diskzip: '-o' option cannot appear more than once");
56       exit 2
57     );
58     try
59       let statbuf = stat path in
60       if statbuf.st_kind = S_DIR then
61         output := Some (Dir path)
62       else
63         output := Some (File path)
64     with
65     (* No such file or directory, assume it's a file output. *)
66     | Unix_error (ENOENT, _, _) -> output := Some (File path)
67   in
68
69   (* By default we don't use any external compression program. *)
70   let extcompress = ref None in
71   let set_extcompress t () =
72     if !extcompress <> None then (
73       prerr_endline (s_"diskzip: '-z' or '-j' cannot appear more than once");
74       exit 2
75     );
76     extcompress := Some t
77   in
78
79   let force = ref false in
80
81   let argspec = Arg.align [
82     "-d", Arg.Clear compressing,
83       " " ^ s_ "Uncompress (default: depends on executable name)";
84     "--debug", Arg.Set Diskimage.debug,
85       " " ^ s_ "Debug mode (default: false)";
86     "-f", Arg.Set force,
87       " " ^ s_"Force compress even if stdout looks like a tty";
88     "-j", Arg.Unit (set_extcompress BZip2),
89       " " ^ s_"Pipe the output/input through bzip2";
90     "-o", Arg.String set_output,
91       "path " ^ s_"Set the output filename or directory name";
92     "-p", Arg.String (fun prog -> set_extcompress (External prog) ()),
93       "prog " ^ s_"Pipe the output/input through external program";
94     "--version", Arg.Unit version,
95       " " ^ s_"Display version and exit";
96     "-z", Arg.Unit (set_extcompress GZip),
97       " " ^ s_"Pipe the output/input through gzip";
98   ] in
99
100   let args = ref [] in
101   let anon_fun str = args := str :: !args in
102   let usage_msg = s_"diskzip: Intelligently compress disk images
103
104 SUMMARY
105   diskzip [-options] disk.img [disk.img ...] > output.dz
106   diskzcat [-options] output.dz > disk.img
107
108 OPTIONS" in
109
110   Arg.parse argspec anon_fun usage_msg;
111
112   (* Turn refs back into normal values. *)
113   let compressing = !compressing in
114   let extcompress = !extcompress in
115   let output = !output in
116   let force = !force in
117   let args = !args in
118
119   (* Check the arguments make sense. *)
120   if compressing && output <> None then (
121     prerr_endline (s_"diskzip: '-o' option cannot be used when compressing");
122     exit 2
123   );
124   if compressing && args = [] then (
125     prerr_endline (s_"diskzip: no input");
126     exit 2
127   );
128   if compressing && not force && isatty stdout then (
129     prerr_endline (s_"diskzip: compressed data not written to a terminal, use '-f' to force");
130     exit 2
131   );
132
133   (* Run the compression or decompression functions. *)
134   if compressing then
135     go_compress extcompress args
136   else
137     go_decompress ?output extcompress args
138
139 (* Do compression. *)
140 and go_compress extcompress images =
141   (* Create a Diskimage machine description from the requested images.  This
142    * also checks that everything we need is readable.
143    *)
144   let machine =
145     Diskimage.open_machine "diskzip" (List.map (fun n -> (n,n)) images) in
146
147   (* Scan the images for filesystems. *)
148   let machine = Diskimage.scan_machine machine in
149
150   (* Create ownership tables. *)
151   let ownership = Diskimage.create_ownership machine in
152
153
154
155   (* Redirect output through external pipe if asked. *)
156   (match extcompress with
157    | None -> ()
158    | Some prog ->
159        let prog, progargs =
160          match prog with
161          | BZip2 -> "bzip2", [|"bzip2"; "-c"|]
162          | GZip -> "gzip", [|"gzip"; "-c"|]
163          | External prog -> "sh", [|"sh"; "-c"; prog |] in
164        let rfd, wfd = pipe () in
165        let pid = fork () in
166        if pid = 0 then (                (* child *)
167          close wfd;
168          dup2 rfd stdin;
169          close rfd;
170          execvp prog progargs
171        ) else (                         (* parent *)
172          close rfd;
173          dup2 wfd stdout;
174          close wfd
175        )
176   )
177
178
179
180
181
182
183
184
185
186
187 and go_decompress ?output extcompress args =
188   (* Read the input, which may be a single named file, or a series of
189    * files (we just concatenate them).  We may have to feed the input
190    * through an external program.
191    *)
192   let () =
193     match args with
194     | [] -> ()                          (* Reading from stdin. *)
195     | [file] ->                         (* Read the named file. *)
196         let fd = openfile file [O_RDONLY] 0 in
197         dup2 fd stdin;
198         close fd
199     | files ->                          (* Concatenate files. *)
200         let rfd, wfd = pipe () in
201         let pid = fork () in
202         if pid = 0 then (               (* child *)
203           close rfd;
204           dup2 wfd stdout;
205           close wfd;
206           execvp "cat" (Array.of_list ("cat" :: "--" :: files))
207         ) else (                        (* parent *)
208           close wfd;
209           dup2 rfd stdin;
210           close rfd
211         )
212   in
213   (match extcompress with
214    | None -> ()
215    | Some prog ->
216        let prog, progargs =
217          match prog with
218          | BZip2 -> "bzip2", [|"bzip2"; "-cd"|]
219          | GZip -> "gzip", [|"gzip"; "-cd"|]
220          | External prog -> "sh", [|"sh"; "-c"; prog |] in
221        let rfd, wfd = pipe () in
222        let pid = fork () in
223        if pid = 0 then (                (* child *)
224          close rfd;
225          dup2 wfd stdout;
226          close wfd;
227          execvp prog progargs
228        ) else (                         (* parent *)
229          close wfd;
230          dup2 rfd stdin;
231          close rfd
232        )
233   )
234
235 (*
236   let header = read_header () in
237   XXX
238
239 *)
240
241
242
243
244
245
246
247
248 (*
249 (* Since we have the wonderful pa_bitmatch, might as well use it to
250  * define a robust binary format for the compressed files.
251  *)
252 and write_header ... =
253   let bs = BITSTRING {
254     0xD152 : 16; 0x01 : 8; 0x00 : 8;    (* file magic, version 1.0 *)
255     nr_disks : 8;                       (* number of disks being packed *)
256     
257
258
259
260   } in
261   
262 and read_header () =
263   (* Diskzip headers are limited to overall max size of 1024 bytes. *)
264   let bs = Bitmatch.bitstring_of_file_descr_max stdin 1024 in
265
266   bitmatch bs with
267   | { 0xD152 : 16;                      (* file magic *)
268       0x01 : 8; (_ as minor) : 8;       (* major, minor versions *)
269     } ->
270
271   (* Is this a later version (major != 1)? *)
272   | { 0xD152 : 16;                      (* file magic *)
273       (_ as major) : 8; (_ as minor) : 8 } when major <> 1 ->
274       eprintf (f_"diskzip: archive version %d.%d, this program only understands version 1.x")
275         major minor;
276       exit 1
277
278   (* If it looks like gzip or bzip2, exit with an informative error. *)
279   | { 0o37 : 8; 0o213 : 8 } ->          (* gzip *)
280       prerr_endline (s_"diskzip: This looks like a gzip archive. Did you mean to pass the '-z' option?");
281       exit 1
282   | { "BZh" : 24 : string } ->          (* bzip2 *)
283       prerr_endline (s_"diskzip: This looks like a bzip2 archive. Did you mean to pass the '-j' option?");
284       exit 1
285
286   (* If it looks like a disk image (MBR), give an error. *)
287   | { _ : 4080 : bitstring; 0x55 : 8; 0xAA : 8 } ->
288       prerr_endline (s_"diskzip: This looks like a disk image. Did you mean to compress it?");
289       exit 1
290
291   | { _ } ->
292       prerr_endline (s_"diskzip: Not a diskzip archive.");
293       exit 1
294 *)
295
296 let () = main ()