Added XML generation and loading. Basic copy operations.
This commit is contained in:
justin
parent
b35b674bdb
commit
878efd05bb
|
|
@ -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))
|
||||
}
|
||||
|
|
@ -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…
Reference in New Issue