Subversion Repositories programming

#!/usr/bin/env python

# Copyright: Ira W. Snyder (devel@irasnyder.com)
# Start Date: 2005-10-13
# End Date:
# License: GNU General Public License v2 (or at your option, any later version)
#
# Changelog Follows:
# - 2005-10-13
# - Added get_par2_filenames() to parse par2 files
# - Added the parset object to represent each parset.
#
# - 2005-10-14
# - Finished the parset object. It will now verify and extract parsets.
# - Small changes to the parset object. This makes the parjoin part
#   much more reliable.
# - Added the OptionParser to make this nice to run at the command line.
# - Made recursiveness an option.
# - Made start directory an option.
# - Check for appropriate programs before starting.
#
# - 2005-10-17
# - Use a regular expression to handle the deletable types.
#
# - 2005-10-18
# - Use regular expressions to handle all finding of files, instead of
#   using the glob module.
# - Add a config class to handle all the default config stuff sanely.
#   This makes it easier to change some of the main parts of the program to
#   your specific configuration.
# - Move the docrcchecks variable inside the get_par2_filenames() function,
#   which is where it belongs anyway.
# - Added command-line option to check for required programs at start.
#
# - 2005-10-20
# - Added a config option to extract with full path.
#

################################################################################
# REQUIREMENTS:
#
# This code requires the programs cfv, par2repair, lxsplit, and rar to be able
# to function properly. I will attempt to check that these are in your path.
################################################################################

class rarslave_config:
    """A simple class to hold the default configs for the whole program"""

    WORK_DIR               = '~/downloads/usenet'
    DELETEABLE_TYPES_REGEX = '^.*\.(rar|r\d\d)$'
    TEMP_REPAIR_REGEX      = '.*\.1$'
    PAR2_REGEX             = '.*\.(PAR|par)2$'
    VIDEO_FILE_REGEX       = '.*\.(AVI|avi|OGM|ogm|MKV|mkv)$'
    RECURSIVE              = True
    CHECK_REQ_PROGS        = False
    EXTRACT_WITH_FULL_PATH = False

    def __init__(self):
        pass

config = rarslave_config()
        
################################################################################
# The PAR2 Parser
#
# This was stolen from cfv (see http://cfv.sourceforge.net/ for a copy)
################################################################################

import struct, errno

def chompnulls(line):
    p = line.find('\0')
    if p < 0: return line
    else:     return line[:p]

def get_par2_filenames(filename):
    """Get all of the filenames that are protected by the par2
    file given as the filename"""

    try:
        file = open(filename, 'rb')
    except:
        print 'Could not open %s' % (filename, )
        return []

    # We always want to do crc checks
    docrcchecks = True

    pkt_header_fmt = '< 8s Q 16s 16s 16s'
    pkt_header_size = struct.calcsize(pkt_header_fmt)
    file_pkt_fmt = '< 16s 16s 16s Q'
    file_pkt_size = struct.calcsize(file_pkt_fmt)
    main_pkt_fmt = '< Q I'
    main_pkt_size = struct.calcsize(main_pkt_fmt)

    seen_file_ids = {}
    expected_file_ids = None
    filenames = []

    while 1:
        d = file.read(pkt_header_size)
        if not d:
            break

        magic, pkt_len, pkt_md5, set_id, pkt_type = struct.unpack(pkt_header_fmt, d)

        if docrcchecks:
            import md5
            control_md5 = md5.new()
            control_md5.update(d[0x20:])
            d = file.read(pkt_len - pkt_header_size)
            control_md5.update(d)

            if control_md5.digest() != pkt_md5:
                raise EnvironmentError, (errno.EINVAL, \
                    "corrupt par2 file - bad packet hash")

        if pkt_type == 'PAR 2.0\0FileDesc':
            if not docrcchecks:
                d = file.read(pkt_len - pkt_header_size)

            file_id, file_md5, file_md5_16k, file_size = \
                struct.unpack(file_pkt_fmt, d[:file_pkt_size])

            if seen_file_ids.get(file_id) is None:
                seen_file_ids[file_id] = 1
                filename = chompnulls(d[file_pkt_size:])
                filenames.append(filename)

        elif pkt_type == "PAR 2.0\0Main\0\0\0\0":
            if not docrcchecks:
                d = file.read(pkt_len - pkt_header_size)

            if expected_file_ids is None:
                expected_file_ids = []
                slice_size, num_files = struct.unpack(main_pkt_fmt, d[:main_pkt_size])
                num_nonrecovery = (len(d)-main_pkt_size)/16 - num_files

                for i in range(main_pkt_size,main_pkt_size+(num_files+num_nonrecovery)*16,16):
                    expected_file_ids.append(d[i:i+16])

        else:
            if not docrcchecks:
                file.seek(pkt_len - pkt_header_size, 1)

    if expected_file_ids is None:
        raise EnvironmentError, (errno.EINVAL, \
            "corrupt or unsupported par2 file - no main packet found")

    for id in expected_file_ids:
        if not seen_file_ids.has_key(id):
            raise EnvironmentError, (errno.EINVAL, \
                "corrupt or unsupported par2 file - " \
                "expected file description packet not found")

    return filenames

