Added XML generation and loading. Basic copy operations.

bmap.go

@@ -0,0 +1,310 @@
+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"
+
+type Block struct {
+	XMLName xml.Name `xml:"Range"`
+	Start   int64
+	End     int64
+	Hash    string `xml:"sha1,attr"`
+}
+
+type Range struct {
+	XMLName xml.Name `xml:"Range"`
+	Range   string   `xml:",chardata"`
+	Hash    string   `xml:"sha1,attr"`
+}
+
+func (b *Block) 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
+}
+
+func (b *Block) 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
+}
+
+type BlockMap struct {
+	Range []Block
+}
+
+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"`
+}
+
+func NewBMap(filename string) BMap {
+	stat := syscall.Stat_t{}
+	syscall.Stat(filename, &stat)
+	blockSize := stat.Blksize
+
+	fd, _ := os.Open(filename)
+
+	defer fd.Close()
+
+	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}}
+}
+
+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
+}
+
+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
+}
+
+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
+}
+
+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
+}
+
+func getBlockMap(fd *os.File, blockSize int64) ([]Block, int64) {
+	//fd, _ := os.Open(filename)
+	f := fibmap.NewFibmapFile(fd)
+
+	holes := f.SeekDataHole()
+	blockMap := make([]Block, len(holes)/2)
+	currentBlock := Block{}
+	mappedBlocks := int64(0)
+	for i, v := range holes {
+		if i%2 == 0 { //Block found @
+			currentBlock = Block{Start: v / blockSize}
+		} else { // Length of block
+			length := v / blockSize
+			currentBlock.End = currentBlock.Start + length - 1
+
+			h := sha1.New()
+			fd.Seek(currentBlock.Start*blockSize, 0)
+			io.CopyN(h, fd, v)
+			currentBlock.Hash = hex.EncodeToString(h.Sum(nil))
+
+			blockMap[i/2] = currentBlock
+			mappedBlocks += length
+		}
+	}
+	return blockMap, mappedBlocks
+}
+
+func getSHA1Hash(r io.Reader) string {
+	h := sha1.New()
+	io.Copy(h, r)
+	return hex.EncodeToString(h.Sum(nil))
+}

tool/bmap.go

@@ -0,0 +1,27 @@
+package main
+
+import "dev.justinjudd.org/justin/bmap"
+
+const imageInputFilename = "/tmp/file.img"
+const GZCompressedImageInputFilename = "/tmp/file.gz"
+const bmapOutputFilename = "/tmp/test.bmap"
+const imageOutputFilename = "/tmp/test.img"
+
+func main() {
+
+	/*
+		b := bmap.NewBMap(imageInputFilename)
+		b.Write("/tmp/test.bmap")
+	*/
+
+	b, err := bmap.Load(bmapOutputFilename)
+	if err != nil {
+		println(err.Error())
+	}
+	//fmt.Printf("%#v\n", b)
+	//err = b.Copy(imageInputFilename, imageOutputFilename)
+	err = b.Copy(GZCompressedImageInputFilename, imageOutputFilename)
+	if err != nil {
+		println(err.Error())
+	}
+}