nbd-chunk-import

#!/usr/bin/env python3
#
# Copyright (C) 2024 by Florian Wesch <fw@dividuum.de>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
import sys, os, binascii, time, argparse, tempfile, errno
from urllib.parse import quote
from tree_chunker import ChunkWriter
from hole_reader import HoleReader

class ChunkWriterDisk(ChunkWriter):
    def setup(self, prefix):
        self._prefix = prefix
        self._unique = 0
        self._reuse = 0
        try:
            os.makedirs(prefix)
        except OSError as err:
            if err.errno != errno.EEXIST:
                raise

    def persist_chunk(self, chunk_hash, chunk_data):
        target = os.path.join(self._prefix, binascii.hexlify(chunk_hash))
        if os.path.exists(target):
            self._reuse += len(chunk_data)
        else:
            with tempfile.NamedTemporaryFile(dir=self._prefix, prefix=b'.temp-') as f:
                f.write(chunk_data)
                os.link(f.name, target)
            self._unique += len(chunk_data)

    @property
    def unique(self):
        return self._unique

    @property
    def reuse(self):
        return self._reuse

    def flush(self, header_chunk_hash, unlock_key):
        print('{}#{}'.format(
            binascii.hexlify(header_chunk_hash).decode('utf8'),
            quote(unlock_key.decode('utf8')),
        ))

MB = 2**20

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog = sys.argv[0],
        description = "Imports a disk image file into a local chunk repository",
    )
    parser.add_argument('input',
        help = "Data to chunk. Most likely a partition image"
    )
    parser.add_argument('unlock_key',
        help = "Unlock key for this import."
    )
    parser.add_argument('repo_key',
        help = "Repository key for this repository."
    )
    parser.add_argument('target_dir',
        help = "Where to place the generated chunks. The directory will be created if it doesn't exists.",
    )
    parser.add_argument('-c', '--chunk_size',
        type = int,
        help = "Chunk size in 2^<chunk_size> bytes. 2^18 = 256KB recommended.",
        default = 18,
    )
    parser.add_argument('-b', '--block_size',
        type = int,
        help = "Block size of the resulting chunk image. (4096 recommended.",
        default = 4096,
    )
    parser.add_argument('-u', '--uncompressed',
        action = 'store_true',
        help = "Turn off compression (Not recommended)"
    )
    args = parser.parse_args()

    started = time.time()

    def print_stats(written, unique, reuse):
        compressed = unique + reuse
        per_s = written / (time.time() - started)
        compression_percent = 100 - 100 / written * compressed
        if compressed:
            reuse_percent = 100 / compressed * reuse
        else:
            reuse_percent = 0
        print('{:.1f}MB/s {:.0f}MB =[ {:.2f}% compression ]=> => {:.0f}MB =[ {:.2f}% chunk reuse ]=> {:.3f}MB written'.format(
            per_s/MB,
            written/MB,
            compression_percent,
            compressed/MB,
            reuse_percent,
            unique/MB,
        ), file=sys.stderr)

    writer = ChunkWriterDisk(
        args.chunk_size,
        args.repo_key.encode('utf8'),
        args.unlock_key.encode('utf8'),
        not args.uncompressed,
    )
    writer.setup(args.target_dir.encode('utf8'))

    last_stats = time.time()
    written = 0
    stream = HoleReader('/dev/stdin' if args.input == '-' else args.input)
    for section in stream:
        if section.all_zero:
            skip = section.skip()
            writer.write_zeros(skip)
            written += skip
        else:
            while 1:
                buf = section.read(2**args.chunk_size)
                if not buf:
                    break
                writer.write(buf)
                written += len(buf)
                if time.time() > last_stats + 0.5:
                    last_stats = time.time()
                    print_stats(written, writer.unique, writer.reuse)
    print_stats(written, writer.unique, writer.reuse)

    print('', file=sys.stderr)
    writer.wrap_up(block_size=args.block_size)