arf

#!/usr/bin/env python2
#####################################################################
"""
A.R.F. (Archive Retrieval Fiend)
--------------------------------------------------------------------
Concept:
    This script exists for two reasons:
        1) Lazy people don't want to type long commands like "tar zxvf" and
	   "unzip" and "unrar", or their respective commandline parameters
        2) Downloading a tarball, extracting it, then deleting it
           is a wasteful and boring process

    This script remedies that situation. Just run "arf (filename or URL)"
    and it'll figure out how to extract it, and put the files in a nicely named
    directory. It works on http, https, and FTP urls.
    
    The directory creation works like this:
       + if the archive contains one directory, extract it as is
       + if the archive contains more than that, make a directory named
	  after the filename with the archive extension removed.

Examples:
    arf http://prdownloads.sourceforge.net/something/mrprogram-1.2.45.6.3.tar.bz2
    arf warez/Awesome.Game.RZR.FLT.DEViANCE/rzrflddev-ag.rar
    arf pile_of_junk.tgz /root

todo:
    + better implementation
        - FileExtractor takes a 'thing' in its __init__ (throwing an exception if it can't handle it), then
          returns an instance that allows the user to '.extract(destination)' the archive.
    + list the contents of the archive 
      - arf l <archive name> will list an archive contents (piped to less)
      - when <archive name> is an url, give the option to extract it afterwards

"""

#####################################################################
## Desired features #################################################
#####################################################################

# == installs packages automatically
# arf something.rar
#   -- rar isn't installed -> prompts to install via apt-get or yum
# arf -i
#   -- install all supported packages

# == list archive contents
# arf -v archive.tar.gz <other archives, including .pk3 and .wad and .cab>
#   -- show contents (colour-coded, piped to less)

# == streaming extraction (uncompress while downloading) 
# arf http://whatever/blah.zip
#   -- supported by all the popular ones: .tar.{gz,bz2}, .zip, .rar

# == stream format detection (format not always obvious from url)
# arf http://site/getazip #=> streams a zip file
#   -- read the first 4096 bytes of the file, and use that to detect the format.  

# == handle weird redirects
# arf http://sf.net/projects/foo/blah/thing.tgz?mirror=what&cookie=324983294&etc.
#   -- if last redirect has a proper filename, use that as archive name.
#      otherwise:
#           1) look at path & guess
#           2) prompt user

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

__version__ = '0.5.3'

import os, sys, re, urllib, time, tempfile
from os.path import join, isfile, isdir, exists, split

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

def madProps():
    s = "  arf! (Archive Retriever, Fo'sho') v%s :: by Chris Gahan (chris@ill-logic.com)  " % __version__
    print '=' * len(s)
    print s
    print '=' * len(s)
    print

def howToUseMe():
    print 'This program takes a compressed file from an url or a filename, figures out'
    print 'what type of file it is, then tries to extract it into a directory.'
    print
    exts = []
    for extractor in extractors:
        exts.extend(extractor.exts)
    print 'Supported filetypes:\n    %s' % ', '.join(exts)
    print
    print 'Usage:'
    print '    %s <input file or url> [output directory]' % sys.argv[0]
    print
    print 'The [output directory] is optional. By default, the current directory is used.'
    print
    
def execute(cmd):
    error = os.system(cmd)
    if error != 0:
        raise Exception, "Error! System command returned errorcode %d.\nCommand was: '%s'" % (error, cmd)
        
def prompt(message='(Y/n)? ', choices=['y','n'], default='y'):
    inp = None
    while not (inp in choices) and inp != '':
        inp = raw_input(message).lower()
    
    if inp == '':
        inp = default

    return default
            
def isurl(thing):
    return re.match(r'^(https?|ftp|git)://', thing) is not None

def isoption(arg):
    if arg[0] != '-' or os.path.exists(arg):
        return False
    else:
        return True

def only_dir_in_path(dir):
    entries = os.listdir(dir)
    dirs = []
    for entry in entries:
        fullentry = join(dir,entry)
        if isfile(fullentry):
            return None
        if isdir(fullentry):
            dirs.append(entry)
        if len(dirs) > 1:
            return None

    return dirs[0]

def is_exe(fpath):
    return os.path.exists(fpath) and os.access(fpath, os.X_OK)
    
def which(program):
    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

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

