Added XML generation and loading. Basic copy operations.
This commit is contained in:
		
							parent
							
								
									b35b674bdb
								
							
						
					
					
						commit
						878efd05bb
					
				
							
								
								
									
										310
									
								
								bmap.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										310
									
								
								bmap.go
									
									
									
									
									
										Normal file
									
								
							@ -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))
 | 
			
		||||
}
 | 
			
		||||
							
								
								
									
										27
									
								
								tool/bmap.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										27
									
								
								tool/bmap.go
									
									
									
									
									
										Normal file
									
								
							@ -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())
 | 
			
		||||
	}
 | 
			
		||||
}
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user