Subversion Repositories programming

#!/usr/bin/env python

# Copyright (c) 2005,2006 Ira W. Snyder (devel@irasnyder.com)

# 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.

#

# - 2005-10-22

# - Re-wrote the config class so that there is a config file, which

#   resides at ~/.config/rarslave/rarslave.conf by default.

# - Added the command-line option -c to write out an updated version

#   of the config file (to fill in any missing options with the defaults)

# - Added the command-line option -f to write out a new default config file,

#   which overwrites any user changes.

# - Made all regexes case insensitive.

# - Made all command-line options override the config file.

#

# - 2005-10-30

# - Added the '-o' option, to output debugging info. Hopefully next time

#   someone finds a bug, they can output this and send it to me with a

#   description of the bug they're seeing.

#

# - 2005-11-05

# - Added an output system to rarslave. This makes a nice status report

#   possible at the end of the program run.

#

# - 2005-11-06

# - Fixed the rar command so that it can extract files whose names begin

#   with a hyphen.

#

# - 2006-03-08

# - Make an interactive mode which asks the user before deleting files.

# - Make multi-file parjoin work.

#

################################################################################

# 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.

################################################################################

import ConfigParser, os

class rarslave_config:

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

    def __read_config(self, filename='~/.config/rarslave/rarslave.conf'):

        """Attempt to open and read the rarslave config file"""

        # Make sure the filename is corrected

        filename = os.path.abspath(os.path.expanduser(filename))

        user_config = {}

        # Write the default config if it doesn't exist

        if not os.path.isfile(filename):

            self.write_config(default=True)

        config = ConfigParser.ConfigParser()

        config.read(filename)

        for section in config.sections():

            for option in config.options(section):

                user_config[(section, option)] = config.get(section, option)

        return user_config

    def write_config(self, filename='~/.config/rarslave/rarslave.conf', default=False):

        """Write out the current config to the config file. If you set default=True, then

        the default config file will be written."""

        config = ConfigParser.ConfigParser()

        # Correct filename

        filename = os.path.abspath(os.path.expanduser(filename))

        # Reset all config to make sure we write the default one, if necessary

        if default:

            self.__user_config = {}

            print 'Writing default config to %s' % (filename, )

        # [directories] section

        config.add_section('directories')

        for (s, k) in self.__defaults.keys():

            if s == 'directories':

                config.set(s, k, self.get_value(s, k))

        # [options] section

        config.add_section('options')

        for (s, k) in self.__defaults.keys():

            if s == 'options':

                config.set(s, k, self.get_value(s, k))

        # [regular_expressions] section

        config.add_section('regular expressions')

        for (s, k) in self.__defaults.keys():

            if s == 'regular expressions':

                config.set(s, k, self.get_value(s, k))

        # Try to make the ~/.config/rarslave/ directory

        if not os.path.isdir(os.path.split(filename)[0]):

            try:

                os.makedirs(os.path.split(filename)[0])

            except:

                print 'Could not make directory: %s' % (os.path.split(filename)[0], )

                sys.exit()

        # Try to write the config file to disk

        try:

            fsock = open(filename, 'w')

            try:

                config.write(fsock)

            finally:

                fsock.close()

        except:

            print 'Could not open: %s for writing' % (filename, )

            sys.exit()

    def __get_default_val(self, section, key):

        return self.__defaults[(section, key)]

    def get_value(self, section, key):

        """Get a config value. Attempts to get the value from the user's

        config first, and then uses the default."""

        try:

            value = self.__user_config[(section, key)]

        except:

            # This should work, unless you write something stupid

            # into the code, so DON'T DO IT

            value = self.__get_default_val(section, key)

        # Convert config options to booleans for easier use

        if value == 'True':

            value = True

        if value == 'False':

            value = False

        return value

    def __init__(self):

        self.__defaults = {

            ('directories', 'working_directory') : '~/downloads/usenet',

            ('options', 'recursive') : True,

            ('options', 'check_required_programs') : False,

            ('options', 'extract_with_full_path') : False,

            ('options', 'interactive') : False,

            ('regular expressions', 'par2_regex') : '.*\.par2$',

            ('regular expressions', 'video_file_regex') : '.*\.(avi|ogm|mkv|mp4)$',

            ('regular expressions', 'temp_repair_regex') : '.*\.1$',

            ('regular expressions', 'remove_regex') : '^.*\.(rar|r\d\d)$' }

        self.__user_config = self.__read_config()

