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   (* Create ownership bitmap for each disk. *)
 154
 155
 156   (* Redirect output through external pipe if asked. *)
 157   (match extcompress with
 158    | None -> ()
 159    | Some prog ->
 160        let prog, progargs =
 161          match prog with
 162          | BZip2 -> "bzip2", [|"bzip2"; "-c"|]
 163          | GZip -> "gzip", [|"gzip"; "-c"|]
 164          | External prog -> "sh", [|"sh"; "-c"; prog |] in
 165        let rfd, wfd = pipe () in
 166        let pid = fork () in
 167        if pid = 0 then (                (* child *)
 168          close wfd;
 169          dup2 rfd stdin;
 170          close rfd;
 171          execvp prog progargs
 172        ) else (                         (* parent *)
 173          close rfd;
 174          dup2 wfd stdout;
 175          close wfd
 176        )
 177   )
 178
 179
 180
 181
 182
 183
 184
 185
 186
 187
 188 and go_decompress ?output extcompress args =
 189   (* Read the input, which may be a single named file, or a series of
 190    * files (we just concatenate them).  We may have to feed the input
 191    * through an external program.
 192    *)
 193   let () =
 194     match args with
 195     | [] -> ()                          (* Reading from stdin. *)
 196     | [file] ->                         (* Read the named file. *)
 197         let fd = openfile file [O_RDONLY] 0 in
 198         dup2 fd stdin;
 199         close fd
 200     | files ->                          (* Concatenate files. *)
 201         let rfd, wfd = pipe () in
 202         let pid = fork () in
 203         if pid = 0 then (               (* child *)
 204           close rfd;
 205           dup2 wfd stdout;
 206           close wfd;
 207           execvp "cat" (Array.of_list ("cat" :: "--" :: files))
 208         ) else (                        (* parent *)
 209           close wfd;
 210           dup2 rfd stdin;
 211           close rfd
 212         )
 213   in
 214   (match extcompress with
 215    | None -> ()
 216    | Some prog ->
 217        let prog, progargs =
 218          match prog with
 219          | BZip2 -> "bzip2", [|"bzip2"; "-cd"|]
 220          | GZip -> "gzip", [|"gzip"; "-cd"|]
 221          | External prog -> "sh", [|"sh"; "-c"; prog |] in
 222        let rfd, wfd = pipe () in
 223        let pid = fork () in
 224        if pid = 0 then (                (* child *)
 225          close rfd;
 226          dup2 wfd stdout;
 227          close wfd;
 228          execvp prog progargs
 229        ) else (                         (* parent *)
 230          close wfd;
 231          dup2 rfd stdin;
 232          close rfd
 233        )
 234   )
 235
 236 (*
 237   let header = read_header () in
 238   XXX
 239
 240 *)
 241
 242
 243
 244
 245
 246
 247
 248
 249 (*
 250 (* Since we have the wonderful pa_bitmatch, might as well use it to
 251  * define a robust binary format for the compressed files.
 252  *)
 253 and write_header ... =
 254   let bs = BITSTRING {
 255     0xD152 : 16; 0x01 : 8; 0x00 : 8;    (* file magic, version 1.0 *)
 256     nr_disks : 8;                       (* number of disks being packed *)
 257
 258
 259
 260
 261   } in
 262
 263 and read_header () =
 264   (* Diskzip headers are limited to overall max size of 1024 bytes. *)
 265   let bs = Bitmatch.bitstring_of_file_descr_max stdin 1024 in
 266
 267   bitmatch bs with
 268   | { 0xD152 : 16;                      (* file magic *)
 269       0x01 : 8; (_ as minor) : 8;       (* major, minor versions *)
 270     } ->
 271
 272   (* Is this a later version (major != 1)? *)
 273   | { 0xD152 : 16;                      (* file magic *)
 274       (_ as major) : 8; (_ as minor) : 8 } when major <> 1 ->
 275       eprintf (f_"diskzip: archive version %d.%d, this program only understands version 1.x")
 276         major minor;
 277       exit 1
 278
 279   (* If it looks like gzip or bzip2, exit with an informative error. *)
 280   | { 0o37 : 8; 0o213 : 8 } ->          (* gzip *)
 281       prerr_endline (s_"diskzip: This looks like a gzip archive. Did you mean to pass the '-z' option?");
 282       exit 1
 283   | { "BZh" : 24 : string } ->          (* bzip2 *)
 284       prerr_endline (s_"diskzip: This looks like a bzip2 archive. Did you mean to pass the '-j' option?");
 285       exit 1
 286
 287   (* If it looks like a disk image (MBR), give an error. *)
 288   | { _ : 4080 : bitstring; 0x55 : 8; 0xAA : 8 } ->
 289       prerr_endline (s_"diskzip: This looks like a disk image. Did you mean to compress it?");
 290       exit 1
 291
 292   | { _ } ->
 293       prerr_endline (s_"diskzip: Not a diskzip archive.");
 294       exit 1
 295 *)
 296
 297 let () = main ()