hacky easy rom loader

coldroms

#!/usr/bin/env python
# -*- coding: utf-8 -*-
try:
    import requests, clint, easydict
except ImportError:
    import pip
    pip.main(["install", "requests", "clint", "easydict"])
    import requests, clint, easydict

import tempfile, re, sys, os, zipfile
from easydict import EasyDict as edict
from clint.textui import colored, puts, progress

__version__ = "0.1"


def _build_session():
    session = requests.Session()
    session.headers.update(
        {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_1)'
                       ' AppleWebKit/537.36 (KHTML, like Gecko) '
                       'Chrome/31.0.1650.63 Safari/537.36'}
    )
    return session


def _choose(games):
    puts("unsure... please choose:")
    for index, game in enumerate(games):
        puts("{}: {game.name} on {platform}".format(
            colored.red(index),
            game=game,
            platform=colored.red(game.platform)
        ))
    user_option = int(raw_input("Option#: "))
    if len(games) < user_option < -1:
        return None
    return games[user_option]


def search(session, term, platform=None):
    params = {'q': term}
    if platform:
        params["system"] = platform
    r = session.get(
        "http://coolrom.com/search",
        params=params)
    results = []
    for line in r.text.splitlines():
        if "/roms" not in line:
            continue
        if "Top 25 Downloaded" in line:
            break
        for match in re.finditer(r'roms/(?P<platform>[\w\-_\d]+)/(?P<game_id>\d+)/(?P<game_name>[\w\-_\d\']+)', line):
            results.append(edict({
                'platform': match.group("platform"),
                'name': match.group("game_name").replace("_"," ").strip(),
                'id': int(match.group("game_id"))
            }))
    return results


def choose_game(search_term, games, force=False):
    # http://en.wikibooks.org/wiki/Algorithm_Implementation/Strings/Levenshtein_distance#Python
    def levenshtein(s1, s2):
        if len(s1) < len(s2):
            return levenshtein(s2, s1)
        # len(s1) >= len(s2)
        if len(s2) == 0:
            return len(s1)
        previous_row = xrange(len(s2) + 1)
        for i, c1 in enumerate(s1):
            current_row = [i + 1]
            for j, c2 in enumerate(s2):
                insertions = previous_row[j + 1] + 1 # j+1 instead of j since previous_row and current_row are one character longer
                deletions = current_row[j] + 1       # than s2
                substitutions = previous_row[j] + (c1 != c2)
                current_row.append(min(insertions, deletions, substitutions))
            previous_row = current_row
        return previous_row[-1]
    if not force:
        for game in games:
            d = levenshtein(game.name, search_term)
            if d <= 5:
                return game
    if not games:
        return
    if len(games) == 1:
        return games[0]
    return _choose(games)


def download(session, game, dst_dir, unpack):
    puts("Preparing download...")
    r = session.get("http://coolrom.com/dlpop.php", params={"id": game.id})
    p = re.compile(r'.*?action="(?P<url>.*?)"',re.DOTALL)
    url = p.match(r.text).group('url')

    dst = None
    dst_path = os.path.join(dst_dir, game.platform)
    if unpack:
        file_obj_func = tempfile.TemporaryFile
        file_obj_func_args = []
    else:
        dst = os.path.join(dst_path, "{game.name}.zip".format(game=game))
        file_obj_func = open
        file_obj_func_args = [dst, "wb"]
    try:
        os.makedirs(dst_path)
    except OSError:
        pass

    puts("Downloading {}".format(url))
    session.headers.update({"Referer": r.url})
    response = session.post(url, stream=True, allow_redirects=False)
    if response.status_code > 300:
        puts(colored.red("error: we got redirected !"))
        return
    with file_obj_func(*file_obj_func_args) as handle:
        file_size_dl = 0
        try:
            file_size = int(response.headers['content-length'])
        except KeyError:
            file_size = 0
        puts(u"Loading %s for %s" % (game.name, game.platform))
        with progress.Bar(
                filled_char=u"👾 ",
                empty_char="  ",
                expected_size=file_size) as bar:
            for block in response.iter_content(1024):
                if not block:
                    break
                file_size_dl += len(block)
                handle.write(block)
                bar.show(file_size_dl)
        if unpack:
            puts("unpacking ...")
            with zipfile.ZipFile(handle) as zip_file:
                files = zip_file.infolist()
                game_file = None
                for file_info in files:
                    if file_info.filename.startswith("readme"):
                        continue
                    game_file = file_info
                    break
                dst = os.path.join(dst_path, game_file.filename)
                zip_file.extract(game_file, dst)
    return dst


def main(search_term, dst_dir=None, platform=None, force_choose=False, unpack=False):
    if platform is None:
        puts("Searching for {}...".format(colored.red(search_term)))
    else:
        puts("Searching for {} on {}...".format(
            colored.red(search_term),
            colored.red(platform),
        ))
    session = _build_session()
    games = search(session, term, platform)
    puts("Found {} result(s)".format(colored.red(len(games))))
    game = choose_game(search_term, games, force_choose)
    if game is None:
        puts(colored.red("Nothing found!"))
        return False
    puts("Choose: {name} on {platform}".format(
        name=colored.red(game.name),
        platform=colored.red(game.platform)
    ))
    location = download(session, game, dst_dir, unpack)
    if not location:
        puts(colored.red("nothing downloaded :("))
    else:
        puts("DONE game at {}".format(colored.red(location)))
    return location


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser(prog='lazy coolrom')
    p.add_argument(
        '-v', '--version', action="version", version=__version__)
    p.add_argument("term", nargs='+', default=None, help="game name")
    p.add_argument(
        "-p", "--platform", dest='platform',
        default=None, help="desired platform")
    p.add_argument(
        '-o', '--out', dest='dst_dir', default="./",
        help="destination derectory, default ./")
    p.add_argument(
        '-c', '--choose', dest='force_choose', action="store_true",
        help="force the choose screen")
    p.add_argument(
        '-u', '--unpack', dest='unpack', action="store_true",
        help="unpack and clean")

    args = p.parse_args()
    term = " ".join(args.term).strip()
    sys.exit(
        int(not main(
            term,
            args.dst_dir,
            platform=args.platform,
            force_choose=args.force_choose,
            unpack=args.unpack))
    )