################################################################################
# The parset object
#
# This is an object based representation of a parset, and will verify itself
# and extract itself, if possible.
################################################################################

import os, glob, re

class parset:
    def __init__(self, par_filename):
        self.parfile = par_filename
        self.extra_pars = []
        self.files = False
        self.used_parjoin = False
        self.verified = False
        self.extracted = False

    def get_filenames(self):
        return get_par2_filenames(parfile)

    def all_there(self):
        """Check if all the files for the parset are present.
        This will help us decide which par2 checker to use first"""
        for f in self.files:
            if not os.path.isfile(f):
                return False

        # The files were all there
        return True

    def verify(self):
        """This will verify the parset by the most efficient method first,
        and then move to a slower method if that one fails"""

        retval = False #not verified yet

        # if all the files are there, try verifying fast
        if self.all_there():
            retval = self.__fast_verify()

            if retval == False:
                # Failed to verify fast, so try it slow, maybe it needs repair
                retval = self.__slow_verify()

        # If we've got a video file, maybe we should try to parjoin it
        elif self.__has_video_file():
            retval = self.__parjoin()

        else: #not all there, maybe we can slow-repair
            retval = self.__slow_verify()

        self.verified = retval
        return self.verified

    def __fast_verify(self):
        retval = os.system('cfv -v -f "%s"' % (self.parfile, ))

        if retval == 0:
            return True #success

        return False #failure

    def __slow_verify(self):
        retval = os.system('par2repair "%s"' % (self.parfile, ))

        if retval == 0:
            return True #success

        return False #failure

    def __parjoin(self):
        retval = os.system('lxsplit -j "%s.001"' % (self.files[0], ))

        retval = self.__fast_verify()

        if retval == False:
            # Failed to verify fast, so try it slow, maybe it needs repair
            retval = self.__slow_verify()

        if retval == False: # failed to verify, so remove the lxsplit created file
            os.remove(self.files[0])

        self.used_parjoin = retval
        self.verified = retval
        return self.verified

    def __has_video_file(self):
        regex = re.compile(config.VIDEO_FILE_REGEX)
    
        for f in self.files:
            if regex.match(f):
                return True

        return False

    def __remove_currentset(self):
        """Remove all of the files that are extractable, as well as the pars.
        Leave everything else alone"""

        if not self.extracted:
            print 'Did not extract yet, not removing currentset'
            return

        # remove the main par
        os.remove(self.parfile)

        # remove all of the extra pars
        for i in self.extra_pars:
            os.remove(i)

        # remove any rars that are associated (leave EVERYTHING else)
        # This regex matches both old and new style rar(s) by default.
        regex = re.compile(config.DELETEABLE_TYPES_REGEX)

        for i in self.files:
            if regex.match(i):
                os.remove(i)

        # remove any .{001,002,...} files (from parjoin)
        if self.used_parjoin:
            for i in os.listdir(os.getcwd()):
                if i != self.files[0] and self.files[0] in i:
                    os.remove(i)

        # remove any temp repair files
        regex = re.compile(config.TEMP_REPAIR_REGEX)
        [os.remove(f) for f in os.listdir(os.getcwd()) if regex.match(f)]

    def __get_extract_file(self):
        """Find the first extractable file"""
        for i in self.files:
            if os.path.splitext(i)[1] == '.rar':
                return i

        return None

    def extract(self):
        """Attempt to extract all of the files related to this parset"""
        if not self.verified:
            self.extracted = False
            print 'Not (successfully) verified, not extracting'
            return False #failed to extract

        extract_file = self.__get_extract_file()

        if extract_file != None:
            if config.EXTRACT_WITH_FULL_PATH:
                retval = os.system('rar x -o+ "%s"' % (extract_file, ))
            else:
                retval = os.system('rar e -o+ "%s"' % (extract_file, ))

            if retval != 0:
                print 'Failed to extract'
                self.extracted = False
                return self.extracted

        # we extracted ok, so remove the currentset
        self.extracted = True
        self.__remove_currentset()

        return self.extracted


