#!/usr/bin/ocamlrun /usr/bin/ocaml
(* -*- tuareg -*- *)
(* iso-attach attaches an updated 'virt-p2v' file to the end of
 * an ISO image.  This is just for quick developer builds because it
 * takes ages to rebuild a full ISO.
 *
 * Copyright (C) 2007-2008 Red Hat Inc.
 * Written by Richard W.M. Jones <rjones@redhat.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

#load "unix.cma";;

open Printf
open Unix

(* CD-ROM format has 2048 byte "sectors" and the final ISO image
 * must be a multiple of 2048 bytes in size.
 *
 * The basic plan is that we will use store the file followed by a
 * trailer, all rounded up to 2048 bytes:
 *
 * +------ - - - - - - - -----+------------+----------------+
 * | file                     | padding    | trailer        |
 * |                          |            | magic | ....   |
 * +------ - - - - - - - -----+------------+----------------+
 * |<----    total size is a multiple of 2048 bytes    ---->|
 *
 * The magic string is used to identify that there is an
 * attachment, and is followed by a few extra fields which
 * identify the file start and true file size.  (Note the
 * original filename is not stored because it is not needed).
 *
 * Attachments can be stacked.  This script only deals with the
 * top-most attachment (ie. the one at the very end of the file).
 * If you want to attach lots of files, a better way is to
 * stuff them into a tarball or ZIP file and attach that.
 *)

let magic = "ISOATTACHMENT003"
let magiclen = String.length magic  (* = 16 bytes *)
let trailerlen = magiclen + 8 + 8   (* magic + file start + true size *)

(* Ugh, would be really nice to have 64 bit pack/unpack functions
 * in the stdlib, instead of this ugliness ...
 *)
let string_of_int64 i =
  let str = String.create 8 in
  let shift_mask i shift =
    Char.chr (Int64.to_int (Int64.shift_right_logical i shift) land 0xff)
  in
  str.[0] <- shift_mask i 56;  str.[1] <- shift_mask i 48;
  str.[2] <- shift_mask i 40;  str.[3] <- shift_mask i 32;
  str.[4] <- shift_mask i 24;  str.[5] <- shift_mask i 16;
  str.[6] <- shift_mask i 8;   str.[7] <- shift_mask i 0;
  str
let int64_of_string str =
  let i = ref 0L in
  let add offs shift =
    i :=
      Int64.logor
	(Int64.shift_left (Int64.of_int (Char.code str.[offs])) shift) !i
  in
  add 0 56;  add 1 48;  add 2 40;  add 3 32;
  add 4 24;  add 5 16;  add 6 8;   add 7 0;
  !i

let rec main () =
  let args = Array.to_list Sys.argv in
  match args with
  | [ _; "add"; isoname; attachment ] -> (* add an attachment *)
      do_add isoname attachment
  | [ _; "delete"; isoname ] ->		(* delete any attachment *)
      do_delete isoname
  | [ _; "get"; isoname; output ] ->	(* get attachment *)
      do_get isoname output
  | [ _; "has"; isoname ] ->		(* is there an attachment? *)
      do_has isoname
  | _ ->
      eprintf "\
NAME

  iso-attach - Attach files to CD-ROM ISO images.

SYNOPSIS

  iso-attach add foo.iso file
    Attach 'file' to 'foo.iso'.

  iso-attach delete foo.iso
    Remove attachment (if any) from 'foo.iso'.

  iso-attach get foo.iso file
    Get attachment from 'foo.iso' and save it as 'file'.

  iso-attach has foo.iso
    Exit with 0 (true) if there is an attachment.
    Exit with 1 (false) if there is no attachment.
    Exit with 2 if there was some other error, eg. file not found.

DESCRIPTION

iso-attach attaches an updated 'virt-p2v' file to the end of
an ISO image.  This is just for quick developer builds because it
takes ages to rebuild a full ISO.

Note that attachments are stacked, so you can add more than one
attachment.  In this case 'get' operation returns the most recently
added and 'delete' operation deletes only the most recently added.
";
      exit 1

and do_has isoname =
  let fd =
    try openfile isoname [O_RDONLY] 0
    with Unix_error (err, syscall, param) ->
      eprintf "%s:%s: %s\n" syscall param (error_message err);
      exit 2 in
  try
    ignore (LargeFile.lseek fd (Int64.of_int ~-trailerlen) SEEK_END);
    let buf = String.create magiclen in
    if read fd buf 0 magiclen <> magiclen then exit 1;
    if buf <> magic then exit 1;
    exit 0
  with
    Unix_error (err, syscall, param) ->
      eprintf "%s:%s: %s\n" syscall param (error_message err);
      exit 1

and do_add isoname attachment =
  let fd = openfile isoname [O_APPEND; O_WRONLY] 0 in

  let iso_size = (LargeFile.fstat fd).LargeFile.st_size in
  if Int64.logand iso_size 2047L <> 0L then
    failwith "ISO image is not a multiple of 2048 bytes in size";

  (* Copy the attachment itself to the end of the file. *)
  let fd2 = openfile attachment [O_RDONLY] 0 in
  let bufsize = 4 * 1024 in
  let buffer = String.create bufsize in
  let rec copy size =
    let n = read fd2 buffer 0 bufsize in
    if n > 0 then (
      ignore (write fd buffer 0 n);
      copy (size + n)
    )
    else size
  in
  let file_size = copy 0 in
  close fd2;

  (* How much padding to use so that file_size + trailer + padding
   * = 2048 bytes multiple?
   *)
  let padding_size = 
    let size = file_size + trailerlen in
    let over = size land 2047 in
    if over > 0 then 2048-over else 0 in
  let attachment_size = padding_size + file_size + trailerlen in
  assert (attachment_size land 2047 = 0);

  (* Write the padding. *)
  ignore (write fd (String.make padding_size 'x') 0 padding_size);

  (* Write the magic. *)
  ignore (write fd magic 0 magiclen);

  (* Write the attachment size and true file size. *)
  let buffer = string_of_int64 (Int64.of_int attachment_size) in
  ignore (write fd buffer 0 8);
  let buffer = string_of_int64 (Int64.of_int file_size) in
  ignore (write fd buffer 0 8);

  close fd

and do_delete isoname =
  let fd = openfile isoname [O_RDWR] 0 in
  ignore (LargeFile.lseek fd (Int64.of_int ~-trailerlen) SEEK_END);
  let buf = String.create magiclen in
  if read fd buf 0 magiclen <> magiclen || buf <> magic then
    failwith "no attachment found";

  (* Read the size of the attachment *)
  let buf = String.create 8 in
  if read fd buf 0 8 <> 8 then
    failwith "cannot read attachment size";
  let attachment_size = int64_of_string buf in

  let orig_size = LargeFile.lseek fd (Int64.neg attachment_size) SEEK_END in

  (* Truncate to start of the file. *)
  LargeFile.ftruncate fd orig_size;

  close fd

and do_get isoname output =
  let fd = openfile isoname [O_RDONLY] 0 in
  ignore (LargeFile.lseek fd (Int64.of_int ~-trailerlen) SEEK_END);
  let buf = String.create magiclen in
  if read fd buf 0 magiclen <> magiclen || buf <> magic then
    failwith "no attachment found";

  (* Read the attachment size and true size. *)
  let buf = String.create 8 in
  if read fd buf 0 8 <> 8 then
    failwith "cannot read attachment size";
  let attachment_size = int64_of_string buf in
  let buf = String.create 8 in
  if read fd buf 0 8 <> 8 then
    failwith "cannot read file true size";
  let size = Int64.to_int (int64_of_string buf) in

  (* Seek to beginning of the attachment. *)
  ignore (LargeFile.lseek fd (Int64.neg attachment_size) SEEK_END);

  (* Copy out the attachment. *)
  let fd2 = openfile output [O_WRONLY; O_CREAT; O_TRUNC] 0o644 in
  let bufsize = 4 * 1024 in
  let buffer = String.create bufsize in
  let rec copy remaining =
    if remaining > 0 then (
      let n = min remaining bufsize in
      let n = read fd buffer 0 n in
      if n = 0 then failwith "corrupted or partial attachment";
      ignore (write fd2 buffer 0 n);
      copy (remaining - n)
    )
  in
  copy size;
  close fd2;

  close fd

let () = main ()

(* This file must end with a newline. *)