entropy.py

import logging
import json
import base64
import numpy as np
import magic
from urllib.parse import urlparse, parse_qs, unquote
from hexdump import hexdump
import binascii
import math
import sp80022suite
from scipy import stats

logger = logging.getLogger(__name__)

ALPHA = 0.01

def _bytes_to_bits(data: bytes) -> bytes:
  """bytes to bitsequence for NIST tests"""
  bits = []
  for byte in data:
    for i in range(7, -1, -1):
      bits.append((byte >> i) & 1)
  return bytes(bits)

def chi2_test(data: bytes) -> float:
  #data = data[int(len(data)/5):] # header data heuristic
  npdata = np.array(list(data), dtype=np.uint8)
  observed_counts = np.histogram(npdata, bins=np.arange(257), range=(0, 256))[0]
  expected_counts = np.full_like(observed_counts, float(len(data) / 256), dtype=np.float64)
  chi2_stat, p_value = stats.chisquare(observed_counts, expected_counts)
  return p_value

def nist_test(data: bytes) -> float:
  """NIST 800-22 statistical tests on raw bytedata

  Returns minimum p-value across all applicable tests.
  Tests selected opportunistically based on minimum data length requirements
  """
  bits = _bytes_to_bits(data)
  n_bits = len(bits)

  p_values = []

  # frequency test
  p_values.append(sp80022suite.frequency(bits))

  # block_frequency (min-block-size assumed to be 64)
  if n_bits >= 64:
    p_values.append(sp80022suite.block_frequency(64, bits))

  # Runs test
  p_values.append(sp80022suite.runs(bits))

  # Cumulative sums test
  p_values.append(sp80022suite.cumulative_sums(bits))

  # approximate entropy: NIST recommends m < floor(log2(n)) - 5
  max_m = int(math.log2(n_bits)) - 5 if n_bits >= 64 else 0
  if max_m >= 2:
    p_values.append(sp80022suite.approximate_entropy(min(max_m, 5), bits))

  # serial test: block_length must be < floor(log2(n))
  max_serial = int(math.log2(n_bits)) if n_bits >= 64 else 0
  if max_serial >= 4:
    p_values.append(sp80022suite.serial(min(max_serial - 1, 16), bits))

  # longest run of ones: needs >= 128 bits
  if n_bits >= 128:
    p_values.append(sp80022suite.longest_run_of_ones(bits))

  # DFT: longer length is better
  if n_bits >= 1000:
    p_values.append(sp80022suite.discrete_fourier_transform(bits))

  if not p_values:
    logger.debug(f"nist_test: data too short ({n_bits} bits), falling back to chi2")
    return chi2_test(data)

  min_p = min(p_values)
  logger.debug(f"nist_test: {len(p_values)} tests, min p-value={min_p:.6f}")
  return min_p

def find_random_bounds(data: bytes, block_size: int = 64, alpha: float = 0.01,
                       test_fn=nist_test) -> tuple[int, int] | None:
  """approximate the largest contiguous block of random bytes in a payload.

  Scans block_size window from tail of payload (assuming that footer data is less common,
  smaller than header data) until sufficiently random tail region is identified.

  binsearch smallest start_idx s.t. nist_test(start_idx, end_idx) > alpha

  Returns (start, end). None if no random block is found.
  """
  length = len(data)
  if length < block_size:
    return None # data too small

  # scan from back to find random tail region
  end = None
  n = 0
  while n + block_size <= length:
    window_start = length - n - block_size
    window = data[window_start : window_start + block_size]
    p = test_fn(window)
    if p > alpha: # Found random tail
      end = window_start + block_size
      break
    n += block_size

  if end is None:
    return None

  # binsearch smallest start_idx s.t. nist_test(start_idx, end_idx) > alpha
  lo = 0
  hi = end - block_size
  while lo < hi:
    mid = (lo + hi) // 2
    p = test_fn(data[mid : end])
    if p > alpha:
      hi = mid  # random extends at least to mid, search left
    else:
      lo = mid + 1  # random block doesn't start this early, search right

  start = lo
  return (start, end)


def decode_any_elem(e):
  if e is None:
    return ""
  if isinstance(e, str):
    datatype = check_entropic(bytes(e, "utf-8"))
  elif isinstance(e, bytes):
    datatype = check_entropic(e)
  elif isinstance(e, list):
    datatype = decode_list(e)
  elif isinstance(e, int):
    return "int"
  elif isinstance(e, float):
    return "float"
  elif isinstance(e, dict):
    datatype = decode_dict(e)
  else:
    raise Exception(f"what is this {e.__class__}")
  return datatype

def decode_list(l):
  for e in l:
    datatype = decode_any_elem(e)
    if "high entropy" in datatype:
      return datatype + " in list"
  return "list"

def decode_dict(d):
  for k, v in d.items():
    if isinstance(k, str):
      k = bytes(k, "ascii")
    datatype = decode_any_elem(v)
    if "high entropy" in datatype:
      return f"key {k.decode('ascii')} has {datatype}"
  return "kv store"

def decode_query_params(queryparams):
  return "query-param " + decode_dict(parse_qs(queryparams))

def try_decode_ascii(data: bytes) -> bytes:
  udata = unquote(data)
  if len(set(udata)) <= 18:
    try:
      unspaced = udata.replace(" ", "").strip()
      hdata = bytes.fromhex(unspaced)
      return "hex-encoded " + check_entropic(hdata)
    except ValueError as e:
      pass
  elif len(set(udata)) <= 66:
    try:
      if "-" in udata or "_" in udata:
        # might be url-safe base64
        udata = udata.replace("-", "+").replace("_", "/")
      unspaced = udata.replace(" ", "").strip()
      bdata = base64.b64decode(unspaced, validate=True)
      return "b64-encoded " + check_entropic(bdata)
    except (binascii.Error, ValueError) as e:
      pass

  if data.count(b"&") + 1 == data.count(b"="):
    return decode_query_params(data)
  return "ASCII"

def check_entropic(data: bytes) -> str:
  datatype = magic.from_buffer(data)
  if datatype == "data" or datatype.startswith("OpenPGP"): # OpenPGP is a common false negative
    block = ""
    if len(data) % 16 == 0:
      block = "128-bit aligned "
    elif len(data) % 8 == 0:
      block = "64-bit aligned "
    p_value = nist_test(data) 
    logger.debug(f"entropy: {p_value}")
    if p_value > 0.01:
      if len(data) < 32:
        return block + "short data"
      return block + "high entropy data"
    bounds = find_random_bounds(data)
    if bounds is not None:
      start, end = bounds
      return block + f"data with {end-start} byte high entropy region [{start:#x}:{end:#x}]"
    return block + "data"
  if datatype.startswith("HTML document"):
    return datatype
  if datatype.startswith("JSON text data"):
    rawdict = json.loads(data)
    if isinstance(rawdict, dict):
      datatype = decode_dict(rawdict)
    else:
      datatype = decode_list(rawdict)
    return "JSON-encoded " + datatype
  if "Image" in datatype:
    return datatype
  if datatype.startswith("ASCII text") or datatype.startswith("OpenPGP"):
    return try_decode_ascii(data)
  return datatype