Script permettant d'effectuer des recherches d'images sur Google

google_pic_search.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Pour les tests :
# python google_pic_search.py -l https://twimg0-a.akamaihd.net/profile_images/2298747498/n9hzj4nx98tgsxnpdagd.png
# python google_pic_search.py -i /home/initbrain/Images/initbrain.png

# Importations
import sys
import os
import re
import urllib2
import mimetypes
import json
import lxml.html as lh
from urllib import quote_plus
from poster.encode import multipart_encode
from poster.streaminghttp import register_openers
#import getpass

# Solution mise en place pour pallier à un problème de buffer avec mingw32...
if not "linux" in sys.platform:
    if not os.getenv("PYTHONUNBUFFERED"):
        print "[!] La variable d'environnement PYTHONUNBUFFERED doit etre declaree avec une valeure non vide !!!"
        sys.exit()


def get_content_type(filename):
    return mimetypes.guess_type(filename)[0] or "application/octet-stream"


def googleSearchPic(url, search_type, target, http_proxy=None):
    if search_type == 1: # URL
        data = dict()
        headers = dict()
    elif search_type == 2: # Fichier
        # Register the streaming http handlers with urllib2
        opener = register_openers()
        # headers contains the necessary Content-Type and Content-Length
        # data is a generator object that yields the encoded parameters
        data, headers = multipart_encode({"encoded_image": open(target, "rb")})

    headers.update({"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
                    "User-Agent": "Mozilla/5.0 (Windows NT 5.1; rv:16.0) Gecko/20100101 Firefox/16.0",
                    "Connection": "keep-alive",
                    "Cache-Control": "no-cache",
                    "Pragma": "no-cache",
                    "Referer": "http://images.google.fr/"})

    if http_proxy:
        # http://username:password@someproxyserver.com:1337
        http_proxy_full_auth_string = "http://%s%s@%s:%s" % (http_proxy["user"],
                                                             ':' + http_proxy["passwd"] if http_proxy.has_key("passwd") else '',
                                                             http_proxy["server"],
                                                             http_proxy["port"])
        proxy_handler = urllib2.ProxyHandler({"http": http_proxy_full_auth_string,
                                              "https": http_proxy_full_auth_string})

        if search_type != 2:
            opener = urllib2.build_opener(proxy_handler)
            urllib2.install_opener(opener)
        else:
            opener.add_handler(proxy_handler)

    if search_type == 1:
        resRequest = req(url, None, headers)
    elif search_type == 2:
        resRequest = req(url, data, headers)
    if resRequest:
        return resRequest
    else:
        return 0

def req(url, data=None, headers=None, timeout=10, retry=2):
    if not headers:
        headers = {"User-Agent": "Mozilla/5.0 (Windows NT 5.1; rv:16.0) Gecko/20100101 Firefox/16.0",
                   "Cache-Control": "no-cache",
                   "Pragma": "no-cache",
                   "Referer": "http://www.google.fr/"}

    request = urllib2.Request(url, data, headers)
    source = False
    loop = 0
    while not source and loop <= retry:
        try:
            connection = urllib2.urlopen(request, timeout=timeout)
        except Exception as err: # TODO
            # Si il y a une erreur de connexion (timeout etc.)
            print "ERREUR : " + url + " ne repond pas ...\n%s" % err
        else:
            if connection.getcode() != 200:
                print "ERREUR : la page " + url + " n'existe plus ..."
            else:
                #print connection.geturl()
                source = connection.read()
                connection.close()
                if not source:
                    print "ERREUR : le code source de " + url + " est vide ..."
                else:
                    return source
        loop += 1

def googleResultParser(source):
    """ Parse les résultats présents sur une page de Google.com
        Exemple de sortie :
        {"http://www.toto.fr/image.jpg": ["http://www.toto.fr/page1.html", "http://www.toto.fr/page2.html"],
         "http://www.tata.fr/image.jpg": ["http://www.tata.fr/page1.html"]}
    """
    tree = lh.fromstring(source)
    pictures_urls = tree.xpath("//h3[@class='r']/a/@href")
    return pictures_urls


def progressbar(i, count, prefix="", size=60):
    x = int(size * i / count)
    sys.stdout.write("%s[%s%s] page %i/%i\r" % (prefix, "#" * x, "." * (size - x), i, count))
    if i == count:
        sys.stdout.write("\n") # Deplacer le cuseur sur la ligne suivante
    sys.stdout.flush()


