from __future__ import division
import requests
import zipfile
import cStringIO as StringIO
import csv
import redis
from urlparse import urlparse
from publicsuffix import PublicSuffixList

from malware_crawl.defaults import redis_db

domains_key = "alexa:domains"


def chunks(l, n):
    """ Yield successive n-sized chunks from l.
    """
    for i in xrange(0, len(l), n):
        yield l[i:i + n]


def dict_chunk(mapping, chunk_size):
    items = mapping.items()
    for chunk in chunks(items, chunk_size):
        yield dict(chunk)


def alexa_top_million():
    """Returns an iterable of objects describing the top million domains from Alexa
    It turns out this list can be used as a set and held in memory fairly
    easily:

    domains = frozenset({item["domain"] for item in alexa_top_million()})
    if "google.com" in domains:
        return u"safe"

    However it might be better to stick this in a *real* database local to
    each worker, and update it each day. Might as well use the Redis DB we
    have lying around. It might be even better to have a BloomFilter before
    even bothering to check the DB.
    """
    alexa_table = csv.DictReader(
        zipfile.ZipFile(
            StringIO.StringIO(
                requests.get("http://s3.amazonaws.com/alexa-static/top-1m.csv.zip").content
            )
        ).open("top-1m.csv"),
        fieldnames=("score", "domain")
    )
    return {item["domain"]: item["score"] for item in alexa_table}


def update_alexa():
    pipe = redis_db["master"].pipeline()
    pipe.delete(domains_key)

    for chunk in dict_chunk(alexa_top_million(), 1000):
        pipe.hmset(domains_key, chunk)

    pipe.execute()


def alexa_malware_scan(url):
    domain = PublicSuffixList().get_public_suffix(urlparse(url).netloc)  # IRIs are going to be a pain here.
    pipe = redis_db["slave"].pipeline()
    pipe.hlen(domains_key)
    pipe.hmget(domains_key, domain)
    total, score = pipe.execute()

    return {
        "type": "clean",
        "confidence": (total - (score - 1) / total)
    }

from __future__ import division
import requests
import zipfile
import cStringIO as StringIO
import csv
import redis
from urlparse import urlparse
from publicsuffix import PublicSuffixList

from celery.utils.functional import chunks

from malware_crawl.defaults import redis_db

domains_key = "alexa:domains"


def dict_chunk(mapping, chunk_size):
    items = mapping.items()
    for chunk in chunks(items, chunk_size):
        yield dict(chunk)


def alexa_top_million():
    """Returns an iterable of objects describing the top million domains from Alexa
    It turns out this list can be used as a set and held in memory fairly
    easily:

    domains = frozenset({item["domain"] for item in alexa_top_million()})
    if "google.com" in domains:
        return u"safe"

    However it might be better to stick this in a *real* database local to
    each worker, and update it each day. Might as well use the Redis DB we
    have lying around. It might be even better to have a BloomFilter before
    even bothering to check the DB.
    """
    alexa_table = csv.DictReader(
        zipfile.ZipFile(
            StringIO.StringIO(
                requests.get("http://s3.amazonaws.com/alexa-static/top-1m.csv.zip").content
            )
        ).open("top-1m.csv"),
        fieldnames=("score", "domain")
    )
    return {item["domain"]: item["score"] for item in alexa_table}


def update_alexa():
    pipe = redis_db["master"].pipeline()
    pipe.delete(domains_key)

    for chunk in dict_chunk(alexa_top_million(), 1000):
        pipe.hmset(domains_key, chunk)

    pipe.execute()


def alexa_malware_scan(url):
    domain = PublicSuffixList().get_public_suffix(urlparse(url).netloc)  # IRIs are going to be a pain here.
    pipe = redis_db["slave"].pipeline()
    pipe.hlen(domains_key)
    pipe.hmget(domains_key, domain)
    total, score = pipe.execute()

    return {
        "type": "clean",
        "confidence": (total - (score - 1) / total)
    }

use chunks function from celery