import hashlib
from itertools import izip_longest

# Big Endian bits
BE_MUL = tuple(2**(7-n) for n in range(8))

def grouper(n, iterable, fillvalue=None):
    "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return izip_longest(fillvalue=fillvalue, *args)


class PackageFileBloomFilter(object):

    def __init__(self, hash_function_name="sha384", records=None):
        super(PackageFileBloomFilter, self).__init__()
        self.hash_function_name = hash_function_name
        self.hasher = getattr(hashlib, hash_function_name)
        self.bloom_pools = tuple(set() for x in self.hasher().digest())

    def _item_hashed_value(self, filename):
        h = self.hasher()
        h.update(filename)
        return map(ord, h.digest())

    def __contains__(self, item):
        return self.check(filename)

    def check(self, filename):
        hashed_value = self._item_hashed_value(filename)
        for bucket, item in zip(self.bloom_pools, hashed_value):
            if item not in bucket:
                return False
        return True

    def add(self, filename):
        if not isinstance(filename, unicode):
            raise ValueError("Filenames must be unicode.")
        if self.hasher is None:
            raise ValueError("This bloom filter is read-only.")

        hashed_value = self._item_hashed_value(filename)

        all_in_set = True
        for bucket, item in zip(self.bloom_pools, hashed_value):
            # Hopefully, one of the buckets doesn't have it set.
            all_in_set = all_in_set and item in bucket
            bucket.add(item)
        if all_in_set:
            warnings.warn("Collision in bloom filter!")

    def read(self, mdns_records):
        self.hasher = None
        self.bloom_pools = tuple(set() for x in self.hasher.digest())

    def publish_to_avahi_service_file(self, destination_file, port_number):
        MDNS_HEADER = """<?xml version="1.0" standalone='no'?><!DOCTYPE service-group SYSTEM "avahi-service.dtd"><service-group><name replace-wildcards="yes">%h APT cache</name><service><type>_apt-http._tcp</type><port>{port_number:d}</port><txt-record>version=1,{hash_function_name!s}-bloom</txt-record>"""
        MDNS_FOOTER = """</service></service-group>"""
        MDNS_RECORD_FORMAT = "<txt-record>octet-{pool_number:03x}={pool_str}</txt-record>"

        hash_function_name = self.hash_function_name  # For formatting.
        with open(destination_file, "w") as f:
            f.write(MDNS_HEADER.format(**locals()))
            for pool_number, pool in enumerate(self.bloom_pools):
                chars = list()
                inside = lambda i: i in pool
                for positions in grouper(8, range(256)):
                    bits = [(1 if inside(x) else 0) for x in positions]
                    chars.append(chr(sum((x*e) for e, x in zip(BE_MUL, bits))))
                pool_str = "".join(chars).encode("base64").strip()

                f.write(MDNS_RECORD_FORMAT.format(**locals()))
            f.write(MDNS_FOOTER)

import os
pfbf = PackageFileBloomFilter()
for filename in os.listdir("/var/cache/apt/archives/"):
    pfbf.add(filename.decode("utf8"))

pfbf.publish_to_avahi_service_file("/var/tmp/avahi-apt-http.service", 5000)