# This is the global config variable.

config = rarslave_config()

# This is the global options variable. (to be set later)

options = None

################################################################################

# The rarslave_output class

#

# This class handles the nice output summary which is printed at the end

# of a run

################################################################################

class rarslave_output:

    # Data structure: list of lists

    # [ [status, filename], ... ]

    #

    # Where status is one of:

    # 0: Verified and Extracted Perfectly

    # 1: Failed to Verify (and therefore Extract)

    # 2: Verified correctly, but failed to Extract

    #

    def __init__(self):

        self.output_list    = []

        self.good_files     = 0

        self.unverified     = 0

        self.unextractable  = 0

        self.corrupt_par2   = 0

    def print_equal_line(self, size=80):

        """Print an 80 character line of equal signs"""

        str = ''

        for i in range(size):

            str += '='

        print str

    def print_results_table(self):

        """Print a nice table of the results from this run"""

        # Print the table of good files (if we have any)

        if self.good_files > 0:

            print

            self.print_equal_line()

            print 'Files that were extracted perfectly'

            self.print_equal_line()

            for entry in self.output_list:

                if entry[0] == 0:

                    print '%s' % (entry[1], )

        # Print the table of unverified files (if we have any)

        if self.unverified > 0:

            print

            self.print_equal_line()

            print 'Files that failed to verify (and extract)'

            self.print_equal_line()

            for entry in self.output_list:

                if entry[0] == 1:

                    print '%s' % (entry[1], )

        # Print the table of unextracted files (if we have any)

        if self.unextractable > 0:

            print

            self.print_equal_line()

            print 'Files that were verified, but failed to extract'

            self.print_equal_line()

            for entry in self.output_list:

                if entry[0] == 2:

                    print '%s' % (entry[1], )

        # Print the table of corrupt PAR2 files (if we have any)

        if self.corrupt_par2 > 0:

            print

            self.print_equal_line()

            print 'Files that had corrupt par2 files'

            self.print_equal_line()

            for entry in self.output_list:

                if entry[0] == 3:

                    print '%s' % (entry[1], )

        # Print a blank line at the end

        print

    def add_file(self, status, filename):

        if status == 0:

            self.good_files += 1

        elif status == 1:

            self.unverified += 1

        elif status == 2:

            self.unextractable += 1

        elif status == 3:

            self.corrupt_par2 += 1

        else:

            # We have a bad value, so raise a ValueError

            raise ValueError

        self.output_list.append([status, filename])

# This is the global output variable

output = rarslave_output()

################################################################################

