#!/usr/bin/env python

from base64 import b64encode, b64decode

import logging
log = logging.getLogger('gglsbl')
log.addHandler(logging.NullHandler())

from gglsbl.utils import to_hex
from gglsbl.protocol import SafeBrowsingApiClient, URL
from gglsbl.storage import SqliteStorage, ThreatList, HashPrefixList


class SafeBrowsingList(object):
    """Interface for Google Safe Browsing API

    supporting partial update of the local cache.
    https://developers.google.com/safe-browsing/v4/
    """

    def __init__(self, api_key, db_path='/tmp/gsb_v4.db', discard_fair_use_policy=False,
                        platforms = None):
        """Constructor.

        Args:
            api_key: string, a key for API authentication.
            db_path: string, path to SQLite DB file to store cached data.
            discard_fair_use_policy: boolean, disable request frequency throttling (only for testing).
            platforms: list, threat lists to look up, default includes all platforms.
        """
        self.api_client = SafeBrowsingApiClient(api_key, discard_fair_use_policy=discard_fair_use_policy)
        self.storage = SqliteStorage(db_path)
        self.platforms = platforms

    def _verify_threat_list_checksum(self, threat_list, remote_checksum):
        local_checksum = self.storage.hash_prefix_list_checksum(threat_list)
        return remote_checksum == local_checksum

    def update_hash_prefix_cache(self):
        """Update locally cached threat lists.
        """
        self.api_client.fair_use_delay()
        self.storage.cleanup_full_hashes()
        threat_lists_to_remove = dict()
        for ts, cs in self.storage.get_threat_lists():
            threat_lists_to_remove[repr(ts)] = ts
        threat_lists = self.api_client.get_threats_lists()
        for entry in threat_lists:
            threat_list = ThreatList.from_api_entry(entry)
            if self.platforms is None or threat_list.platform_type in self.platforms:
                self.storage.add_threat_list(threat_list)
                try:
                    del threat_lists_to_remove[repr(threat_list)]
                except KeyError:
                    pass
        for ts in threat_lists_to_remove.values():
            self.storage.delete_hash_prefix_list(ts)
            self.storage.delete_threat_list(ts)
        del threat_lists_to_remove

        self.api_client.fair_use_delay()
        threat_lists = self.storage.get_threat_lists()
        client_state = dict([(t.as_tuple(), s) for t,s in threat_lists])
        for response in self.api_client.get_threats_update(client_state):
            response_threat_list = ThreatList(response['threatType'], response['platformType'], response['threatEntryType'])
            if response['responseType'] == 'FULL_UPDATE':
                self.storage.delete_hash_prefix_list(response_threat_list)
            for r in response.get('removals', []):
                self.storage.remove_hash_prefix_indices(response_threat_list, r['rawIndices']['indices'])
            for a in response.get('additions', []):
                hash_prefix_list = HashPrefixList(a['rawHashes']['prefixSize'], b64decode(a['rawHashes']['rawHashes']))
                self.storage.populate_hash_prefix_list(response_threat_list, hash_prefix_list)
            expected_checksum = b64decode(response['checksum']['sha256'])
            if self._verify_threat_list_checksum(response_threat_list, expected_checksum):
                log.info('Local cache checksum matches the server: {}'.format(to_hex(expected_checksum)))
                self.storage.update_threat_list_client_state(response_threat_list, response['newClientState'])
            else:
                raise Exception('Local cache checksum does not match the server: "{}". Consider removing {}'.format(to_hex(expected_checksum), self.storage.db_path))

    def _sync_full_hashes(self, hash_prefixes):
        """Download full hashes matching hash_prefixes.

        Also update cache expiration timetsamps.
        """
        threat_lists = self.storage.get_threat_lists()
        client_state = dict([(t.as_tuple(), s) for t,s in threat_lists])
        self.api_client.fair_use_delay()
        fh_response = self.api_client.get_full_hashes(hash_prefixes, client_state)

        # update negative cache for each hash prefix
        # store full hash (insert or update) with positive cache bumped up
        for m in fh_response['matches']:
            threat_list = ThreatList(m['threatType'], m['platformType'], m['threatEntryType'])
            hash_value = b64decode(m['threat']['hash'])
            cache_duration = int(m['cacheDuration'].rstrip('s'))
            malware_threat_type = None
            for metadata in m['threatEntryMetadata'].get('entries', []):
                k = b64decode(metadata['key'])
                v = b64decode(metadata['value'])
                if k == 'malware_threat_type':
                    malware_threat_type = v
            self.storage.store_full_hash(threat_list, hash_value, cache_duration, malware_threat_type)

        negative_cache_duration = int(fh_response['negativeCacheDuration'].rstrip('s'))
        for prefix_value in hash_prefixes:
            for threat_list in threat_lists:
                self.storage.update_hash_prefix_expiration(threat_list[0], prefix_value, negative_cache_duration)

    def lookup_url(self, url):
        """Look up specified URL in Safe Browsing threat lists."""
        if type(url) is not str:
            url = url.encode('utf8')
        url_hashes = URL(url).hashes
        list_names = self._lookup_hashes(url_hashes)
        if list_names:
            return list_names
        return None

    def _lookup_hashes(self, full_hashes):
        """Lookup URL hash in blacklists

        Returns names of lists it was found in.
        """
        full_hashes = list(full_hashes)
        cues = [to_hex(fh[0:4]) for fh in full_hashes]
        result = []
        try:
            matching_prefixes = {}
            matching_full_hashes = set()
            is_potential_threat = False
            # First lookup hash prefixes which match full URL hash
            for (threat_list, hash_prefix, negative_cache_expired) in self.storage.lookup_hash_prefix(cues):
                for full_hash in full_hashes:
                    if full_hash.startswith(hash_prefix):
                        is_potential_threat = True
                        # consider hash prefix negative cache as expired if it is expired in at least one threat list
                        matching_prefixes[hash_prefix] = matching_prefixes.get(hash_prefix, False) or negative_cache_expired
                        matching_full_hashes.add(full_hash)
            # if none matches, URL hash is clear
            if not is_potential_threat:
                return []
            # if there is non-expired full hash, URL is blacklisted
            matching_expired_threat_lists = set()
            for threat_list, has_expired in self.storage.lookup_full_hashes(matching_full_hashes):
                if has_expired:
                    matching_expired_threat_lists.add(threat_list)
                else:
                    result.append(threat_list)
            if result:
                return result

            # If there are no matching expired full hash entries
            # and negative cache is still current for all prefixes, consider it safe
            if len(matching_expired_threat_lists) == 0 and sum(map(int, matching_prefixes.values())) == 0:
                log.info('Negative cache hit.')
                return []

            # Now we can assume that there are expired matching full hash entries and/or
            # cache prefix entries with expired negative cache. Both require full hash sync.
            self._sync_full_hashes(matching_prefixes.keys())
            # Now repeat full hash lookup
            for threat_list, has_expired in self.storage.lookup_full_hashes(matching_full_hashes):
                if not has_expired:
                    result.append(threat_list)
        except:
            self.storage.db.rollback()
            raise
        return result
