borg/dedupestore/archiver.py

import os
import hashlib
import logging
import zlib
import cPickle
import argparse
import sys

from chunkifier import chunkify
from cache import Cache, NS_ARCHIVES, NS_CHUNKS
from bandstore import BandStore
from helpers import location_validator, pretty_size

CHUNK_SIZE = 55001

class LevelFilter(logging.Filter):

    def __init__(self, *args, **kwargs):
        super(LevelFilter, self).__init__(args, **kwargs)
        self.count = {}

    def filter(self, record):
        self.count.setdefault(record.levelname, 0)
        self.count[record.levelname] += 1
        return record


class Archive(object):

    def __init__(self, store, cache, name=None):
        self.store = store
        self.cache = cache
        self.items = []
        self.chunks = []
        self.chunk_idx = {}
        if name:
            self.open(name)

    def open(self, name):
        id = self.cache.archives[name]
        data = self.store.get(NS_ARCHIVES, id)
        if hashlib.sha256(data).digest() != id:
            raise Exception('Archive hash did not match')
        archive = cPickle.loads(zlib.decompress(data))
        self.items = archive['items']
        self.name = archive['name']
        self.chunks = archive['chunks']
        for i, (id, csize, osize) in enumerate(archive['chunks']):
            self.chunk_idx[i] = id

    def save(self, name):
        archive = {'name': name, 'items': self.items, 'chunks': self.chunks}
        data = zlib.compress(cPickle.dumps(archive))
        self.id = hashlib.sha256(data).digest()
        self.store.put(NS_ARCHIVES, self.id, data)
        self.store.commit()

    def add_chunk(self, id, csize, osize):
        try:
            return self.chunk_idx[id]
        except KeyError:
            idx = len(self.chunks)
            self.chunks.append((id, csize, osize))
            self.chunk_idx[id] = idx
            return idx

    def stats(self, cache):
        total_osize = 0
        total_csize = 0
        total_usize = 0
        chunk_count = {}
        for item in self.items:
            if item['type'] == 'FILE':
                total_osize += item['size']
                for idx in item['chunks']:
                    id = self.chunk_idx[idx]
                    chunk_count.setdefault(id, 0)
                    chunk_count[id] += 1
        for id, c in chunk_count.items():
            count, csize, osize = cache.chunkmap[id]
            total_csize += csize
            if  c == count:
                total_usize += csize
        return dict(osize=total_osize, csize=total_csize, usize=total_usize)

    def list(self):
        for item in self.items:
            print item['path']

    def extract(self, dest=None):
        dest = dest or os.getcwdu()
        for item in self.items:
            assert item['path'][0] not in ('/', '\\', ':')
            path = os.path.join(dest, item['path'])
            logging.info(path)
            if item['type'] == 'DIR':
                if not os.path.exists(path):
                    os.makedirs(path)
            if item['type'] == 'FILE':
                if not os.path.exists(os.path.dirname(path)):
                    os.makedirs(os.path.dirname(path))
                with open(path, 'wb') as fd:
                    for chunk in item['chunks']:
                        id = self.chunk_idx[chunk]
                        data = self.store.get(NS_CHUNKS, id)
                        cid = data[:32]
                        data = data[32:]
                        if hashlib.sha256(data).digest() != cid:
                            raise Exception('Invalid chunk checksum')
                        if hashlib.sha256(zlib.decompress(data)).digest() != id:
                            raise Exception('Invalid chunk checksum')
                        fd.write(zlib.decompress(data))

    def verify(self):
        for item in self.items:
            if item['type'] == 'FILE':
                for chunk in item['chunks']:
                    id = self.chunk_idx[chunk]
                    data = self.store.get(NS_CHUNKS, id)
                    data = self.store.get(NS_CHUNKS, id)
                    cid = data[:32]
                    data = data[32:]
                    if (hashlib.sha256(data).digest() != cid or
                        hashlib.sha256(zlib.decompress(data)).digest() != id):
                        logging.error('%s ... ERROR', item['path'])
                        break
                else:
                    logging.info('%s ... OK', item['path'])

    def delete(self, cache):
        self.store.delete(NS_ARCHIVES, self.cache.archives[self.name])
        for item in self.items:
            if item['type'] == 'FILE':
                for c in item['chunks']:
                    id = self.chunk_idx[c]
                    cache.chunk_decref(id)
        self.store.commit()
        del cache.archives[self.name]
        cache.save()

    def create(self, name, paths, cache):
        if name in cache.archives:
            raise NameError('Archive already exists')
        for path in paths:
            for root, dirs, files in os.walk(path):
                for d in dirs:
                    p = os.path.join(root, d)
                    self.items.append(self.process_dir(p, cache))
                for f in files:
                    p = os.path.join(root, f)
                    entry = self.process_file(p, cache)
                    if entry:
                        self.items.append(entry)
        self.save(name)
        cache.archives[name] = self.id
        cache.save()

    def process_dir(self, path, cache):
        path = path.lstrip('/\\:')
        logging.info(path)
        return {'type': 'DIR', 'path': path}

    def process_file(self, path, cache):
        try:
            fd = open(path, 'rb')
        except IOError, e:
            logging.error(e)
            return
        with fd:
            path = path.lstrip('/\\:')
            logging.info(path)
            chunks = []
            size = 0
            for chunk in chunkify(fd, CHUNK_SIZE, 30):
                size += len(chunk)
                chunks.append(self.add_chunk(*cache.add_chunk(chunk)))
        return {'type': 'FILE', 'path': path, 'chunks': chunks, 'size': size}


