#!/usr/bin/env bash

##############################################################################
# dotest - A Randomness Quality Testing Script
# ---------------------------------------------------------------------------
#
# OVERVIEW:
# This script tests the quality of a random byte generator program by:
# 1. Generating a stream of random bytes using the provided entropy source.
# 2. Testing the generated random data using various randomness test
#    utilities:
#    - `dieharder` (if installed): A comprehensive battery of statistical
#      tests.
#    - `rngtest` (if installed): A FIPS 140-2 compliance tester.
#    - `ent` (if installed): A simple entropy analysis tool.
# 3. Summarizing the results with clear, color-coded output for easy
#    interpretation.
#
# ---------------------------------------------------------------------------
#
# SCRIPT COMPONENTS:
#
# 1. **ENTROPY SOURCE**:
#    The random byte stream is generated using an external program
#    (`entropy_test`) which outputs raw random bytes to standard output.
#    The script captures this output and writes it to a file
#    (`random_data.bin`) for further analysis.
#
# 2. **TEST UTILITIES**:
#    - `ent`: Measures basic entropy statistics, such as Shannon entropy,
#             Chi-square tests, and arithmetic mean of the byte
#             distribution.
#    - `dieharder`: Executes a comprehensive suite of statistical tests on
#                   the generated random data. Selected tests include:
#                   * Birthday Spacings (-d 0)
#                   * Overlapping Permutations (-d 1)
#                   * Rank of 32x32 Binary Matrices (-d 2)
#    - `rngtest`: A utility that performs FIPS 140-2 tests, which include:
#                 * Monobit, Poker, Runs, Long run, Continuous run
#
# 3. **QUALITY RATIO FOR RNGTEST**:
#    The script calculates a "Success-to-Failure Ratio" from `rngtest`
#    results:
#       Quality Ratio = Successes / (Failures + 1)
#    - Adding 1 to the denominator avoids division by zero.
#    - A high ratio indicates better randomness quality.
#    - A warning is triggered if the ratio falls below the predefined
#      minimum threshold (`min_ratio`, default = 100).
#
# 4. **HIGHLIGHTING**:
#    - Results are color-coded for better readability:
#      * Green: Indicates success or excellent quality.
#      * Orange: Indicates minor issues, within acceptable limits.
#      * Red: Indicates significant issues or unacceptable quality.
#
# ---------------------------------------------------------------------------
#
# USAGE:
# 1. Ensure the `entropy_test` program is compiled. This requires linking
#    against the `ici` library and providing the path to the main ION
#    headers. From the directory containing this script, use:
#
#        gcc -o entropy_test entropy_test.c -I../../include -lici
#
# 2. Run the script:
#       ./dotest
#
# 3. Results will be displayed on the console with a summary.
#
# ---------------------------------------------------------------------------
#
# LIMITATIONS:
# - If `ent`, `dieharder`, and `rngtest` are all uninstalled, the script
#   will exit with an error message.
# - By default, the script generates 10 MB of random data. Adjust the size
#   in the script's configuration section if needed.
#
# ---------------------------------------------------------------------------
#
# CUSTOMIZATION:
# - The `min_ratio` for `rngtest` can be modified in the `run_rngtest`
#   function to suit specific application requirements.
# - Additional `dieharder` tests can be added or removed by modifying the
#   `run_dieharder` function.
#
# ---------------------------------------------------------------------------
#
# Modification History:
#  MM/DD/YY  AUTHOR         DESCRIPTION
#  --------  ------------   ---------------------------------------------
#  12/11/24  S. DeBaun      Initial implementation
#  06/05/25  S. DeBaun      Update replaces grep -P and adds missing ent
#                           function.
#
##############################################################################


ENTROPY_TEST_PROGRAM="entropy_test"
OUTPUT_FILE="random_data.bin"
LOG_FILE="randomness_test.log"
BYTE_SIZE=${BYTE_SIZE:-10M}

# Color codes for test result highlighting
# Use simple string assignment. The `printf` command will interpret the octal escape sequences.
GREEN='\033[1;32m'
ORANGE='\033[1;33m'
RED='\033[1;31m'
RESET='\033[0m'

# Overall status flag
OVERALL_SUCCESS=1

#======================================================================
# Helper Functions
#======================================================================
check_binary() {
    if [ ! -f "$ENTROPY_TEST_PROGRAM" ]; then
        printf "\n${RED}Error: Required binary '%s' not found.${RESET}\n" "$ENTROPY_TEST_PROGRAM"
        printf "\nPlease build it with the following command (run from this directory):\n"
        printf "${GREEN}gcc -o %s entropy_test.c -I../../include -lici${RESET}\n\n" "$ENTROPY_TEST_PROGRAM"
        exit 1
    fi
}