class URLGrabber(urllib.FancyURLopener):
    lastdisplaystring = ''

    def reporthook(self, blockcount, blocksize, filesize):
        """ Print the current transfer statistics """
    
        bytes = blockcount * blocksize
        if bytes > filesize:
            bytes = filesize
    
        percent = (bytes*100.0)//filesize
        now = time.time()
        if now - self.starttime == 0.0:
            now += 0.0001
        speed = (bytes / (now - self.starttime)) / 1024
    
        lastlen = len(self.lastdisplaystring)
        #dmesg('lastline = %d' % lastlen)
        print '\x08' * (lastlen+2),
        nextline = '%0.1f%% complete (%d/%d bytes) - %0.2f k/sec' % (percent, bytes, filesize, speed)
        #dmesg('nextline = %d' % len(nextline))
        if len(nextline) < lastlen:
            nextline += ' '*(lastlen - len(nextline))
    
        #dmesg('newnextline = %d' % len(nextline))
        print nextline,
    
        self.lastdisplaystring = nextline
    
    def retrieve(self, url):
        def reporthookwrapper(blockcount, blocksize, filesize):
            self.reporthook(blockcount, blocksize, filesize)
        
        self.lastdisplaystring = ''
        self.starttime = time.time()
        return urllib.FancyURLopener.retrieve(self, url, reporthook=reporthookwrapper)

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

class FileExtractor:
    """ Mother class!! """
    
    exts            = []
    streamable      = False
    binary          = None
    package         = None

    def canHandle(self, filename):
        """ Can you handle the file? CAN YOU HANDLE THE FILE?!!? """
        filename = filename.lower()
        for ext in self.exts:
            if filename.endswith(ext):
                return True
        return False

    def dirname_from_archivename(self, archivename):
        lowername = archivename.lower()
        for ext in self.exts:
            if lowername.endswith(ext):
                return archivename[0:-len(ext)]
        return None

    def find_binary(self):
        if isinstance(self.binary, list):
            binaries = self.binary
        else:
            binaries = [self.binary]
            
        for binary in binaries:
            if which(binary) is not None:
                return binary
        
        return None
    

    def check_binary(self):
        binary = self.find_binary()
        if binary:
            return binary
        else:
            answer = prompt("* %s not found. Would you like to install it? [Y/n] " % self.binary, choices=['y','n'], default='y')
            if answer == 'n':
                print "+ Aborting..."
                return
            self.install_package()
            
            return self.find_binary()

    def install_package(self):
        pkgtools = {
            "Ubuntu": "apt-get install",
            "Debian": "apt-get install",
            "Fedora": "yum install",
        }
        
        result = re.findall(r"Distributor ID:\t(\w+)", os.popen("lsb_release -i").read())
        if result:
            distro = result[0]
            print "Distribution: %s" % distro

            pkgtool = pkgtools.get(distro)
            if pkgtool is None:
                print "Error: Unknown distro!"
                return
            
            print "Package: %s" % self.package
            print
            
            command = 'sudo %s "%s"' % (pkgtool, self.package)
            
            print "Executing: %s" % command
            print
            execute(command)
            print
            
        else:
            print "Error: Could not find your Linux version via the 'lsb_release' tool."
        

    def extract_file(self, filename, destdir='.'):
        raise Exception, "extract_file() not implemented!"
    
    def extract_url(self, url, destdir='.'):
        opener = URLGrabber()
        filename, headers = opener.retrieve(url)
        print
        self.extract_file(filename, destdir)
        os.remove(filename)
    
    def extract(self, thing, destdir='.'):
        print "thing:", thing
        if isurl(thing):
            archivename = split(urllib.url2pathname(thing))[1]
        elif isfile(thing):
            archivename = split(thing)[1]
        else:
            raise "Unknown thing: %s" % thing
            
        print "archivename:", archivename
        print
            
        self.check_binary()
            
        print "+ Extracting..."
        print "    thing: %s" % thing
        print "    archive name: %s" % archivename
        print "    destination: %s" % destdir
        print "    type: %s" % self.desc
        
        if not exists(destdir):
            answer = prompt("* Destination doesn't exist. Create? [Y/n] ", choices=['y','n'], default='y')
            if answer == 'n':
                print "+ Aborting..."
                return
            print "+ Creating: %s" % destdir
            os.mkdir(destdir)

        tempdir = tempfile.mkdtemp('', 'arf-', destdir)
        print "+ Extracting archive to: %s" % tempdir

        print "-"*50
                            
        if isurl(thing):
            self.extract_url(thing, tempdir)
        elif isfile(thing):
            self.extract_file(thing, tempdir)
        else:
            print "+ I don't know what the crap this thing is:\n\t%s\n\n...what're you trying to pull here, huh?!?!" % thing
            return
        
        print "-"*50
        
        print "archivename: %s" % repr(archivename)
        only_dir = only_dir_in_path(tempdir)
        if only_dir:
            target = join(destdir, only_dir)
            print "+ Placing contained directory in '%s'..." % target
            os.rename(join(tempdir, only_dir), target)
            os.removedirs(tempdir)
        else:
            dest = join(destdir, self.dirname_from_archivename(archivename))
            count = 1
            while exists(dest):
                dest = join(destdir, self.dirname_from_archivename(archivename)) + "(%d)"%count
                count += 1
            print "+ Renaming '%s' to '%s'" % (tempdir, dest)
            os.rename(tempdir, dest)
        
        print "* Done. ARF!"
        print


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

