from urlparse import urlparse
from datetime import timedelta

from six.moves import zip, map
from six import binary_type

from celery import task, current_app
from celery.contrib.batches import Batches

from malware_crawl.defaults import redis_db, REQUESTS_SESSIONS
requests = REQUESTS_SESSIONS["api"]
from django.conf import settings

import cPickle as pickle

wot_api_target = "https://api.mywot.com/0.4/public_link_json"


def redis_cache(key_prefix, prepare, request, items):
    items = list(map(prepare, items))
    unique_items = list(set(items))
    total_response = {}

    # get all the items that are already in cache
    pipe = redis_db["slave"].pipeline()
    for item in unique_items:
        pipe.get(key_prefix.format(item=item))

    cache_misses = set()
    # add those items we already know about to the total_response
    for item, cache in zip(unique_items, pipe.execute()):
        if cache is not None:
            total_response[item] = pickle.loads(cache)
        else:
            cache_misses.add(item)

    # use the given function to collect the missing values
    total_response.update(request(cache_misses))

    # add to the cache the result from the INTERNET
    pipe = redis_db["master"].pipeline()
    cache_length = timedelta(minutes=30).seconds
    for item in cache_misses:
        pipe.setex(key_prefix.format(item=item),
            pickle.dumps(total_response[item], -1),
            cache_length
        )
    pipe.execute()

    return [total_response[item] for item in items]


@task(base=Batches, flush_every=100, flush_interval=((10, 0.1)[settings.DEBUG]))
def wot_malware_scan(requests):
    sig = lambda url: url
    reponses = wot_malware_scan_real(
        (sig(*request.args, **request.kwargs) for request in requests)
    )
    # use mark_as_done to manually return response data
    for response, request in zip(reponses, requests):
        current_app.backend.mark_as_done(request.id, response)


def prepare_requests(domains):
    "ensure requests to the WOT api are less than 8KiB"
    prepared = list(map(lambda domain: binary_type(domain + '/'), set(domains)))
    while prepared:
        req = binary_type(wot_api_target + "?hosts=")
        while prepared and (len(req) + len(prepared[-1])) <= (8 * 1024):
            req += prepared.pop()
        yield req


def wot_malware_scan_real(urls):
    def prepare(url):
        return urlparse(url).netloc

    def request(domains):
        total_response = {}
        # get those we didn't from the INTERNET
        for api_url in prepare_requests(domains):
            try:
                total_response.update(requests.get(api_url).json())
            except:
                print api_url
                raise
        return total_response

    responses = redis_cache(
        key_prefix="malucrawl:web_api:wot:{item}",
        prepare=prepare,
        request=request,
        items=urls
    )

    return [{
        "type": "generic",
        "confidence": (
            response.get('0', [0, 0])[0] / 100
         )
    } for response in responses]


if __name__ == "__main__":
    print wot_malware_scan_real(("http://google.com", "http://yahoo.com"))

from urlparse import urlparse
from datetime import timedelta

from six.moves import zip, map
from six import binary_type

from celery import task, current_app
from celery.contrib.batches import Batches

from malware_crawl.defaults import redis_db, REQUESTS_SESSIONS
requests = REQUESTS_SESSIONS["api"]
from django.conf import settings

import cPickle as pickle

wot_api_target = "https://api.mywot.com/0.4/public_link_json"


def redis_cache(key_prefix, prepare, request, items):
    items = list(map(prepare, items))
    unique_items = list(set(items))
    total_response = {}

    # get all the items that are already in cache
    pipe = redis_db["slave"].pipeline()
    for item in unique_items:
        pipe.get(key_prefix.format(item=item))

    cache_misses = set()
    # add those items we already know about to the total_response
    for item, cache in zip(unique_items, pipe.execute()):
        if cache is not None:
            total_response[item] = pickle.loads(cache)
        else:
            cache_misses.add(item)

    # use the given function to collect the missing values
    total_response.update(request(cache_misses))

    # add to the cache the result from the INTERNET
    pipe = redis_db["master"].pipeline()
    cache_length = timedelta(minutes=30).seconds
    for item in cache_misses:
        pipe.setex(key_prefix.format(item=item),
            pickle.dumps(total_response[item], -1),
            cache_length
        )
    pipe.execute()

    return [total_response[item] for item in items]


@task(base=Batches, flush_every=100, flush_interval=((10, 0.1)[settings.DEBUG]))
def wot_malware_scan(requests):
    sig = lambda url: url
    reponses = wot_malware_scan_real(
        (sig(*request.args, **request.kwargs) for request in requests)
    )
    # use mark_as_done to manually return response data
    for response, request in zip(reponses, requests):
        current_app.backend.mark_as_done(request.id, response)


def prepare_requests(domains):
    "ensure requests to the WOT api are less than 8KiB"
    prepared = list(map(lambda domain: binary_type(domain + '/'), set(domains)))
    while prepared:
        req = binary_type(wot_api_target + "?hosts=")
        while prepared and (len(req) + len(prepared[-1])) <= (8 * 1024):
            req += prepared.pop()
        yield req


def wot_malware_scan_real(urls):
    def prepare(url):
        return urlparse(url).netloc

    def request(domains):
        total_response = {}
        # get those we didn't from the INTERNET
        for api_url in prepare_requests(domains):
            try:
                total_response.update(requests.get(api_url).json())
            except:
                print api_url
                raise
        return total_response

    responses = redis_cache(
        key_prefix="malucrawl:web_api:wot:{item}",
        prepare=prepare,
        request=request,
        items=urls
    )

    return [{
        "type": "generic",
        "confidence": (
            response.get('0', [0, 0])[0] / 100
         )
    } for response in responses]


if __name__ == "__main__":
    print wot_malware_scan_real(("http://google.com", "http://yahoo.com"))

# get all the items that are already in cache # add those items we already know about to the total_response # use the given function to collect the missing values # add to the cache the result from the INTERNET # use mark_as_done to manually return response data # get those we didn't from the INTERNET