bmap/bmap.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)
}