class TarballExtractor(FileExtractor):
    exts            = ['.tar.gz', '.tar.bz2',
                       '.tar', '.tgz', '.tar.lzma', '.tar.lrz',
                       '.tar.xz', '.tpxz', '.gem', '.pixz']
    streamable      = True
    desc            = "TAR file"
    binary          = "tar"
    package         = binary

    # interface
        
    def extract_file(self, filename, destdir=None):
        cmd = self.getcommand(filename)
        cmd = '%s "%s" -C "%s"' % (cmd, filename, destdir)
        print "executing: %s" % cmd
        execute(cmd)

    # utilities
    
    def getcommand(self, filename):
    	lower = filename.lower()
    	
        if lower.endswith(".gz") or lower.endswith(".tgz"):
            cmd = "tar --xattrs -zxvf"
        elif lower.endswith(".bz2"):
            cmd = "tar --xattrs -jxvf"
        elif lower.endswith(".xz") or lower.endswith(".lzma"):
            cmd = "tar --xattrs -Jxvf"
        elif lower.endswith(".lrz"):
            cmd = "tar --xattrs -Ilrzip -xvf"
        elif lower.endswith(".tpxz") or lower.endswith(".pixz"):
            cmd = "tar --xattrs -Ipixz -xvf"
        else:
            cmd = "tar --xattrs -xvf"
        
        return cmd        
                
    
class ZipExtractor(FileExtractor):
    exts            = ['.zip', '.pk3', '.jar', '.war', '.xpi', '.epub', '.apk', '.k3b']
    streamable      = False
    desc            = "ZIP file"
    binary          = "unzip"
    package         = binary

    # interface
    
    def extract_file(self, filename, destdir=None):
        execute('unzip -o "%s" -d "%s"' % (filename, destdir))
        
class RarExtractor(FileExtractor):
    exts            = ['.rar', '.001']
    streamable      = True
    desc            = "RAR file"
    binary          = "unrar"
    package         = binary
    
    def extract_file(self, filename, destdir=None):
        execute('unrar x "%s" "%s"' % (filename, destdir))

class SevenZipExtractor(FileExtractor):
    exts            = ['.7z',]
    streamable      = False
    desc            = "7zip file"
    binary          = ["7z", "7za", "7zr"]
    package         = "p7zip"

    # interface
    
    def extract_file(self, filename, destdir=None):
        binary = self.find_binary()
        execute('%s x "%s" -o"%s"' % (binary, filename, destdir))

class DebExtractor(FileExtractor):
    # Note: alternate method is "ar p mypackage.deb data.tar.gz | tar zx"
    exts            = ['.deb']
    streamable      = True
    desc            = "DEBian package"
    binary          = "dpkg"
    def extract_file(self, filename, destdir=None):
        execute('dpkg -x "%s" "%s"' % (filename, destdir))

class GitExtractor(FileExtractor):
    exts            = ['.git']
    desc            = "Git Repository"
    binary          = "git"
    package         = "git-core"
    
    def extract_url(self, filename, destdir=None):
        execute('git clone "%s" "%s"' % (filename, destdir))

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

extractors = [
                TarballExtractor(),
                ZipExtractor(),
                RarExtractor(),
                DebExtractor(),
                SevenZipExtractor(),
                GitExtractor(),
             ]

def extract(filename, destdir):
    for extractor in extractors:
        if extractor.canHandle(filename):
            extractor.extract(filename, destdir)
            return
    
    print "+ I can't extract THAT! Why you gotta play me like this?!"

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

if __name__ == '__main__':
    madProps()

    if len(sys.argv) <= 1:
        howToUseMe()
    
    else:
        args = sys.argv[1:]
        numargs = len(args)

        filename = args[0]
        
        if numargs == 1:
            destdir = '.'
        elif numargs == 2:
            destdir = args[1]

        extract(filename, destdir)