#!/bin/bash

# Pure Bash implementation of Keccak-256
# This implements the Keccak-f[1600] permutation with r=1088, c=512, output=256 bits

# ---- Constants ----

# Round constants for Iota step
IOTA_RC=(
    0x0000000000000001 0x0000000000008082 0x800000000000808A 0x8000000080008000
    0x000000000000808B 0x0000000080000001 0x8000000080008081 0x8000000000008009
    0x000000000000008A 0x0000000000000088 0x0000000080008009 0x000000008000000A
    0x000000008000808B 0x800000000000008B 0x8000000000008089 0x8000000000008003
    0x8000000000008002 0x8000000000000080 0x000000000000800A 0x800000008000000A
    0x8000000080008081 0x8000000000008080 0x0000000080000001 0x8000000080008008
)

# Rotation offsets
RHO_OFFSETS=(
    0  1  62 28 27
    36 44 6  55 20
    3  10 43 25 39
    41 45 15 21 8
    18 2  61 56 14
)

# ---- Helper functions ----

# Split 64-bit value into high and low 32-bit parts
# This helps bash handle large integers better
to_u32_pair() {
    local value=$1
    local low=$(( value & 0xFFFFFFFF ))
    local high=$(( (value >> 32) & 0xFFFFFFFF ))
    echo "$high $low"
}

# Combine high and low 32-bit parts into a 64-bit value
from_u32_pair() {
    local high=$1
    local low=$2
    echo $(( (high << 32) | low ))
}

# Rotate a 64-bit value left by n bits - uses 32-bit parts for better bash handling
rotate_left64() {
    local val=$1
    local shift=$2
    
    # Handle zero rotation
    if (( shift == 0 )); then
        echo $val
        return
    fi
    
    # Split into high and low 32-bit parts
    local val_parts=( $(to_u32_pair $val) )
    local high=${val_parts[0]}
    local low=${val_parts[1]}
    
    # Rotation within 32-bit range
    if (( shift < 32 )); then
        local high_shift=$(( high << shift ))
        local high_overflow=$(( high >> (32 - shift) ))
        local low_shift=$(( low << shift ))
        local low_overflow=$(( low >> (32 - shift) ))
        
        # Combine with overflow
        local new_high=$(( (high_shift & 0xFFFFFFFF) | low_overflow ))
        local new_low=$(( (low_shift & 0xFFFFFFFF) | high_overflow ))
        
        echo $(from_u32_pair $new_high $new_low)
    else
        # Handle large rotations
        local adjusted_shift=$(( shift - 32 ))
        local high_shift=$(( low << adjusted_shift ))
        local high_overflow=$(( low >> (32 - adjusted_shift) ))
        local low_shift=$(( high << adjusted_shift ))
        local low_overflow=$(( high >> (32 - adjusted_shift) ))
        
        local new_high=$(( (high_shift & 0xFFFFFFFF) | low_overflow ))
        local new_low=$(( (low_shift & 0xFFFFFFFF) | high_overflow ))
        
        echo $(from_u32_pair $new_high $new_low)
    fi
}

# ---- State management functions ----

# Initialize Keccak state
init_state() {
    # Initialize all 25 lanes to zero
    for i in {0..24}; do
        state[$i]=0
    done
}

# Get state index from x,y coordinates (convert to 1D array index)
index() {
    local x=$1
    local y=$2
    echo $(( x + 5 * y ))
}

# Get rotation offset for position x,y
get_rho_offset() {
    local x=$1
    local y=$2
    echo ${RHO_OFFSETS[$(( x + 5 * y ))]}
}

# ---- Keccak-f permutation and its steps ----

# Implement the theta step of Keccak-f
step_theta() {
    local -a C D
    
    # Compute column parity in C[x]
    for x in {0..4}; do
        C[$x]=0
        for y in {0..4}; do
            local idx=$(index $x $y)
            C[$x]=$(( C[$x] ^ state[$idx] ))
        done
    done
    
    # Compute D[x] = C[x-1] ^ rot(C[x+1], 1)
    for x in {0..4}; do
        local x_minus_1=$(( (x + 4) % 5 ))
        local x_plus_1=$(( (x + 1) % 5 ))
        local rot_val=$(rotate_left64 ${C[$x_plus_1]} 1)
        D[$x]=$(( C[$x_minus_1] ^ rot_val ))
    done
    
    # Update state with D[x]
    for x in {0..4}; do
        for y in {0..4}; do
            local idx=$(index $x $y)
            state[$idx]=$(( state[$idx] ^ D[$x] ))
        done
    done
}