class Archiver(object):

    def open_store(self, location):
        store = BandStore(location.path)
        cache = Cache(store)
        return store, cache

    def exit_code_from_logger(self):
        if not self.level_filter.count.get('ERROR'):
            return 0
        else:
            return 1

    def do_create(self, args):
        store, cache = self.open_store(args.archive)
        archive = Archive(store, cache)
        archive.create(args.archive.archive, args.paths, cache)
        return self.exit_code_from_logger()

    def do_extract(self, args):
        store, cache = self.open_store(args.archive)
        archive = Archive(store, cache, args.archive.archive)
        archive.extract(args.dest)
        return self.exit_code_from_logger()

    def do_delete(self, args):
        store, cache = self.open_store(args.archive)
        archive = Archive(store, cache, args.archive.archive)
        archive.delete(cache)
        return self.exit_code_from_logger()

    def do_list(self, args):
        store, cache = self.open_store(args.src)
        if args.src.archive:
            archive = Archive(store, cache, args.src.archive)
            archive.list()
        else:
            for archive in sorted(cache.archives):
                print archive
        return self.exit_code_from_logger()

    def do_verify(self, args):
        store, cache = self.open_store(args.archive)
        archive = Archive(store, cache, args.archive.archive)
        archive.verify()
        return self.exit_code_from_logger()

    def do_info(self, args):
        store, cache = self.open_store(args.archive)
        archive = Archive(store, cache, args.archive.archive)
        stats = archive.stats(cache)
        print 'Original size:', pretty_size(stats['osize'])
        print 'Compressed size:', pretty_size(stats['csize'])
        print 'Unique data:', pretty_size(stats['usize'])
        return self.exit_code_from_logger()

    def run(self, args=None):
        parser = argparse.ArgumentParser(description='Dedupestore')
        parser.add_argument('-v', '--verbose', dest='verbose', action='store_true',
                            default=False,
                            help='Verbose output')

        subparsers = parser.add_subparsers(title='Available subcommands')
        subparser = subparsers.add_parser('create')
        subparser.set_defaults(func=self.do_create)
        subparser.add_argument('archive', metavar='ARCHIVE',
                               type=location_validator(archive=True),
                               help='Archive to create')
        subparser.add_argument('paths', metavar='PATH', nargs='+', type=str,
                               help='Paths to add to archive')

        subparser = subparsers.add_parser('extract')
        subparser.set_defaults(func=self.do_extract)
        subparser.add_argument('archive', metavar='ARCHIVE',
                               type=location_validator(archive=True),
                               help='Archive to create')
        subparser.add_argument('dest', metavar='DEST', type=str, nargs='?',
                               help='Where to extract files')

        subparser = subparsers.add_parser('delete')
        subparser.set_defaults(func=self.do_delete)
        subparser.add_argument('archive', metavar='ARCHIVE',
                               type=location_validator(archive=True),
                               help='Archive to delete')

        subparser = subparsers.add_parser('list')
        subparser.set_defaults(func=self.do_list)
        subparser.add_argument('src', metavar='SRC', type=location_validator(),
                               help='Store/Archive to list contents of')

        subparser= subparsers.add_parser('verify')
        subparser.set_defaults(func=self.do_verify)
        subparser.add_argument('archive', metavar='ARCHIVE',
                               type=location_validator(archive=True),
                               help='Archive to verity integrity of')

        subparser= subparsers.add_parser('info')
        subparser.set_defaults(func=self.do_info)
        subparser.add_argument('archive', metavar='ARCHIVE',
                               type=location_validator(archive=True),
                               help='Archive to display information about')

        args = parser.parse_args(args)
        if args.verbose:
            logging.basicConfig(level=logging.INFO, format='%(message)s')
        else:
            logging.basicConfig(level=logging.WARNING, format='%(message)s')
        self.level_filter = LevelFilter()
        logging.getLogger('').addFilter(self.level_filter)
        return args.func(args)

def main():
    archiver = Archiver()
    sys.exit(archiver.run())

if __name__ == '__main__':
    main()