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 Int63.Operators
  24 open Diskzip_gettext.Gettext
  25 module Bitmap = Diskzip_bitmap
  26
  27 type output = File of string | Dir of string
  28 type extcompress = BZip2 | GZip | External of string
  29
  30 let rec main () =
  31   (* Program name changes behaviour. *)
  32   let compressing =
  33     let name = Sys.argv.(0) in
  34     let name = Filename.basename name in (* just the executable name *)
  35     let name =                           (* remove .opt or .exe *)
  36       try Filename.chop_extension name
  37       with Invalid_argument("Filename.chop_extension") -> name in
  38     let name = String.lowercase name in
  39     match name with
  40     | "diskzcat" -> false
  41     | "diskzip" -> true
  42     | _ ->
  43         eprintf
  44           (f_"diskzip: unknown executable name '%s', assuming 'diskzip'\n")
  45           name;
  46         true in
  47   let compressing = ref compressing in
  48
  49   (* Command line argument parsing. *)
  50   let version () =
  51     printf "diskzip\n"; (* XXX version XXX *)
  52     exit 0
  53   in
  54
  55   let output = ref None in
  56   let set_output path =
  57     if !output <> None then (
  58       prerr_endline (s_"diskzip: '-o' option cannot appear more than once");
  59       exit 2
  60     );
  61     try
  62       let statbuf = stat path in
  63       if statbuf.st_kind = S_DIR then
  64         output := Some (Dir path)
  65       else
  66         output := Some (File path)
  67     with
  68     (* No such file or directory, assume it's a file output. *)
  69     | Unix_error (ENOENT, _, _) -> output := Some (File path)
  70   in
  71
  72   (* By default we don't use any external compression program. *)
  73   let extcompress = ref None in
  74   let set_extcompress t () =
  75     if !extcompress <> None then (
  76       prerr_endline (s_"diskzip: '-z' or '-j' cannot appear more than once");
  77       exit 2
  78     );
  79     extcompress := Some t
  80   in
  81
  82   let force = ref false in
  83
  84   let argspec = Arg.align [
  85     "-d", Arg.Clear compressing,
  86       " " ^ s_ "Uncompress (default: depends on executable name)";
  87     "--debug", Arg.Set Diskimage.debug,
  88       " " ^ s_ "Debug mode (default: false)";
  89     "-f", Arg.Set force,
  90       " " ^ s_"Force compress even if stdout looks like a tty";
  91     "-j", Arg.Unit (set_extcompress BZip2),
  92       " " ^ s_"Pipe the output/input through bzip2";
  93     "-o", Arg.String set_output,
  94       "path " ^ s_"Set the output filename or directory name";
  95     "-p", Arg.String (fun prog -> set_extcompress (External prog) ()),
  96       "prog " ^ s_"Pipe the output/input through external program";
  97     "--version", Arg.Unit version,
  98       " " ^ s_"Display version and exit";
  99     "-z", Arg.Unit (set_extcompress GZip),
 100       " " ^ s_"Pipe the output/input through gzip";
 101   ] in
 102
 103   let args = ref [] in
 104   let anon_fun str = args := str :: !args in
 105   let usage_msg = s_"diskzip: Intelligently compress disk images
 106
 107 SUMMARY
 108   diskzip [-options] disk.img [disk.img ...] > output.dz
 109   diskzcat [-options] output.dz > disk.img
 110
 111 OPTIONS" in
 112
 113   Arg.parse argspec anon_fun usage_msg;
 114
 115   (* Turn refs back into normal values. *)
 116   let compressing = !compressing in
 117   let extcompress = !extcompress in
 118   let output = !output in
 119   let force = !force in
 120   let args = !args in
 121
 122   (* Check the arguments make sense. *)
 123   if compressing && output <> None then (
 124     prerr_endline (s_"diskzip: '-o' option cannot be used when compressing");
 125     exit 2
 126   );
 127   if compressing && args = [] then (
 128     prerr_endline (s_"diskzip: no input");
 129     exit 2
 130   );
 131   if compressing && not force && isatty stdout then (
 132     prerr_endline (s_"diskzip: compressed data not written to a terminal, use '-f' to force");
 133     exit 2
 134   );
 135
 136   (* Run the compression or decompression functions. *)
 137   if compressing then
 138     go_compress extcompress args
 139   else
 140     go_decompress ?output extcompress args
 141
 142 (* Do compression. *)
 143 and go_compress extcompress images =
 144   (* Create a Diskimage machine description from the requested images.  This
 145    * also checks that everything we need is readable.
 146    *)
 147   let machine =
 148     Diskimage.open_machine "diskzip" (List.map (fun n -> (n,n)) images) in
 149
 150   (* Scan the images for filesystems. *)
 151   let machine = Diskimage.scan_machine machine in
 152
 153   (* Create ownership tables. *)
 154   let ownership = Diskimage.create_ownership machine in
 155
 156   (* Create ownership bitmap for each disk. *)
 157   List.iter (
 158     fun { Diskimage.d_name = name; d_dev = disk } ->
 159       let blocksize = disk#blocksize in
 160       let size = disk#size in           (* Size in bytes. *)
 161       let nr_blocks = size /^ blocksize in (* Number of disk sectors. *)
 162
 163       if !Diskimage.debug then
 164         eprintf "Creating bitmap for %s (%s sectors) ...\n%!"
 165           disk#name (Int63.to_string nr_blocks);
 166
 167       (* Create an empty bitmap, one bit per sector. *)
 168       let bitmap = Bitmap.create nr_blocks in
 169
 170       (* Get the lookup function for this disk. *)
 171       let lookup = Diskimage.get_owners_lookup machine ownership disk in
 172
 173       (* Lookup each sector. *)
 174       Bitmap.iter_set (
 175         fun blk _ ->
 176           let owners = lookup blk in
 177           false
 178       ) bitmap
 179   ) machine.Diskimage.m_disks;
 180
 181   (* Redirect output through external pipe if asked. *)
 182   (match extcompress with
 183    | None -> ()
 184    | Some prog ->
 185        let prog, progargs =
 186          match prog with
 187          | BZip2 -> "bzip2", [|"bzip2"; "-c"|]
 188          | GZip -> "gzip", [|"gzip"; "-c"|]
 189          | External prog -> "sh", [|"sh"; "-c"; prog |] in
 190        let rfd, wfd = pipe () in
 191        let pid = fork () in
 192        if pid = 0 then (                (* child *)
 193          close wfd;
 194          dup2 rfd stdin;
 195          close rfd;
 196          execvp prog progargs
 197        ) else (                         (* parent *)
 198          close rfd;
 199          dup2 wfd stdout;
 200          close wfd
 201        )
 202   )
 203
 204
 205
 206
 207
 208
 209
 210
 211
 212
 213 and go_decompress ?output extcompress args =
 214   (* Read the input, which may be a single named file, or a series of
 215    * files (we just concatenate them).  We may have to feed the input
 216    * through an external program.
 217    *)
 218   let () =
 219     match args with
 220     | [] -> ()                          (* Reading from stdin. *)
 221     | [file] ->                         (* Read the named file. *)
 222         let fd = openfile file [O_RDONLY] 0 in
 223         dup2 fd stdin;
 224         close fd
 225     | files ->                          (* Concatenate files. *)
 226         let rfd, wfd = pipe () in
 227         let pid = fork () in
 228         if pid = 0 then (               (* child *)
 229           close rfd;
 230           dup2 wfd stdout;
 231           close wfd;
 232           execvp "cat" (Array.of_list ("cat" :: "--" :: files))
 233         ) else (                        (* parent *)
 234           close wfd;
 235           dup2 rfd stdin;
 236           close rfd
 237         )
 238   in
 239   (match extcompress with
 240    | None -> ()
 241    | Some prog ->
 242        let prog, progargs =
 243          match prog with
 244          | BZip2 -> "bzip2", [|"bzip2"; "-cd"|]
 245          | GZip -> "gzip", [|"gzip"; "-cd"|]
 246          | External prog -> "sh", [|"sh"; "-c"; prog |] in
 247        let rfd, wfd = pipe () in
 248        let pid = fork () in
 249        if pid = 0 then (                (* child *)
 250          close rfd;
 251          dup2 wfd stdout;
 252          close wfd;
 253          execvp prog progargs
 254        ) else (                         (* parent *)
 255          close wfd;
 256          dup2 rfd stdin;
 257          close rfd
 258        )
 259   )
 260
 261 (*
 262   let header = read_header () in
 263   XXX
 264
 265 *)
 266
 267
 268
 269
 270
 271
 272
 273
 274 (*
 275 (* Since we have the wonderful pa_bitmatch, might as well use it to
 276  * define a robust binary format for the compressed files.
 277  *)
 278 and write_header ... =
 279   let bs = BITSTRING {
 280     0xD152 : 16; 0x01 : 8; 0x00 : 8;    (* file magic, version 1.0 *)
 281     nr_disks : 8;                       (* number of disks being packed *)
 282
 283
 284
 285
 286   } in
 287
 288 and read_header () =
 289   (* Diskzip headers are limited to overall max size of 1024 bytes. *)
 290   let bs = Bitmatch.bitstring_of_file_descr_max stdin 1024 in
 291
 292   bitmatch bs with
 293   | { 0xD152 : 16;                      (* file magic *)
 294       0x01 : 8; (_ as minor) : 8;       (* major, minor versions *)
 295     } ->
 296
 297   (* Is this a later version (major != 1)? *)
 298   | { 0xD152 : 16;                      (* file magic *)
 299       (_ as major) : 8; (_ as minor) : 8 } when major <> 1 ->
 300       eprintf (f_"diskzip: archive version %d.%d, this program only understands version 1.x")
 301         major minor;
 302       exit 1
 303
 304   (* If it looks like gzip or bzip2, exit with an informative error. *)
 305   | { 0o37 : 8; 0o213 : 8 } ->          (* gzip *)
 306       prerr_endline (s_"diskzip: This looks like a gzip archive. Did you mean to pass the '-z' option?");
 307       exit 1
 308   | { "BZh" : 24 : string } ->          (* bzip2 *)
 309       prerr_endline (s_"diskzip: This looks like a bzip2 archive. Did you mean to pass the '-j' option?");
 310       exit 1
 311
 312   (* If it looks like a disk image (MBR), give an error. *)
 313   | { _ : 4080 : bitstring; 0x55 : 8; 0xAA : 8 } ->
 314       prerr_endline (s_"diskzip: This looks like a disk image. Did you mean to compress it?");
 315       exit 1
 316
 317   | { _ } ->
 318       prerr_endline (s_"diskzip: Not a diskzip archive.");
 319       exit 1
 320 *)
 321
 322 let () = main ()