diff --git a/bmap.go b/bmap.go new file mode 100644 index 0000000..faff06b --- /dev/null +++ b/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)) +} diff --git a/tool/bmap.go b/tool/bmap.go new file mode 100644 index 0000000..9e99df4 --- /dev/null +++ b/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()) + } +}