update-iso.ml

   1 #!/usr/bin/ocamlrun /usr/bin/ocaml
   2 #load "unix.cma";;
   3
   4 (* update-iso.ml attaches an updated 'virt-p2v.ml' file to the end of
   5  * an ISO image.  This is just for quick developer builds because it
   6  * takes ages to rebuild a full ISO.
   7  *
   8  * Copyright (C) 2007-2008 Red Hat Inc.
   9  * Written by Richard W.M. Jones <rjones@redhat.com>
  10  *
  11  * This program is free software; you can redistribute it and/or modify
  12  * it under the terms of the GNU General Public License as published by
  13  * the Free Software Foundation; either version 2 of the License, or
  14  * (at your option) any later version.
  15  *
  16  * This program is distributed in the hope that it will be useful,
  17  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  18  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  19  * GNU General Public License for more details.
  20  *
  21  * You should have received a copy of the GNU General Public License
  22  * along with this program; if not, write to the Free Software
  23  * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  24  *)
  25
  26 open Printf
  27 open Unix
  28
  29 (* CD-ROM format has 2048 byte "sectors" and the final ISO image
  30  * must be a multiple of 2048 bytes in size.
  31  *
  32  * The basic plan is that we will use store the file followed by a
  33  * trailer, all rounded up to 2048 bytes:
  34  *
  35  * +------ - - - - - - - -----+------------+----------------+
  36  * | file                     | padding    | trailer        |
  37  * |                          |            | magic | ....   |
  38  * +------ - - - - - - - -----+------------+----------------+
  39  * |<----    total size is a multiple of 2048 bytes    ---->|
  40  *
  41  * The magic string is used to identify that there is an
  42  * attachment, and is followed by a few extra fields which
  43  * identify the file start and true file size.  (Note the
  44  * original filename is not stored because it is not needed).
  45  *
  46  * Attachments can be stacked.  This script only deals with the
  47  * top-most attachment (ie. the one at the very end of the file).
  48  * If you want to attach lots of files, a better way is to
  49  * stuff them into a tarball or ZIP file and attach that.
  50  *)
  51
  52 let magic = "ISOATTACHMENT002"
  53 let magiclen = String.length magic  (* = 16 bytes *)
  54 let trailerlen = magiclen + 8 + 8   (* magic + file start + true size *)
  55
  56 (* Ugh, would be really nice to have 64 bit pack/unpack functions
  57  * in the stdlib, instead of this ugliness ...
  58  *)
  59 let string_of_int64 i =
  60   let str = String.create 8 in
  61   let shift_mask i shift =
  62     Char.chr (Int64.to_int (Int64.shift_right_logical i shift) land 0xff)
  63   in
  64   str.[0] <- shift_mask i 56;  str.[1] <- shift_mask i 48;
  65   str.[2] <- shift_mask i 40;  str.[3] <- shift_mask i 32;
  66   str.[4] <- shift_mask i 24;  str.[5] <- shift_mask i 16;
  67   str.[6] <- shift_mask i 8;   str.[7] <- shift_mask i 0;
  68   str
  69 let int64_of_string str =
  70   let i = ref 0L in
  71   let add offs shift =
  72     i :=
  73       Int64.logor
  74         (Int64.shift_left (Int64.of_int (Char.code str.[offs])) shift) !i
  75   in
  76   add 0 56;  add 1 48;  add 2 40;  add 3 32;
  77   add 4 24;  add 5 16;  add 6 8;   add 7 0;
  78   !i
  79
  80 let rec main () =
  81   let args = Array.to_list Sys.argv in
  82   match args with
  83   | [ _; "add"; isoname; attachment ] -> (* add an attachment *)
  84       do_add isoname attachment
  85   | [ _; "delete"; isoname ] ->         (* delete any attachment *)
  86       do_delete isoname
  87   | [ _; "get"; isoname; output ] ->    (* get attachment *)
  88       do_get isoname output
  89   | [ _; "has"; isoname ] ->            (* is there an attachment? *)
  90       do_has isoname
  91   | _ ->
  92       eprintf "\
  93 update-iso.ml add foo.iso file
  94   Attach 'file' to 'foo.iso'.
  95
  96 update-iso.ml delete foo.iso
  97   Remove attachment (if any) from 'foo.iso'.
  98
  99 update-iso.ml get foo.iso file
 100   Get attachment from 'foo.iso' and save it as 'file'.
 101
 102 update-iso.ml has foo.iso
 103   Exit with 0 (true) if there is an attachment.
 104   Exit with 1 (false) if there is no attachment.
 105   Exit with 2 if there was some other error, eg. file not found.
 106
 107 Note that attachments are stacked, so you can add more than one
 108 attachment.  In this case 'get' operation returns the most recently
 109 added and 'delete' operation deletes only the most recently added.
 110 ";
 111       exit 1
 112
 113 and do_has isoname =
 114   let fd =
 115     try openfile isoname [O_RDONLY] 0
 116     with Unix_error (err, syscall, param) ->
 117       eprintf "%s:%s: %s\n" syscall param (error_message err);
 118       exit 2 in
 119   try
 120     ignore (LargeFile.lseek fd (Int64.of_int ~-trailerlen) SEEK_END);
 121     let buf = String.create magiclen in
 122     if read fd buf 0 magiclen <> magiclen then exit 1;
 123     if buf <> magic then exit 1;
 124     exit 0
 125   with
 126     Unix_error (err, syscall, param) ->
 127       eprintf "%s:%s: %s\n" syscall param (error_message err);
 128       exit 1
 129
 130 and do_add isoname attachment =
 131   let fd = openfile isoname [O_APPEND; O_WRONLY] 0 in
 132
 133   let iso_size = (LargeFile.fstat fd).LargeFile.st_size in
 134   if Int64.logand iso_size 2047L <> 0L then
 135     failwith "ISO image is not a multiple of 2048 bytes in size";
 136
 137   (* Copy the attachment itself to the end of the file. *)
 138   let fd2 = openfile attachment [O_RDONLY] 0 in
 139   let bufsize = 4 * 1024 in
 140   let buffer = String.create bufsize in
 141   let rec copy size =
 142     let n = read fd2 buffer 0 bufsize in
 143     if n > 0 then (
 144       ignore (write fd buffer 0 n);
 145       copy (size + n)
 146     )
 147     else size
 148   in
 149   let file_size = copy 0 in
 150   close fd2;
 151
 152   (* How much padding to use so that file_size + trailer + padding
 153    * = 2048 bytes multiple?
 154    *)
 155   let padding_size =
 156     let size = file_size + trailerlen in
 157     let over = size land 2047 in
 158     if over > 0 then 2048-over else 0 in
 159   assert ((padding_size + file_size + trailerlen) land 2047 = 0);
 160
 161   (* Write the padding. *)
 162   ignore (write fd (String.make padding_size 'x') 0 padding_size);
 163
 164   (* Write the magic. *)
 165   ignore (write fd magic 0 magiclen);
 166
 167   (* Write the file start and true size. *)
 168   let buffer = string_of_int64 iso_size in
 169   ignore (write fd buffer 0 8);
 170   let buffer = string_of_int64 (Int64.of_int file_size) in
 171   ignore (write fd buffer 0 8);
 172
 173   close fd
 174
 175 and do_delete isoname =
 176   let fd = openfile isoname [O_RDWR] 0 in
 177   ignore (LargeFile.lseek fd (Int64.of_int ~-trailerlen) SEEK_END);
 178   let buf = String.create magiclen in
 179   if read fd buf 0 magiclen <> magiclen || buf <> magic then
 180     failwith "no attachment found";
 181
 182   (* Read the start offset of the file. *)
 183   let buf = String.create 8 in
 184   if read fd buf 0 8 <> 8 then
 185     failwith "cannot read attachment size";
 186   let offset = int64_of_string buf in
 187
 188   (* Truncate to start of the file. *)
 189   LargeFile.ftruncate fd offset;
 190
 191   close fd
 192
 193 and do_get isoname output =
 194   let fd = openfile isoname [O_RDONLY] 0 in
 195   ignore (LargeFile.lseek fd (Int64.of_int ~-trailerlen) SEEK_END);
 196   let buf = String.create magiclen in
 197   if read fd buf 0 magiclen <> magiclen || buf <> magic then
 198     failwith "no attachment found";
 199
 200   (* Read the start and size. *)
 201   let buf = String.create 8 in
 202   if read fd buf 0 8 <> 8 then
 203     failwith "cannot read attachment offset";
 204   let offset = int64_of_string buf in
 205   let buf = String.create 8 in
 206   if read fd buf 0 8 <> 8 then
 207     failwith "cannot read attachment size";
 208   let size = Int64.to_int (int64_of_string buf) in
 209
 210   (* Seek to beginning of the attachment. *)
 211   ignore (LargeFile.lseek fd offset SEEK_SET);
 212
 213   (* Copy out the attachment. *)
 214   let fd2 = openfile output [O_WRONLY; O_CREAT; O_TRUNC] 0o644 in
 215   let bufsize = 4 * 1024 in
 216   let buffer = String.create bufsize in
 217   let rec copy remaining =
 218     if remaining > 0 then (
 219       let n = min remaining bufsize in
 220       let n = read fd buffer 0 n in
 221       if n = 0 then failwith "corrupted or partial attachment";
 222       ignore (write fd2 buffer 0 n);
 223       copy (remaining - n)
 224     )
 225   in
 226   copy size;
 227   close fd2;
 228
 229   close fd
 230
 231 let () = main ()