################################################################################
# The rarslave program itself
################################################################################

import os, sys
from optparse import OptionParser

def check_required_progs():
    """Check if the required programs are installed"""

    shell_not_found = 32512
    needed = []

    if os.system('cfv --help > /dev/null 2>&1') == shell_not_found:
        needed.append('cfv')

    if os.system('par2repair --help > /dev/null 2>&1') == shell_not_found:
        needed.append('par2repair')

    if os.system('lxsplit --help > /dev/null 2>&1') == shell_not_found:
        needed.append('lxpsplit')

    if os.system('rar --help > /dev/null 2>&1') == shell_not_found:
        needed.append('rar')

    if needed:
        for n in needed:
            print 'Needed program "%s" not found in $PATH' % (n, )

        sys.exit(1)

def get_parsets():
    """Get a representation of each parset in the current directory, and
    return them as a list of parset instances"""

    regex = re.compile(config.PAR2_REGEX)
    par2files = [f for f in os.listdir(os.getcwd()) if regex.match(f)]

    parsets = []

    for i in par2files:
        try:
            filenames = get_par2_filenames(i)
            create_new = True
        except EnvironmentError:
            print 'CORRUPT PARFILE: %s' % (i, )
            continue

        # if we already have an instance for this set, append
        # this par file to the extra_pars field
        for j in parsets:
            if j.files == filenames:
                j.extra_pars.append(i)
                create_new = False

        # we haven't seen this set yet, so we'll create it now
        if create_new == True:
            cur = parset(i)
            cur.files = filenames
            parsets.append(cur)

    return parsets

def directory_worker(dir):
    """Attempts to find, verify, and extract every parset in the directory
    given as a parameter"""

    cwd = os.getcwd()
    os.chdir(dir)

    parsets = get_parsets()

    # Verify each parset
    for p in parsets:
        p.verify()

    # Attempt to extract each parset
    for p in parsets:
        p.extract()

    os.chdir(cwd)

def main():

    # Build the OptionParser
    parser = OptionParser()
    parser.add_option('-n', '--not-recursive',
                      action='store_false', dest='recursive',
                      default=config.RECURSIVE, help="Don't run recursively")
    
    parser.add_option('-d', '--work-dir',
                      dest='work_dir', default=config.WORK_DIR,
                      help="Start running at DIR", metavar='DIR')
    
    parser.add_option('-p', '--check-required-programs',
                       action='store_true', dest='check_progs',
                       default=config.CHECK_REQ_PROGS, help="Don't check for required programs")

    # Parse the given options
    (options, args) = parser.parse_args()

    # Fix up the working directory
    options.work_dir = os.path.abspath(os.path.expanduser(options.work_dir))

    # Check that we have the required programs installed
    if options.check_progs:
        check_required_progs()

    # Run rarslave!
    if options.recursive:
        for root, dirs, files in os.walk(options.work_dir):
            directory_worker(root)
    else:
        directory_worker(options.work_dir)

if __name__ == '__main__':
    main()