(* 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

open Int63.Operators

let debug = ref false

class virtual device =
object (self)
  method virtual size : int63
  method virtual name : string
  method virtual blocksize : int63
  method virtual map_block : int63 -> (device * int63) list
  method virtual contiguous : Int63.t -> Int63.t

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

    let blocksize = self#blocksize in

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

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

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

    (* Copy the first block (partial). *)
    (match self#map_block first_blk with
     | [] -> not_mapped_error ()
     | (dev, base) :: _ ->
	 let len =
	   min len (blocksize -^ 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#map_block blk with
	 | [] -> not_mapped_error ()
	 | (dev, base) :: _ ->
	     let str = dev#read ~^0 self#blocksize in
	     Buffer.add_string buf str
	);
	loop (Int63.succ blk)
      )
    in
    loop (Int63.succ first_blk);

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

    assert (Int63.to_int 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, String.length str lsl 3)
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 = Int63.of_int64 (LargeFile.fstat fd).LargeFile.st_size in
object (self)
  inherit device
  method read offset len =
    let offset = Int63.to_int64 offset in
    let len = Int63.to_int len in
    ignore (LargeFile.lseek fd offset SEEK_SET);
    let str = String.make len '\000' in
    ignore (read fd str 0 len);
    str
  method size = size
  method name = filename
  method blocksize = blocksize
  method map_block _ = []
  method contiguous offset =
    size -^ offset
  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 < ~^0 || len < ~^0 || offset +^ len > size then
      invalid_arg (
	sprintf "%s: tried to read outside device boundaries (%s/%s/%s)"
	  name (Int63.to_string offset) (Int63.to_string len)
	  (Int63.to_string size)
      );
    dev#read (start+^offset) len
  method blocksize = blocksize
  method map_block i = [dev, i *^ blocksize +^ start]
  method contiguous offset =
    size -^ offset
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 = dev#read
  method blocksize = new_blocksize
  method map_block new_blk =
    let orig_blk = new_blk *^ new_blocksize /^ dev#blocksize in
    dev#map_block orig_blk
  method contiguous offset = dev#size -^ offset
end

(* The null device.  Any attempt to read generates an error. *)
let null_device : device =
object
  inherit device
  method read _ _ = assert false
  method size = ~^0
  method name = "null"
  method blocksize = ~^1
  method map_block _ = assert false
  method contiguous _ = ~^0
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_dev : device;			(* Partitions (whole) device. *)
  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 : int63;			(* Block size (bytes). *)
  fs_blocks_total : int63;		(* Total blocks. *)
  fs_is_swap : bool;			(* If swap, following not valid. *)
  fs_blocks_reserved : int63;		(* Blocks reserved for super-user. *)
  fs_blocks_avail : int63;		(* Blocks free (available). *)
  fs_blocks_used : int63;		(* Blocks in use. *)
  fs_inodes_total : int63;		(* Total inodes. *)
  fs_inodes_reserved : int63;		(* Inodes reserved for super-user. *)
  fs_inodes_avail : int63;		(* Inodes free (available). *)
  fs_inodes_used : int63;		(* Inodes in use. *)
}

(* Physical volumes. *)
and pv = {
  lvm_plugin_id : lvm_plugin_id;        (* The LVM plug-in. *)
  pv_dev : device;			(* Device covering whole PV. *)
  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 uniq ?(cmp = Pervasives.compare) = function
  | [] -> []
  | [x] -> [x]
  | x :: y :: xs when cmp x y = 0 ->
      uniq (x :: xs)
  | x :: y :: xs ->
      x :: uniq (y :: xs)

let sort_uniq ?cmp xs =
  let xs = ExtList.List.sort ?cmp xs in
  let xs = uniq ?cmp xs in
  xs

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