(* Diskimage library for reading disk images.
   (C) Copyright 2007-2008 Richard W.M. Jones, Red Hat Inc.
   http://libvirt.org/

   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., 675 Mass Ave, Cambridge, MA 02139, USA.
 *)

open Printf
open Unix

let debug = ref false

let ( +* ) = Int32.add
let ( -* ) = Int32.sub
let ( ** ) = Int32.mul
let ( /* ) = Int32.div

let ( +^ ) = Int64.add
let ( -^ ) = Int64.sub
let ( *^ ) = Int64.mul
let ( /^ ) = Int64.div

class virtual device =
object (self)
  method virtual size : int64
  method virtual name : string
  method virtual blocksize : int
  method virtual mapblock : int64 -> (device * int64) list

  (* Block-based read.  Inefficient so normally overridden in subclasses. *)
  method read offset len =
    if offset < 0L || len < 0 then
      invalid_arg "device: read: negative offset or length";

    let blocksize64 = Int64.of_int self#blocksize in

    (* Break the request into blocks.
     * Find the first and last blocks of this request.
     *)
    let first_blk = offset /^ blocksize64 in
    let offset_in_first_blk = offset -^ first_blk *^ blocksize64 in
    let last_blk = (offset +^ Int64.of_int (len-1)) /^ blocksize64 in

    (* Buffer for the result. *)
    let buf = Buffer.create len in

    let not_mapped_error () = invalid_arg "device: read: block not mapped" in

    (* Copy the first block (partial). *)
    (match self#mapblock first_blk with
     | [] -> not_mapped_error ()
     | (dev, base) :: _ ->
	 let len =
	   min len (Int64.to_int (blocksize64 -^ offset_in_first_blk)) in
	 let str = dev#read (base +^ offset_in_first_blk) len in
	 Buffer.add_string buf str
    );

    (* Copy the middle blocks. *)
    let rec loop blk =
      if blk < last_blk then (
	(match self#mapblock blk with
	 | [] -> not_mapped_error ()
	 | (dev, base) :: _ ->
	     let str = dev#read 0L self#blocksize in
	     Buffer.add_string buf str
	);
	loop (Int64.succ blk)
      )
    in
    loop (Int64.succ first_blk);

    (* Copy the last block (partial). *)
    if first_blk < last_blk then (
      match self#mapblock last_blk with
      | [] -> not_mapped_error ()
      | (dev, base) :: _ ->
	  let len = (offset +^ Int64.of_int len) -^ last_blk *^ blocksize64 in
	  let len = Int64.to_int len in
	  let str = dev#read 0L len in
	  Buffer.add_string buf str
    );

    assert (len = Buffer.length buf);
    Buffer.contents buf

  (* Helper method to read a chunk of data into a bitstring. *)
  method read_bitstring offset len =
    let str = self#read offset len in
    (str, 0, len * 8)
end

(* A concrete device which just direct-maps a file or /dev device. *)
class block_device filename blocksize =
  let fd = openfile filename [ O_RDONLY ] 0 in
  let size = (LargeFile.fstat fd).LargeFile.st_size in
object (self)
  inherit device
  method read offset len =
    ignore (LargeFile.lseek fd offset SEEK_SET);
    let str = String.make len '\000' in
    read fd str 0 len;
    str
  method size = size
  method name = filename
  method blocksize = blocksize
  method mapblock _ = []
  method close () = close fd
end

(* A linear offset/size from an underlying device. *)
class offset_device name start size blocksize (dev : device) =
object
  inherit device
  method name = name
  method size = size
  method read offset len =
    if offset < 0L || len < 0 || offset +^ Int64.of_int len > size then
      invalid_arg (
	sprintf "%s: tried to read outside device boundaries (%Ld/%d/%Ld)"
	  name offset len size
      );
    dev#read (start+^offset) len
  method blocksize = blocksize
  method mapblock i = [dev, i *^ Int64.of_int blocksize +^ start]
end

(* A device with just a modified block size. *)
class blocksize_overlay new_blocksize (dev : device) =
object
  inherit device
  method name = dev#name
  method size = dev#size
  method read offset len = dev#read offset len
  method blocksize = new_blocksize
  method mapblock new_blk =
    let orig_blk =
      new_blk *^ Int64.of_int new_blocksize /^ Int64.of_int dev#blocksize in
    dev#mapblock orig_blk
end

(* The null device.  Any attempt to read generates an error. *)
let null_device : device =
object
  inherit device
  method read _ _ = assert false
  method size = 0L
  method name = "null"
  method blocksize = 1
  method mapblock _ = assert false
end

type machine = {
  m_name : string;			(* Machine name. *)
  m_disks : disk list;			(* Machine disks. *)
  m_lv_filesystems :
    (lv * filesystem) list;		(* Machine LV filesystems. *)
}
and disk = {
  d_name : string;			(* Device name (eg "hda") *)

  (* About the device itself. *)
  d_dev : block_device;			(* Disk device. *)
  d_content : disk_content;		(* What's on it. *)
}
and disk_content =
  [ `Unknown				(* Not probed or unknown. *)
  | `Partitions of partitions		(* Contains partitions. *)
  | `Filesystem of filesystem		(* Contains a filesystem directly. *)
  | `PhysicalVolume of pv		(* Contains an LVM PV. *)
  ]

(* Partitions. *)

and partitions = {
  parts_plugin_id : parts_plugin_id;	(* Partitioning scheme. *)
  parts : partition list		(* Partitions. *)
}
and partition = {
  part_status : partition_status;	(* Bootable, etc. *)
  part_type : int;			(* Partition filesystem type. *)
  part_dev : device;			(* Partition device. *)
  part_content : partition_content;	(* What's on it. *)
}
and partition_status = Bootable | Nonbootable | Malformed | NullEntry
and partition_content =
  [ `Unknown				(* Not probed or unknown. *)
  | `Filesystem of filesystem		(* Filesystem. *)
  | `PhysicalVolume of pv		(* Contains an LVM PV. *)
  ]

(* Filesystems (also swap devices). *)
and filesystem = {
  fs_plugin_id : fs_plugin_id;		(* Filesystem. *)
  fs_dev : device;			(* Device containing the filesystem. *)
  fs_blocksize : int;			(* Block size (bytes). *)
  fs_blocks_total : int64;		(* Total blocks. *)
  fs_is_swap : bool;			(* If swap, following not valid. *)
  fs_blocks_reserved : int64;		(* Blocks reserved for super-user. *)
  fs_blocks_avail : int64;		(* Blocks free (available). *)
  fs_blocks_used : int64;		(* Blocks in use. *)
  fs_inodes_total : int64;		(* Total inodes. *)
  fs_inodes_reserved : int64;		(* Inodes reserved for super-user. *)
  fs_inodes_avail : int64;		(* Inodes free (available). *)
  fs_inodes_used : int64;		(* Inodes in use. *)
}

(* Physical volumes. *)
and pv = {
  lvm_plugin_id : lvm_plugin_id;        (* The LVM plug-in. *)
  pv_uuid : string;			(* UUID. *)
}

(* Logical volumes. *)
and lv = {
  lv_dev : device;			(* Logical volume device. *)
}

and parts_plugin_id = string
and fs_plugin_id = string
and lvm_plugin_id = string

(* Convert a UUID (containing '-' chars) to canonical form. *)
let canonical_uuid uuid =
  let uuid' = String.make 32 ' ' in
  let j = ref 0 in
  for i = 0 to String.length uuid - 1 do
    if !j >= 32 then invalid_arg "canonical_uuid";
    let c = uuid.[i] in
    if c <> '-' then ( uuid'.[!j] <- c; incr j )
  done;
  if !j <> 32 then invalid_arg "canonical_uuid";
  uuid'

(* This version by Isaac Trotts. *)
let group_by ?(cmp = Pervasives.compare) ls =
  let ls' =
    List.fold_left
      (fun acc (day1, x1) ->
         match acc with
             [] -> [day1, [x1]]
           | (day2, ls2) :: acctl ->
               if cmp day1 day2 = 0
               then (day1, x1 :: ls2) :: acctl
               else (day1, [x1]) :: acc)
      []
      ls
  in
  let ls' = List.rev ls' in
  List.map (fun (x, xs) -> x, List.rev xs) ls'

let rec range a b =
  if a < b then a :: range (a+1) b
  else []