def main():
    http_proxy = None

    # Variables
    usage = """
    Usage :
        python """ + sys.argv[0] + """ [Type de recherche] \"URL/FILE_PATH\"

    Options :
        --lien/-lien/-l :       Effectuer une recherche a partir du lien d'une image
        --image/-image/-i :     Effectuer une recherche a partir d'un fichier image
        --proxy/-proxy/-p :     Utiliser un proxy
        --verbose/-verbose/-v : Afficher les chemins des fichiers enregistres
        --help/-help/-h :       Afficher ce message d'aide

    Exemples :
        python """ + sys.argv[0] + """ -l http://www.toto.fr/images/test.jpg
            Realiser une recherche a partir d'une url sans passer par un proxy

        python """ + sys.argv[0] + """ -p -i images/toto.jpg
            Realiser une recherche a partir d'un fichier image en passant par un proxy
    """

    # Verification qu'il y a au moins 2 arguments passe au script
    if len(sys.argv) < 2:
        print "ERREUR : Vous devez renseigner au moins 2 arguments\n" + usage
        sys.exit()
    else:
        # Verification du nombre de parametres et de leur type
        error = ""
        search_type = 0
        target = False
        for param in sys.argv[1:]:
            if param in ["--help", "-help", "-h"]:
                print usage
                sys.exit()
            elif param in ["--proxy", "-proxy", "-p"]:
                # Créer un tableau qui permettra d'accueillir les informations saisies par l'utilisateur par la suite
                http_proxy = {}
                # Possibilité d'hardcoder les infos d'un proxy => PAS BIEN !
                #http_proxy = {"server": "myproxy.tld",
                #              "port": "8080",
                #              "user": "myuser",
                #              "passwd": "mypasswd"}
            elif param in ["--verbose", "-verbose", "-v"]:
                verbose = True # Not used but KEEP IT!!!!
            elif param in ["--lien", "-lien", "-l"]:
                if not search_type:
                    if target:
                        error += "ERREUR : le type de recherche (url/fichier) doit etre precise en premier\n"
                    else:
                        search_type = 1
                else:
                    error += "ERREUR : le type de recherche (url/fichier) a deja ete renseigne\n"
            elif param in ["--image", "-image", "-i"]:
                if target:
                    error += "ERREUR : le type de recherche (url/fichier) doit etre precise en premier\n"
                else:
                    search_type = 2
            elif param[0:7] == "http://" or param[0:8] == "https://":
                if search_type == 2:
                    error = error + "ERREUR : parametre invalide '" + param + "'\n"
                else:
                    target = param
            elif search_type == 2:
                if not os.path.isfile(param): # check si fichier existe
                    error = error + "ERREUR : le fichier '" + param + "' est introuvable\n"
                else:
                    # check si le fichier est bien une image et qu'il correspond
                    # aux formats supportes par le moteur de recherche (TinEye)
                    image_format = get_content_type(param)
                    if image_format not in ["image/png", "image/gif", "image/jpeg"]:
                        error = error + "ERREUR : format '" + image_format + "' non pris en charge\n"
                    else:
                        target = param
            else:
                error = error + "ERREUR : parametre invalide '" + param + "'\n"

        if not search_type:
            error += "ERREUR : veuillez preciser le type de recherche\n"
        if not target:
            error += "ERREUR : veuillez indiquer l'image a rechercher\n"
        if error:
            print error + usage
            sys.exit()

        # Parametrage du proxy en DIRECT ou via PROXY
        if http_proxy is not None:
            # On passe par un proxy
            if not http_proxy.has_key("user") or not http_proxy.has_key("passwd") or not http_proxy.has_key("server") or not http_proxy.has_key("port"):
                http_proxy["server"] = raw_input('[Proxy] Serveur : ')
                http_proxy["port"] = raw_input('[Proxy] Port : ')
                http_proxy["user"] = raw_input('[Proxy] Identifiant (laisser vide si aucun) : ')
                # getpass permet de ne pas afficher la saisie
                http_proxy["passwd"] = raw_input("[Proxy] Mot de passe (laisser vide si aucun) : ")
                # TODO getpass.getpass("blah") fait foirer le unbuffered (mingw32...)
            else:
                print "[!] Attention, des identifiants sont visibles dans les variables d'environnement !"

    print "Envoi de l'image au moteur de recherche Google"
    if search_type == 1: # URL
        url = "http://www.google.fr/searchbyimage?image_url=%s&hl=fr&authuser=0&filter=0" % quote_plus(target)
    elif search_type == 2: # Fichier
        url = "http://images.google.fr/searchbyimage/upload"

    source = googleSearchPic(url, search_type, target, http_proxy)
    if not source:
        sys.exit()
        # TODO adapter les erreurs à celles de google
    if "Your image is too simple to create a unique fingerprint." in source:
        print "Aucun resultat (image trop simple) !"
    elif "Could not read the image." in source:
        print "Aucun resultat (image corrompue) !"
    else:
        resRegExpResults = re.compile("id=\"resultStats\">Environ (\d*)", re.MULTILINE).findall(source)
        if len(resRegExpResults) == 1:
            if resRegExpResults[0] == 0:
                print "Aucun resultat !"
            else:
                resRegExpSearchHash = re.compile("&tbs=(.*?)\">Recherche avancée</a>", re.MULTILINE).findall(source)
                if len(resRegExpSearchHash):
                    startPage = 0
                    urlPage = "http://www.google.fr/search?hl=fr&authuser=0&tbs=" + resRegExpSearchHash[0] + "&start=%d&filter=0" % startPage
                    finalResult = list()
                    finalResult += googleResultParser(source)
                    while "start=" + str(startPage + 10) in source in source:
                        startPage += 10
                        #progressbar(startPage, int(resRegExpResults[0]), "Recuperation : ")
                        print "Google: Page %d en cours (%d lien(s) trouve(s) jusqu'ici)" % (startPage / 10 + 1, len(finalResult))
                        #print "Page " + str(numPage) + " en cours (" + str(len(finalResult)) + " resultats pour l'instant)"
                        urlPage = "http://www.google.fr/search?hl=fr&authuser=0&tbs=" + resRegExpSearchHash[0] + "&start=%d&filter=0" % startPage
                        source = req(urlPage)
                        finalResult += googleResultParser(source)
                    else:
                        print "Recherche terminee (" + str(len(finalResult)) + " resultats) !"
                        if len(finalResult) > 0:
                            print "Creation du fichier google_search_result.txt"
                            with open('google_search_result.txt', 'w') as bakupFile:
                               bakupFile.write(str(json.dumps(finalResult, indent=4)))
        else:
            print "[!] L'envoi de l'image a Google a echoue !"


if __name__ == '__main__':
    main()