source: misc/filearchive/filearchive @ 41:9753fcb9f750

#!/usr/bin/env python
#
# filearchive
# A simple file archiver with encryption possibilities
# For encryption it relies on 'mcrypt' being installed
#
# Copyright (C) 2008 Daniele Tricoli aka Eriol <eriol@mornie.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

__version__ = '0.1'

import getpass
import md5
import optparse
import os
import subprocess
import sys
import zipfile

SETTINGS = {}

def md5sum(f):
    m = md5.new()
    while True:
        block = f.read(8096)
        if not block:
            break
        m.update(block)
    return m.hexdigest()

def process_command_line(argv):
    if argv is None:
        argv = sys.argv[1:]

    usage = 'usage: %prog [options] file_dir'
    version = '%%prog %s' % __version__
    parser = optparse.OptionParser(usage, version=version)

    parser.add_option('-e', '--encrypt', action='store_true',
                      default=False, dest='encrypt',
                      help='encrypt archives using mcrypt')

    parser.add_option('-o', '--output', action='store',
                      default=False, dest='output_directory',
                      help='''put archives in OUTPUT_DIRECTORY. If not set
                      OUTPUT_DIRECTORY=tmp/''')

    parser.add_option('-s', '--size', action='store', type='int',
                      default=False, dest='size_archive',
                      help='max archives size')

    settings, args = parser.parse_args(argv)

    if len(args) != 1:
        parser.error('You must specify a directory with files to archive')

    for opt, value in settings.__dict__.items():
        SETTINGS[opt] = value

    file_dir = args[0]
    if file_dir[-1] == '/': # remove trailing /
        file_dir = file_dir[:-1]

    SETTINGS['file_dir'] = file_dir

def list_archives():
    files_list = [[]]
    mb_counter = 0
    archive_counter = 0
    if not SETTINGS['size_archive']:
        size_archive = 18 # 18 MB
    else:
        size_archive = SETTINGS['size_archive']

    size_archive_bytes = size_archive * 1024**2

    try:
        ld = os.listdir(SETTINGS['file_dir'])
    except:
        print "Can't open directory: %s" % SETTINGS['file_dir']
        sys.exit(1)

    files = sorted(ld)

    for f in files:
        file_path = os.path.join(SETTINGS['file_dir'], f)
        size = os.stat(file_path)[6]
        mb_counter += size

        if mb_counter > size_archive_bytes:
            files_list.append([])
            archive_counter += 1
            mb_counter = size

        files_list[archive_counter].append(file_path)

    return files_list

def create_archives(file_list):
    archive_name = os.path.basename(SETTINGS['file_dir'])

    if not SETTINGS['output_directory']:
        output_directory = 'tmp/'
    else:
        output_directory = SETTINGS['output_directory']

    if not os.path.exists(output_directory):
        try:
            os.mkdir(output_directory)
        except OSError, data:
            print 'Error:', data.strerror
            sys.exit(1)

    for i, item in enumerate(file_list):
        zip_name = archive_name + '_' + str(i) + '.zip'
        zip_archive = os.path.join(output_directory, zip_name)
        print 'Creating... %s' % zip_archive
        z = zipfile.ZipFile(zip_archive, 'w')
        for f in item:
            z.write(f)

    if SETTINGS['encrypt']:
        archive_ext = '.nc'

        password = getpass.getpass()
        path = os.path.join(output_directory, '*.zip')
        command = 'mcrypt %s -k %s -u' % (path, password)

        subprocess.call(command, shell=True)
    else:
        archive_ext = '.zip'

    print 'Generating md5sum...'
    md5name = os.path.join(output_directory, archive_name + '.md5sum')
    try:
        md5file = open(md5name, 'w')

        l = os.listdir(output_directory)
        for f in sorted(l):
            name, ext = os.path.splitext(f)
            name = name[:len(archive_name)]
            if name == archive_name and ext == archive_ext:
                file_path = os.path.join(output_directory, f)
                fi = open(file_path)
                m = md5sum(fi)
                md5file.write('%s  %s\n' % (m, f))
    finally:
        md5file.close()

def main(argv=None):
    process_command_line(argv)
    photos = list_archives()
    create_archives(photos)

    print 'Archiving complete!'

    return 0

if __name__ == '__main__':
    status = main()
    sys.exit(status)