# 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 print_debug_info(self):

        """Special function for debugging"""

        print '========== DEBUG INFO STARTS HERE =========='

        print '=== parfile ==='

        print self.parfile

        print

        print '=== extra_pars ==='

        for f in self.extra_pars:

            print f

        print

        print '=== files ==='

        for f in self.files:

            print f

        print '=========== DEBUG INFO ENDS HERE ==========='

    def get_filenames(self):

        return get_par2_filenames(self.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 video files, maybe we should try to parjoin them

        elif self.__number_of_video_files() > 0:

            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):

        # Join all of the files

        for f in self.files:

            if self.__is_video_file(f):

                self.__parjoin_single(f)

        # Re-verify the parset

        retval = self.__fast_verify()

        if retval == False:

            # Failed to verify fast, so try it slow, maybe it needs repair

            retval = self.__slow_verify()

        # Failed to verify, so remove all of the lxsplit created files

        if retval == False:

            for f in self.files:

                if self.__is_video_file(f):

                    try:

                        os.remove(f)

                    except OSError:

                        print 'Failed to remove file: %s' % (f, )

        self.used_parjoin = retval

        self.verified = retval

        return self.verified

    def __parjoin_single(self, base_file):

        parjoinable = False

        regex = re.compile("^%s\.001$" % (re.escape(base_file), ), re.IGNORECASE)

        for f in os.listdir(os.getcwd()):

            if regex.match(f):

                parjoinable = True

                break

        if parjoinable:

            retval = os.system('lxsplit -j "%s.001"' % (base_file, ))

            return retval

        return False

    def __is_video_file(self, file):

        regex = re.compile(

                config.get_value('regular expressions', 'video_file_regex'),

                re.IGNORECASE)

        if regex.match(file):

            return True

        return False

    def __number_of_video_files(self):

        count = 0

        for f in self.files:

            if self.__is_video_file(f):

                count += 1

        return count

    def __remove_split_files(self, base_name, files_to_remove):

        """Store base_name.{001,002, ... , 999} files in the list

        $files_to_remove"""

        base_name_regex = "^%s\.\d\d\d$" % (re.escape(base_name), )

        regex = re.compile(base_name_regex, re.IGNORECASE)

        for f in os.listdir(os.getcwd()):

            if regex.match(f):

                files_to_remove.append(f)

    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

        files_to_remove = []

        # remove the main par

        files_to_remove.append(self.parfile)

        # remove all of the extra pars

        for i in self.extra_pars:

            files_to_remove.append(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.get_value('regular expressions', 'remove_regex'),

                re.IGNORECASE)

        for i in self.files:

            if regex.match(i):

                files_to_remove.append(i)

        # remove any .{001,002,...} files (from parjoin)

        if self.used_parjoin:

            for f in self.files:

                self.__remove_split_files(f, files_to_remove)

        # remove any temp repair files

        regex = re.compile(

                config.get_value('regular expressions', 'temp_repair_regex'),

                re.IGNORECASE)

        [files_to_remove.append(f) for f in os.listdir(os.getcwd()) if regex.match(f)]

        # interactively remove files

        if options.interactive:

            print # blank line

            for f in files_to_remove:

                print f

            print '========================================'

            done = False

            while not done:

                s = raw_input("Delete files [y,n]: ")

                s.lower()

                if s == 'y' or s == 'yes':

                    done = True

                    self.__remove_list_of_files(files_to_remove)

                elif s == 'n' or s == 'no':

                    done = True

                    print 'Not removing files'

                else:

                    print 'Bad selection, try again...'

        else:

            self.__remove_list_of_files(files_to_remove)

    def __remove_list_of_files(self, files_to_remove):

        """Remove all files in the list"""

        # remove duplicates from the list

        temp = []

        for f in files_to_remove:

            if f not in temp:

                temp.append(f)

        files_to_remove = temp

        # remove the files

        for f in files_to_remove:

            try:

                os.remove(f)

            except OSError:

                print 'WW: Problem deleting: %s' % 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

            output.add_file(1, self.parfile)

            return False #failed to extract

        extract_file = self.__get_extract_file()

        if extract_file != None:

            if config.get_value('options', '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:

                output.add_file(2, self.parfile)

                self.extracted = False

                return self.extracted

        # we extracted ok, so remove the currentset

        self.extracted = True

        self.__remove_currentset()

        output.add_file(0, self.parfile)

        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.get_value('regular expressions', 'par2_regex'),

            re.IGNORECASE)

    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:

            output.add_file(3, 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()

    # Print debug info if we're supposed to

    if options.debug_info:

        for p in parsets:

            p.print_debug_info()

    # No debug info

    else:

        # 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.get_value('options', 'recursive'),

                      help="Don't run recursively")

    parser.add_option('-d', '--work-dir',

                      dest='work_dir',

                      default=config.get_value('directories', 'working_directory'),

                      help="Start running at DIR", metavar='DIR')

    parser.add_option('-p', '--check-required-programs',

                       action='store_true', dest='check_progs',

                       default=config.get_value('options', 'check_required_programs'),

                       help="Check for required programs")

    parser.add_option('-f', '--write-default-config',

                      action='store_true', dest='write_def_config',

                      default=False, help="Write out a new default config")

    parser.add_option('-c', '--write-new-config',

                      action='store_true', dest='write_config',

                      default=False, help="Write out the current config")

    parser.add_option('-o', '--output-debug-info',

                       action='store_true', dest='debug_info',

                       default=False,

                       help="Output debug info for every parset, then exit")

    parser.add_option('-i', '--interactive', dest='interactive', action='store_true',

                      default=config.get_value('options', 'interactive'),

                      help="Confirm before removing files")

    # Parse the given options

    global 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()

    # Write out a new default config, if we need it

    if options.write_def_config:

        config.write_config(default=True)

    # Write out the current config (adds new options to an existing config)

    if options.write_config:

        config.write_config()

    # Run rarslave!

    if options.recursive:

        for root, dirs, files in os.walk(options.work_dir):

            directory_worker(root)

    else:

        directory_worker(options.work_dir)

    # Print the results

    output.print_results_table()

if __name__ == '__main__':

    main()