'''
Script to grab pics from a photobucket album automatically
and save them locally. Provide password if the album is protected
and the album URL.
Author: Kulverstukas
Website: http://9v.lt
Shouts to Evilzone.org and Programisiai.lt
Version: 0.4
http://9v.lt/blog/update-photobucket-ripper/
'''
import os
import sys
import re
import urllib
import mechanize
from argparse import ArgumentParser
#============================================
class ImageMethods:
    def downloadImages(self, links, browser):
        linksList = [i.strip() for i in re.split(',\s{2,}', links)]
        print "* Found "+str(len(linksList))+" images..."
    
        print "* Compiling regex patterns and downloading the pictures..."
        picUrl = re.compile('url: "(https?://(.*?)\.photobucket\.com/albums/(.*?))",')
        picName = re.compile('title: "(.*?)"')
    
        counter = 1
        for link in linksList:
            name = picName.search(link).group(0).replace("title: \"","")[:-1]
            if (name == ""):
                name = noNamePic
            if (os.path.exists(mainFolder+'/'+name) == False):
                os.mkdir(mainFolder+'/'+name)
            picLink = picUrl.search(link)
            picLink = picLink.group(0).replace("url: \"", "")[:-2]
            fileName = os.path.basename(picLink)
            if (os.path.exists(mainFolder+'/'+name+'/'+fileName)):
                print '%d. Retrieving "%s" into "%s" folder' % (counter, fileName, name)
                print "*** "+name+'/'+fileName+" exists. Skipping..."
            else:
                try:
                    size = CalculateSize().calculateSize(browser.open(picLink).info().get("Content-Length"))
                    print '%d. Retrieving "%s" into "%s" folder -- Size: %s' % (counter, fileName, name, size)
                    urllib.urlretrieve(picLink, mainFolder+'/'+name+'/'+fileName)
                except KeyboardInterrupt:
                    print " Terminating..."
                    sys.exit(0)
                except Exception as e:
                    print " Terminating with message: %s" % e
                    sys.exit(0)
            counter += 1
    def grabSlideshowData(self, htmlCode):
        data = re.search("PB\.Slideshow\.data \= \[\n.*\];", htmlCode)
        if (data == None):
            print "*** Something went wrong grabbing picture data. Terminating..."
            sys.exit(0)
        data = data.group(0).replace("PB.Slideshow.data = [", "").replace("];", "").strip()
        
        return data
#============================================
class VideoMethods:
    def downloadVideos(self, list, browser):
        counter = 1
        print "* Found "+str(len(list))+" videos..."
        for item in list:
            url = item[0]
            name = item[1]
            fileName = os.path.basename(url)
            if (os.path.exists(mainFolder+'/'+name) == False):
                os.mkdir(mainFolder+'/'+name)
            if (os.path.exists(mainFolder+'/'+name+'/'+fileName)):
                print '%d. Retrieving "%s" into "%s" folder' % (counter, fileName, name)
                print "*** "+name+'/'+fileName+" exists. Skipping..."
            else:
                try:
                    size = CalculateSize().calculateSize(browser.open(url).info().get("Content-Length"))
                    print '%d. Retrieving "%s" into "%s" folder -- Size: %s' % (counter, fileName, name, size)
                    urllib.urlretrieve(url, mainFolder+'/'+name+'/'+fileName)
                except KeyboardInterrupt:
                    print " Terminating..."
                    sys.exit(0)
                # except Exception as e:
                    # print " Terminating with message: %s" % e
                    # sys.exit(0)
            counter += 1
    
    def grabVideoLinks(self, html):
        list = []
        pattern = "<img src\=\""+album.replace(".", "\.")+".*/>"
        matchObj = re.search("http://(.*?)\.", pattern)
        pattern = pattern.replace(matchObj.group(0), "http://[\w\d]*\.")
        rawList = re.findall(pattern, html)
        videoName = ""
        videoUrl = ""
        for link in rawList:
            # grab the video name and trim crap from it
            videoName = re.search("title=\"(.*?)\"", link).group(0).replace("title=\"","")[:-1]
            # grab the video url and leave only the URL to video
            videoUrl = re.search("alt=\"(.*?)\"", link).group(0).replace("alt=\"","")
            videoUrl = videoUrl.replace(re.search("\s(.*?)\"", videoUrl).group(0), "")
            videoUrl = os.path.join(album, videoUrl)
            list.append((videoUrl, videoName))
            
        return list
#============================================
class CalculateSize:
    def calculateSize(self, bytes):
        abbrevs = ["kB", "mB", "gB"]
        if (bytes == None):
            size = "0 kB"
        else:
            bytes = float(bytes)
            if (bytes < 1024.0):
                size = "%d B" % (bytes)
            else:
                for abbrev in abbrevs:
                    if (bytes >= 1024.0):
                        bytes = bytes / 1024.0
                        size = "%.2f %s" % (bytes, abbrev)
        return size
#============================================
def begin():
    print '\n* Creating "%s" folder...' % mainFolder
    if (os.path.exists(mainFolder)):
        print "*** Folder exists. Skipping..."
    else:
        os.mkdir(mainFolder)
    
    print "* Initiating connection to Photobucket..."
    browser = mechanize.Browser()
    browser.addheaders = [('User-Agent', 'Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)')]
    browser.set_handle_equiv(False)
    #browser.set_debug_http(True)
    browser.set_handle_robots(False)
    try:
        browser.open(album)
    except KeyboardInterrupt:
        print " Terminating..."
        sys.exit(0)
    except Exception as e:
        print " Terminating with message: %s" % e
        sys.exit(0)
    
    # see if the album has a password field
    rawHtml = ""
    for form in browser.forms():
        if (form.name == "frmLogin"):
            if (passwd == ""):
                print "*** Album requires password, none given. Terminating..."
                sys.exit(0)
            print "* Album requires password... using '"+passwd+"'"
            browser.select_form(name="frmLogin")
            browser.form["loginForm[password]"] = passwd
            print "* Submitting password..."
            browser.submit()
            break
    
    if ((filter == "pic") or (filter == "all")):
        print "* Reading image HTML code..."
        rawHtml = browser.open(album+slideshowFilter).read()
        imgMethods = ImageMethods()
        slideshowData = imgMethods.grabSlideshowData(rawHtml)
        imgMethods.downloadImages(slideshowData, browser)
        print "     Done grabbing images!"
    if ((filter == "vid") or (filter == "all")):
        print "* Reading video HTML code..."
        rawHtml = browser.open(album+videoFilter).read()
        vidMethods = VideoMethods()
        links = vidMethods.grabVideoLinks(rawHtml)
        vidMethods.downloadVideos(links, browser)
        print "     Done grabbing videos!"
#=============================================
parser = ArgumentParser(description="Script to grab and save pics from a photobucket album automatically")
parser.add_argument('-u', '--url', help='Album URL', required=True)
parser.add_argument('-p', '--passwd', help='Album password (if any)', default="")
parser.add_argument('-f', '--filter', help='What to download (pic/vid/all)', default="all")
parser.add_argument('-d', '--dir', help='Where to download (folder name)', default="PhotobucketGetter")
args = parser.parse_args()
#====== global vars, change values here ======
noNamePic = 'NoName'
slideshowFilter = "?albumview=slideshow"
videoFilter = "?mediafilter=videos"
mainFolder = args.dir
album = args.url
passwd = args.passwd
filter = args.filter
if ((filter != 'pic') and (filter != 'vid') and (filter != 'all')):
    filter = 'all'
#=============================================
begin()