bmap/bmap.go

package bmap

import (
	"archive/tar"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"crypto/sha1"
	"encoding/hex"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"
	"syscall"

	"github.com/frostschutz/go-fibmap"
)

const defaultBMAPHash = "0000000000000000000000000000000000000000"

// BlockRange represents a range of blocks of written data
type BlockRange struct {
	XMLName xml.Name `xml:"Range"`
	Start   int64
	End     int64
	Hash    string `xml:"sha1,attr"`
}

// Range is an intermediate data type for xml conversion
type Range struct {
	XMLName xml.Name `xml:"Range"`
	Range   string   `xml:",chardata"`
	Hash    string   `xml:"sha1,attr"`
}

// MarshalXML is used to convert BlockRange to a Range object for XML output
func (b *BlockRange) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	//fmt.Println(e)
	//fmt.Println(start)

	r := ""
	if b.Start == b.End {
		r = fmt.Sprintf("%d", b.Start)
	} else {
		r = fmt.Sprintf("%d-%d", b.Start, b.End)
	}
	bl := Range{Range: r, Hash: b.Hash}
	err := e.Encode(bl)

	return err
}

// UnmarshalXML is used to convert a Range object from XML input into a BlockRange object
func (b *BlockRange) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {

	r := Range{}
	err := d.DecodeElement(&r, &start)
	if err != nil {
		fmt.Println(err.Error())
		return nil
	}
	b.Hash = r.Hash
	type pair struct {
		First  int64
		Second int64
	}
	//p := []int64{}
	p := pair{}

	n, err := fmt.Sscanf(r.Range, "%d-%d", &p.First, &p.Second)
	if err != nil || n != 2 {

		n, err := fmt.Sscanf(r.Range, "%d", &p.First)
		if err != nil || n != 1 { //more than one block
			fmt.Println("Multiple Blocks:", err.Error())
			return err
		}

		b.Start = p.First
		b.End = p.First

	} else {
		b.Start = p.First
		b.End = p.Second
	}

	return nil
}

// BlockMap represents a list of Block Ranges
type BlockMap struct {
	Range []BlockRange
}

// BMap contains all of the information that needs to be represented in the bmap file
type BMap struct {
	XMLName           xml.Name `xml:"bmap"`
	Version           string   `xml:"version,attr"`
	ImageSize         int64
	BlockSize         int64
	BlocksCount       int64
	MappedBlocksCount int64
	BmapFileSHA1      string
	BlockMap          BlockMap `xml:"BlockMap"`
}

// NewBMap creates a new BMap representation by reading an image file
func NewBMap(filename string) (BMap, error) {
	stat := syscall.Stat_t{}
	syscall.Stat(filename, &stat)
	blockSize := stat.Blksize

	fd, err := os.Open(filename)
	defer fd.Close()
	if err != nil {
		return BMap{}, err
	}

	blockMap, total := getBlockMap(fd, blockSize)
	return BMap{Version: "1.3", ImageSize: stat.Size, BlockSize: blockSize, BlocksCount: stat.Size / blockSize, MappedBlocksCount: total, BmapFileSHA1: defaultBMAPHash, BlockMap: BlockMap{blockMap}}, nil
}

// XMLOutput converts a BMap object to it's corresponding XML output
func (b *BMap) XMLOutput() ([]byte, error) {
	output, err := xml.MarshalIndent(b, "", "    ")
	if err != nil {
		fmt.Printf("error: %v\n", err)
	}
	return append([]byte(xml.Header), output...), err
}

// Write generates initial XML output, takes the SHA signature, and writes updated XML to file
func (b *BMap) Write(outputFilename string) error {
	output, err := b.XMLOutput()
	if err != nil {
		return err
	}

	bmapHash := getSHA1Hash(bytes.NewBuffer(output))

	b.BmapFileSHA1 = bmapHash
	output, err = b.XMLOutput()
	if err != nil {
		return err
	}
	outputFile, err := os.Create(outputFilename)
	defer outputFile.Close()
	if err != nil {
		return err
	}

	outputFile.Write(output)

	return err
}