# Implement the rho and pi steps combined
step_rho_pi() {
    local -a temp_state
    
    # Copy current state
    for i in {0..24}; do
        temp_state[$i]=${state[$i]}
    done
    
    # Apply rho and pi transformations
    for y in {0..4}; do
        for x in {0..4}; do
            local src_idx=$(index $x $y)
            local rot_val=$(get_rho_offset $x $y)
            
            # Calculate pi transformation
            local new_x=$y
            local new_y=$(( (2*x + 3*y) % 5 ))
            local dst_idx=$(index $new_x $new_y)
            
            # Apply rotation (rho) and lane permutation (pi)
            state[$dst_idx]=$(rotate_left64 ${temp_state[$src_idx]} $rot_val)
        done
    done
}

# Implement the chi step of Keccak-f
step_chi() {
    local -a temp_state
    
    # Copy current state
    for i in {0..24}; do
        temp_state[$i]=${state[$i]}
    done
    
    # Apply chi transformation
    for y in {0..4}; do
        for x in {0..4}; do
            local idx=$(index $x $y)
            local x1=$(( (x + 1) % 5 ))
            local x2=$(( (x + 2) % 5 ))
            local idx1=$(index $x1 $y)
            local idx2=$(index $x2 $y)
            
            # Calculate ~temp_state[x1,y] & temp_state[x2,y]
            local not_x1=$(( (~temp_state[$idx1]) & 0xFFFFFFFFFFFFFFFF ))
            local and_result=$(( not_x1 & temp_state[$idx2] ))
            
            # XOR with original state
            state[$idx]=$(( temp_state[$idx] ^ and_result ))
        done
    done
}

# Implement the iota step of Keccak-f
step_iota() {
    local round=$1
    state[0]=$(( state[0] ^ ${IOTA_RC[$round]} ))
}

# Perform a full Keccak-f[1600] permutation
keccak_f() {
    for round in {0..23}; do
        step_theta
        step_rho_pi
        step_chi
        step_iota $round
    done
}

# ---- Input processing functions ----

