526 lines
11 KiB
Go
526 lines
11 KiB
Go
package bmap
|
|
|
|
import (
|
|
"archive/tar"
|
|
"bytes"
|
|
"compress/bzip2"
|
|
"compress/gzip"
|
|
"crypto/sha1"
|
|
"encoding/hex"
|
|
"encoding/xml"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"io/ioutil"
|
|
"net/http"
|
|
"os"
|
|
"regexp"
|
|
"strings"
|
|
"syscall"
|
|
|
|
"github.com/frostschutz/go-fibmap"
|
|
//"github.com/uli-go/xz/lzma"
|
|
"code.google.com/p/lzma"
|
|
)
|
|
|
|
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,omitempty"`
|
|
}
|
|
|
|
// 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, addChecksum bool) (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, addChecksum)
|
|
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
|
|
}
|
|
|
|
// LoadFromFile reads and converts a bmap file to a BMap object
|
|
func LoadFromFile(filename string, verify bool) (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
|
|
}
|
|
|
|
if verify {
|
|
// 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
|
|
}
|
|
|
|
// LoadFromReader reads and converts a bmap in an io.Reader to a BMap object
|
|
func LoadFromReader(r io.Reader, verify bool) (BMap, error) {
|
|
b := BMap{}
|
|
data, err := ioutil.ReadAll(r)
|
|
if err != nil {
|
|
return b, err
|
|
}
|
|
|
|
err = xml.Unmarshal(data, &b)
|
|
if err != nil {
|
|
return b, err
|
|
}
|
|
|
|
if verify {
|
|
// 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
|
|
}
|
|
|
|
// loadInput is designed to handle loading local and remote files
|
|
func loadInput(input string) (io.ReadCloser, error) {
|
|
var reader io.ReadCloser
|
|
|
|
inFile, err := os.Open(input)
|
|
//defer inFile.Close()
|
|
reader = inFile
|
|
if err != nil { // Not local file. Try to grab remote file
|
|
switch {
|
|
case strings.HasPrefix(input, "http"):
|
|
//fmt.Println("Get URL", input)
|
|
resp, err := http.Get(input)
|
|
if err != nil {
|
|
return reader, errors.New("Unable to download image from " + input)
|
|
}
|
|
reader = resp.Body
|
|
//fmt.Println("Retrieved", resp.ContentLength, "bytes")
|
|
|
|
}
|
|
//return reader, err
|
|
}
|
|
|
|
return reader, nil
|
|
}
|
|
|
|
// decompressInput "unwraps" - aka decompresses and unarchives image files
|
|
func decompressInput(input string, r io.Reader) (io.Reader, error) {
|
|
var reader io.Reader
|
|
var err error
|
|
|
|
switch {
|
|
case strings.HasSuffix(input, ".tar.gz"), strings.HasSuffix(input, ".tgz"):
|
|
|
|
reader, err = getTarGZReader(r)
|
|
|
|
case strings.HasSuffix(input, ".gz"), strings.HasSuffix(input, ".gzip"):
|
|
|
|
reader, err = getGZReader(r)
|
|
|
|
case strings.HasSuffix(input, ".xz"), strings.HasSuffix(input, ".lzma"):
|
|
|
|
reader, err = getLZMAReader(r)
|
|
|
|
default:
|
|
reader = r
|
|
|
|
}
|
|
|
|
return reader, 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, verify bool) error {
|
|
|
|
var reader io.Reader
|
|
|
|
r, err := loadInput(input)
|
|
defer r.Close()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
reader, err = decompressInput(input, r)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
outFile, err := os.Create(output)
|
|
defer outFile.Close()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
block, err := isBlockDevice(outFile)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if !block { // Can't truncate block devices
|
|
err = outFile.Truncate(b.ImageSize)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
place := int64(0)
|
|
for _, block := range b.BlockMap.Range {
|
|
|
|
diff := block.Start*b.BlockSize - place
|
|
place = block.Start * b.BlockSize
|
|
|
|
_, err := moveReaderForward(reader, diff)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = outFile.Seek(diff, os.SEEK_CUR)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
length := (block.End - block.Start + 1) * b.BlockSize
|
|
place += length
|
|
if len(block.Hash) != 0 && verify {
|
|
//println("checking hash")
|
|
// Verify hash sum
|
|
//var checksumReader bytes.Buffer
|
|
checksumReader := io.TeeReader(reader, outFile)
|
|
h := sha1.New()
|
|
io.CopyN(h, checksumReader, length)
|
|
if block.Hash != hex.EncodeToString(h.Sum(nil)) {
|
|
return fmt.Errorf("Checksum mismatch for blockrange %d-%d", block.Start, block.End)
|
|
}
|
|
} else {
|
|
written, err := io.CopyN(outFile, reader, length)
|
|
if err != nil {
|
|
return err
|
|
|
|
}
|
|
if written != length {
|
|
return errors.New("Unable to copy")
|
|
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
outFile.Sync()
|
|
return nil
|
|
}
|
|
|
|
// Copy is used to copy images to a destination when a bmap file is unavailable
|
|
func Copy(input string, output string) error {
|
|
|
|
reader, err := loadInput(input)
|
|
defer reader.Close()
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
|
|
seeker, err := decompressInput(input, reader)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
outFile, err := os.Create(output)
|
|
defer outFile.Close()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
io.Copy(outFile, seeker)
|
|
|
|
outFile.Sync()
|
|
|
|
return nil
|
|
}
|
|
|
|
func getTarReader(reader io.Reader) (io.ReadSeeker, error) {
|
|
tarReader := tar.NewReader(reader)
|
|
tarReader.Next()
|
|
|
|
data, err := ioutil.ReadAll(tarReader)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return bytes.NewReader(data), nil
|
|
|
|
}
|
|
|
|
func getTarGZReader(reader io.Reader) (io.Reader, error) {
|
|
gzReader, err := gzip.NewReader(reader)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
tarReader := tar.NewReader(gzReader)
|
|
tarReader.Next()
|
|
|
|
return tarReader, nil
|
|
|
|
}
|
|
|
|
func getGZReader(reader io.Reader) (io.Reader, error) {
|
|
return gzip.NewReader(reader)
|
|
}
|
|
|
|
func getTarBZReader(reader io.Reader) (io.Reader, error) {
|
|
bzReader := bzip2.NewReader(reader)
|
|
|
|
tarReader := tar.NewReader(bzReader)
|
|
tarReader.Next()
|
|
|
|
return tarReader, nil
|
|
|
|
}
|
|
|
|
func getBZReader(reader io.Reader) (io.Reader, error) {
|
|
bzReader := bzip2.NewReader(reader)
|
|
|
|
return bzReader, nil
|
|
}
|
|
|
|
func getLZMAReader(reader io.Reader) (io.Reader, error) {
|
|
|
|
/*
|
|
return lzma.NewReader(reader)
|
|
|
|
*/
|
|
xzReader := lzma.NewReader(reader)
|
|
|
|
return xzReader, nil
|
|
}
|
|
|
|
func getXZReader(reader io.Reader) (io.ReadSeeker, error) {
|
|
//return xz.NewSeekReader(reader)
|
|
return nil, nil
|
|
}
|
|
|
|
// getBlockMap finds all of the ranges of written blocks in a file/image
|
|
func getBlockMap(fd *os.File, blockSize int64, addChecksum bool) ([]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
|
|
|
|
if addChecksum {
|
|
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))
|
|
}
|
|
|
|
// moveReaderForward provides seek capabilities for an io.Reader, even ones that aren't an io.Seeker
|
|
func moveReaderForward(r io.Reader, count int64) (int64, error) {
|
|
|
|
seeker, ok := r.(io.Seeker)
|
|
if ok {
|
|
return seeker.Seek(count, os.SEEK_CUR)
|
|
} else {
|
|
return io.CopyN(ioutil.Discard, r, count)
|
|
}
|
|
|
|
}
|
|
|
|
// isBlockDevice checks if a file descriptor is a block device
|
|
func isBlockDevice(fd *os.File) (bool, error) {
|
|
block := false
|
|
s, err := fd.Stat()
|
|
if err != nil {
|
|
return block, err
|
|
}
|
|
block = (s.Mode() & os.ModeDevice) != 0
|
|
return block, nil
|
|
}
|
|
|
|
// CleanBMap cleans spaces out of bmaptools bmap file and updates SHA Hash
|
|
func CleanBMap(input string) error {
|
|
data, err := ioutil.ReadFile(input)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
re := regexp.MustCompile(`>\s*([\d\w-]+)\s*<`)
|
|
fixed := re.ReplaceAllString(string(data), ">${1}<")
|
|
b, err := LoadFromReader(bytes.NewReader([]byte(fixed)), false)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
b.BmapFileSHA1 = defaultBMAPHash
|
|
return b.Write(input)
|
|
|
|
}
|