// Load reads and converts a bmap file to a BMap object
func Load(filename string) (BMap, error) {
	b := BMap{}
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return b, err
	}

	err = xml.Unmarshal(data, &b)
	if err != nil {
		return b, err
	}

	// Verify SHA Hash

	bmapHash := b.BmapFileSHA1

	b.BmapFileSHA1 = defaultBMAPHash
	output, err := b.XMLOutput()
	if err != nil {
		//fmt.Printf("error: %v\n", err)
		return b, err
	}

	if bmapHash != getSHA1Hash(bytes.NewBuffer(output)) {
		return b, errors.New("XML SHA Signature does not match")
	}

	b.BmapFileSHA1 = bmapHash

	return b, err
}

// Copy copies an input file and uses the BMap data to create a new image file
func (b *BMap) Copy(input string, output string) error {
	inFile, err := os.Open(input)
	defer inFile.Close()
	if err != nil {
		return err
	}

	outFile, err := os.Create(output)
	defer outFile.Close()
	if err != nil {
		return err
	}
	err = outFile.Truncate(b.ImageSize)
	if err != nil {
		return err
	}

	var reader io.ReadSeeker

	switch {
	case strings.HasSuffix(input, ".gz"):

		reader, err = getGZReader(inFile)
		if err != nil {
			return err
		}

	default:
		reader = inFile
	}

	for _, block := range b.BlockMap.Range {
		reader.Seek(block.Start*b.BlockSize, 0)
		outFile.Seek(block.Start*b.BlockSize, 0)
		length := (block.End - block.Start + 1) * b.BlockSize
		written, err := io.CopyN(outFile, reader, length)
		if err != nil {
			return err

		}
		if written != length {
			return errors.New("Unable to copy")

		}

	}

	return nil
}

func getTarReader(reader io.Reader) (io.ReadSeeker, error) {
	tarReader := tar.NewReader(reader)

	data, err := ioutil.ReadAll(tarReader)
	if err != nil {
		return nil, err
	}
	return bytes.NewReader(data), nil

}

func getTarGZReader(reader io.Reader) (io.ReadSeeker, error) {
	gzReader, err := gzip.NewReader(reader)
	if err != nil {
		return nil, err
	}
	tarReader := tar.NewReader(gzReader)

	data, err := ioutil.ReadAll(tarReader)
	if err != nil {
		return nil, err
	}
	return bytes.NewReader(data), nil

}

func getGZReader(reader io.Reader) (io.ReadSeeker, error) {
	gzReader, err := gzip.NewReader(reader)
	if err != nil {
		return nil, err
	}
	data, err := ioutil.ReadAll(gzReader)
	if err != nil {
		return nil, err
	}
	return bytes.NewReader(data), nil
}

func getBZReader(reader io.Reader) (io.ReadSeeker, error) {
	bzReader := bzip2.NewReader(reader)

	data, err := ioutil.ReadAll(bzReader)
	if err != nil {
		return nil, err
	}
	return bytes.NewReader(data), nil
}

// getBlockMap finds all of the ranges of written blocks in a file/image
func getBlockMap(fd *os.File, blockSize int64) ([]BlockRange, int64) {
	//fd, _ := os.Open(filename)
	f := fibmap.NewFibmapFile(fd)

	holes := f.SeekDataHole()
	blockMap := make([]BlockRange, len(holes)/2)
	currentBlockRange := BlockRange{}
	mappedBlocks := int64(0)
	for i, v := range holes {
		if i%2 == 0 { //BlockRange found @
			currentBlockRange = BlockRange{Start: v / blockSize}
		} else { // Length of block
			length := v / blockSize
			currentBlockRange.End = currentBlockRange.Start + length - 1

			h := sha1.New()
			fd.Seek(currentBlockRange.Start*blockSize, 0)
			io.CopyN(h, fd, v)
			currentBlockRange.Hash = hex.EncodeToString(h.Sum(nil))

			blockMap[i/2] = currentBlockRange
			mappedBlocks += length
		}
	}
	return blockMap, mappedBlocks
}

// getSHA1Hash returns a hex encoded hash from an io.Reader
func getSHA1Hash(r io.Reader) string {
	h := sha1.New()
	io.Copy(h, r)
	return hex.EncodeToString(h.Sum(nil))
}