Using aes128 for encryption

[python-delta-tar] / rescue_tar.py

#!/usr/bin/env python

import argparse
import os
import tempfile
from functools import partial

from deltatar import tarfile
import filesplit

def rescue(tar_files, rescue_dir=None):
    '''
    Rescues a multivolume tarfile. Checks file name extension to detect
    format (compression, etc). Assumes it to be multivolume tar.
    '''
    # setup rescue_dir
    if isinstance(tar_files, basestring):
        tar_files = [tar_files]

    if not isinstance(tar_files, list):
        raise Exception("tar_files must be a list")

    for f in tar_files:
        if not isinstance(f, basestring):
            raise Exception("tar_files must be a list of strings")
        if not os.path.exists(f):
            raise Exception("tar file '%s' doesn't exist" % f)

    if rescue_dir is None:
        rescue_dir = os.path.dirname(tar_files[0])
    elif rescue_dir is None:
        rescue_dir = tempfile.mkdtemp()

    # autodetect file type by extension
    first_tar_file = tar_files[0]
    if first_tar_file.endswith(".tar.gz"):
        mode = "r#gz"
    elif first_tar_file.endswith(".tar"):
        mode = "r"

    base_name = os.path.basename(first_tar_file)
    extract_files = tar_files

    # num the number of files used in rescue mode. Used to name those files
    # when creating them. We put num in an object so that it can be referenced
    # instead of copied inside new_gz partial
    context = dict(num=0)

    # divide in compressed tar block files if it's r#gz
    if mode == "r#gz":
        extract_files = []
        # function used to create each chunk file
        def new_gz(context, extract_files, prefix, i):
            path = "%s.%d" %(prefix, context['num'])
            extract_files.append(path)
            context['num'] += 1
            return open(path, 'w')
        new_gz = partial(new_gz, context, extract_files)

        # split in compressed chunks
        for f in tar_files:
            filesplit.split_file('\x1f\x8b',
                os.path.join(rescue_dir, base_name), f, new_gz)

    # includes volumes already extracted with new_volume_handler
    already_extracted_vols = []

    def new_volume_handler(already_extracted_vols, next_num, tarobj, base_name, volume_number):
        '''
        Handles the new volumes when extracting
        '''

        # handle the special case where the first file is whatever.tar.gz and
        # the second is whatever.tar.gz.0
        base_name_split = base_name.split('.')
        next_num = 0
        try:
            next_num = int(base_name_split[-1]) + 1
            base_name = ".".join(base_name_split[:-1])
        except ValueError as e:
            pass

        volume_path = "%s.%d" % (base_name, next_num)
        already_extracted_vols.append(volume_path)
        tarobj.open_volume(volume_path)

    new_volume_handler = partial(new_volume_handler, already_extracted_vols)

    # extract files, as much as possible
    for f in extract_files:
        if f in already_extracted_vols:
            continue
        try:
            tarobj = tarfile.TarFile.open(f, mode=mode,
                new_volume_handler=new_volume_handler)
            tarobj.extractall()
            tarobj.close()
        except:
            pass

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--rescue_dir", help="directory where rescue files "
        "should be created. /tmp by default")
    parser.add_argument("tar_files", nargs="+", help="list of files of a "
        "multitar file to rescue. Assumes format first.extension "
        "second.extension.0 third.extension.1 ...")

    args = parser.parse_args()
    rescue(tar_files=args.tar_files, rescue_dir=args.rescue_dir)