# Convert a hex string to an array of bytes
hex_to_bytes() {
    local hex="$1"
    local -a bytes
    
    # Handle empty input
    if [[ -z "$hex" ]]; then
        echo ""
        return
    fi
    
    # Ensure even number of characters
    if (( ${#hex} % 2 != 0 )); then
        hex="0$hex"
    fi
    
    # Convert pairs of hex digits to bytes
    for (( i=0; i<${#hex}; i+=2 )); do
        local byte="0x${hex:$i:2}"
        bytes+=( "$byte" )
    done
    
    echo "${bytes[@]}"
}

# Pad the input according to Keccak padding rule (0x01, ..., 0x80)
apply_padding() {
    local -a input=("$@")
    local len=${#input[@]}
    local rate_bytes=136   # 1088 bits / 8
    
    # Calculate padding length
    local mod_len=$(( len % rate_bytes ))
    local pad_len=$(( rate_bytes - mod_len ))
    
    # Ensure minimum padding of 2 bytes
    if (( pad_len < 2 )); then
        pad_len=$(( pad_len + rate_bytes ))
    fi
    
    # Add input bytes
    local -a result=("${input[@]}")
    
    # Add 0x01 byte
    result+=( "0x01" )
    
    # Add zeros if needed
    for (( i=1; i<pad_len-1; i++ )); do
        result+=( "0x00" )
    done
    
    # Add final 0x80 byte
    result+=( "0x80" )
    
    echo "${result[@]}"
}

# ---- Keccak sponge construction functions ----

# Absorb data into the Keccak state
absorb_data() {
    local -a data=("$@")
    local data_len=${#data[@]}
    local rate_bytes=136   # 1088 bits / 8
    
    # Initialize state
    init_state
    
    # Process data in blocks
    for (( block_start=0; block_start<data_len; block_start+=rate_bytes )); do
        local block_end=$(( block_start + rate_bytes ))
        if (( block_end > data_len )); then
            block_end=$data_len
        fi
        
        # Absorb block data into state
        for (( i=block_start; i<block_end; i++ )); do
            local rel_pos=$(( i - block_start ))
            local lane_idx=$(( rel_pos / 8 ))
            local byte_pos=$(( rel_pos % 8 ))
            local byte=${data[$i]}
            
            # Convert byte position to bit shift
            local bit_shift=$(( byte_pos * 8 ))
            
            # Calculate state indices
            local x=$(( lane_idx % 5 ))
            local y=$(( lane_idx / 5 ))
            local state_idx=$(index $x $y)
            
            # XOR byte into appropriate position in state
            state[$state_idx]=$(( state[$state_idx] ^ ( byte << bit_shift ) ))
        done
        
        # Apply Keccak permutation after each block
        if (( block_start + rate_bytes <= data_len )); then
            keccak_f
        fi
    done

    # Final permutation if not yet applied
    if (( data_len % rate_bytes != 0 )); then
        keccak_f
    fi
}

# Squeeze output from the Keccak state
squeeze_output() {
    local output_bytes=32  # 256 bits (for Keccak-256)
    local result=""
    
    # Extract bytes from state (output is limited to 32 bytes for Keccak-256)
    for (( i=0; i<output_bytes; i++ )); do
        local lane_idx=$(( i / 8 ))
        local byte_pos=$(( i % 8 ))
        local x=$(( lane_idx % 5 ))
        local y=$(( lane_idx / 5 ))
        local state_idx=$(index $x $y)
        
        # Extract byte from lane
        local bit_shift=$(( byte_pos * 8 ))
        local byte=$(( (state[$state_idx] >> bit_shift) & 0xFF ))
        
        # Format as two hex digits
        printf -v hex_byte "%02x" $byte
        result+="$hex_byte"
    done
    
    echo "$result"
}

# ---- Main functions ----

# Compute Keccak-256 hash of a hex string
keccak256() {
    local input_hex="$1"
    
    # Convert input to byte array
    local -a input_bytes
    input_bytes=( $(hex_to_bytes "$input_hex") )
    
    # Apply padding
    local -a padded_data
    padded_data=( $(apply_padding "${input_bytes[@]}") )
    
    # Process data through sponge construction
    absorb_data "${padded_data[@]}"
    
    # Get hash result
    local hash
    hash=$(squeeze_output)
    
    echo "$hash"
}

# ---- Test functions ----

# Test Keccak-256 with known test vectors
test_keccak256() {
    echo "Running Keccak-256 tests..."
    
    # Test vector 1: Empty string
    local expected1="c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470"
    local result1=$(keccak256 "")
    
    # Test vector 2: "hello"
    local expected2="1c8aff950685c2ed4bc3174f3472287b56d9517b9c948127319a09a7a36deac8"
    local result2=$(keccak256 "68656c6c6f")
    
    # Test vector 3: "The quick brown fox jumps over the lazy dog"
    local quick_fox="54686520717569636b2062726f776e20666f78206a756d7073206f76657220746865206c617a7920646f67"
    local expected3="4d741b6f1eb29cb2a9b9911c82f56fa8d73b04959d3d9d222895df6c0b28aa15"
    local result3=$(keccak256 "$quick_fox")
    
    # Display results
    echo "Test 1 (empty string):"
    echo "Expected: $expected1"
    echo "Got:      $result1"
    echo "Status:   $([ "$result1" == "$expected1" ] && echo "PASS" || echo "FAIL")"
    echo
    
    echo "Test 2 (\"hello\"):"
    echo "Expected: $expected2"
    echo "Got:      $result2"
    echo "Status:   $([ "$result2" == "$expected2" ] && echo "PASS" || echo "FAIL")"
    echo
    
    echo "Test 3 (\"The quick brown fox...\"):"
    echo "Expected: $expected3"
    echo "Got:      $result3"
    echo "Status:   $([ "$result3" == "$expected3" ] && echo "PASS" || echo "FAIL")"
}

# ---- Command-line interface ----

if [[ "$1" == "-x" ]]; then
    # Hex input
    keccak256 "$2"
elif [[ "$1" == "-s" ]]; then
    # String input (convert to hex first)
    hex=$(echo -n "$2" | xxd -p | tr -d '\n')
    keccak256 "$hex"
elif [[ "$1" == "--test" ]]; then
    # Run tests
    test_keccak256
else
    # Show usage
    echo "Usage: $0 [-x hex_string] [-s ascii_string] [--test]"
    echo "  -x: Compute Keccak-256 of hex string (e.g. 68656c6c6f for 'hello')"
    echo "  -s: Compute Keccak-256 of ASCII string (e.g. 'hello')"
    echo "  --test: Run test vectors to verify implementation"
fi 