"""The entropy calculation implementation."""

import collections
import math

from plaso.analyzers.hashers import interface
from plaso.analyzers.hashers import manager


class EntropyHasher(interface.BaseHasher):
    """Calculates the byte entropy of input files."""

    NAME = "entropy"
    ATTRIBUTE_NAME = "file_entropy"
    DESCRIPTION = "Calculates the byte entropy of input data."

    def __init__(self):
        """Initializes the entropy hasher."""
        super().__init__()
        self._byte_frequency_counter = collections.Counter()
        self._file_length = 0

    def GetStringDigest(self):
        """Calculates the byte entropy value.

        Byte entropy is a value between 0.0 and 8.0, and is returned as a string
        to match the Plaso analyzer and storage APIs.

        Returns:
          str: byte entropy formatted as a floating point number with 6 decimal
              places calculated over the data blocks passed to Update().
        """
        if self._file_length == 0:
            return "0.000000"

        entropy = 0.0
        for byte_frequency in self._byte_frequency_counter.values():
            byte_probability = byte_frequency / self._file_length
            if byte_probability:
                entropy += -byte_probability * math.log(byte_probability, 2)
        return f"{entropy:.6f}"

    def Update(self, data):
        """Updates the state of the entropy calculator with a new block of data.

        Repeated calls to update are equivalent to one single call with the
        concatenation of the arguments.

        Args:
          data(bytes): block of data with which to update the context of the entropy
              calculator.
        """
        # The call to update() determines the number of occurrences of a byte value
        # within data.
        self._byte_frequency_counter.update(data)
        self._file_length += len(data)


manager.HashersManager.RegisterHasher(EntropyHasher)