check_utility() {
    if command -v "$1" &>/dev/null; then
        printf "%s found.\n" "$1"
        return 0
    else
        printf "%s not found.\n" "$1"
        return 1
    fi
}

#======================================================================
# Test Runners
#======================================================================
generate_random_bytes() {
    printf "\nGenerating random bytes (%s)...\n" "$BYTE_SIZE"
    "./$ENTROPY_TEST_PROGRAM" | head -c "$BYTE_SIZE" >"$OUTPUT_FILE"
    if [ $? -ne 0 ] || [ ! -s "$OUTPUT_FILE" ]; then
        printf "${RED}Error: Failed to generate random bytes.${RESET}\n"
        exit 1
    fi
    printf "Random bytes saved to %s\n" "$OUTPUT_FILE"
}

run_ent() {
    printf "\nRunning 'ent' statistical analysis...\n"
    ent "$OUTPUT_FILE"
}


run_dieharder() {
    printf "\nRunning 'dieharder' tests...\n"
    # A simplified dieharder runner:
    dieharder -a -f "$OUTPUT_FILE"
}

run_rngtest() {
    printf "\nRunning 'rngtest'...\n"

    local successes=0
    local failures=0
    local min_ratio=100

    # Capture full rngtest output
    local rngtest_output
    rngtest_output=$(rngtest < "$OUTPUT_FILE" 2>&1)

    # Display full rngtest output
    printf -- "-----------------------\n"
    printf "%s\n" "$rngtest_output"

    # CORRECTED and portable parsing logic
    # Use printf "%s" to avoid adding an extra newline before piping
    successes=$(printf "%s" "$rngtest_output" | awk -F': ' '/successes/ {print $3}')
    failures=$(printf "%s" "$rngtest_output" | awk -F': ' '/failures/ {print $3}')

    # Default to 0 if parsing somehow fails
    successes=${successes:-0}
    failures=${failures:-0}

    # Calculate quality ratio
    local quality_ratio=$((successes / (failures + 1)))

    printf "\n${GREEN}RNGTEST RESULTS SUMMARY${RESET}\n"
    printf -- "-----------------------\n"
    printf "FIPS 140-2 successes: ${GREEN}%s${RESET}\n" "${successes}"
    printf "FIPS 140-2 failures:  ${RED}%s${RESET}\n" "${failures}"
    printf "Success-to-Failure Ratio: ${ORANGE}%s${RESET}\n" "${quality_ratio}"

    if [ "$failures" -eq 0 ]; then
        printf "${GREEN}All FIPS 140-2 tests passed.${RESET}\n"
    elif [ "$quality_ratio" -lt "$min_ratio" ]; then
        printf "${RED}Warning: Random data quality is below acceptable limits.${RESET}\n"
        OVERALL_SUCCESS=0
    else
        printf "${ORANGE}Minor issues detected, but within acceptable limits.${RESET}\n"
    fi
}

#======================================================================
# Cleanup
#======================================================================
cleanup_files() {
    if [ $OVERALL_SUCCESS -eq 1 ]; then
        printf "\n${GREEN}All tests passed. Cleaning up temporary files...${RESET}\n"
        rm -f "$OUTPUT_FILE"
    else
        printf "\n${ORANGE}Some tests failed or were weak. Keeping %s for review.${RESET}\n" "$OUTPUT_FILE"
    fi
}

#======================================================================
# Main Script Logic
#======================================================================
# Check for our test binary
check_binary

printf "Checking for recommended utilities...\n"
# CORRECTED: Initialize variables to 0 before the checks
ENT_INSTALLED=0
DIEHARDER_INSTALLED=0
RNGTEST_INSTALLED=0

if check_utility "ent"; then
    ENT_INSTALLED=1
fi
if check_utility "dieharder"; then
    DIEHARDER_INSTALLED=1
fi
if check_utility "rngtest"; then
    RNGTEST_INSTALLED=1
fi

if [ $ENT_INSTALLED -eq 0 ] && [ $DIEHARDER_INSTALLED -eq 0 ] && [ $RNGTEST_INSTALLED -eq 0 ]; then
    printf "No randomness testing utilities found. Please install one or more of: ent, dieharder, rngtest\n"
    exit 1
fi

generate_random_bytes

# Run the tests
if [ $RNGTEST_INSTALLED -eq 1 ]; then
    run_rngtest
fi

if [ $DIEHARDER_INSTALLED -eq 1 ]; then
    run_dieharder
fi

if [ $ENT_INSTALLED -eq 1 ]; then
    run_ent
fi


# Perform cleanup based on overall test results
cleanup_files
