#!/bin/bash

# Function to print usage information
usage() {
    echo "Usage: $0 <image> <start_tag> <end_tag>"
    echo "Example: $0 tabledevil/sep 230101 230916"
    exit 1
}

# Function to get today's date in YYMMDD format
get_today() {
    date +%y%m%d
}

# Assign arguments to variables or use defaults
image="${1:-tabledevil/sep}"
start_tag="$2"
end_tag="${3:-$(get_today)}"
IMAGE_REPOSITORY=$(echo "$image" | cut -d'/' -f1)
IMAGE_NAME=$(echo "$image" | cut -d'/' -f2)

PATTERN="Threat Found!"

# Validate that start_tag and end_tag are in the correct format
if ! [[ "$start_tag" =~ ^[0-9]{6}$ ]] || ! [[ "$end_tag" =~ ^[0-9]{6}$ ]]; then
    echo "Error: Tags must be in YYMMDD format."
    usage
fi

# Function to get Docker Hub token using PAT
get_token() {
    if [ -z "$DOCKER_USERNAME" ]; then
        read -p "Enter Docker Hub username: " DOCKER_USERNAME
    fi
    if [ -z "$DOCKER_PAT" ]; then
        read -sp "Enter Docker Hub token: " DOCKER_PAT
        echo
    fi
    TOKEN=$(curl -s -H "Content-Type: application/json" -X POST -d '{"username": "'${DOCKER_USERNAME}'", "password": "'${DOCKER_PAT}'"}' https://hub.docker.com/v2/users/login/ | jq -r .token)
    if [ "$TOKEN" == "null" ]; then
        echo "Failed to get token. Please check your credentials."
        exit 1
    fi
}

# Function to get tags for a repository
get_tags() {
    PAGE=1
    PAGE_SIZE=100
    TAGS=()

    while true; do
        if [ -n "$TOKEN" ]; then
            RESPONSE=$(curl -s -H "Authorization: JWT ${TOKEN}" "https://hub.docker.com/v2/repositories/${IMAGE_REPOSITORY}/${IMAGE_NAME}/tags/?page_size=${PAGE_SIZE}&page=${PAGE}")
        else
            RESPONSE=$(curl -s "https://hub.docker.com/v2/repositories/${IMAGE_REPOSITORY}/${IMAGE_NAME}/tags/?page_size=${PAGE_SIZE}&page=${PAGE}")
        fi
        TAGS_PAGE=$(echo $RESPONSE | jq -r '.results[].name')
        TAGS+=($TAGS_PAGE)

        NEXT=$(echo $RESPONSE | jq -r '.next')
        if [ "$NEXT" == "null" ]; then
            break
        fi
        PAGE=$((PAGE + 1))
    done

    echo "${TAGS[@]}"
}

# Function to filter tags based on date range
filter_tags_by_date() {
    local tags=("$@")
    local filtered=()
    local start="$start_tag"
    local end="$end_tag"

    for tag in "${tags[@]}"; do
        if [[ "$tag" =~ ^[0-9]{6}$ ]]; then
            if [[ "$tag" -ge "$start" && "$tag" -le "$end" ]]; then
                filtered+=("$tag")
            fi
        fi
    done

    echo "${filtered[@]}"
}

# Function to check for malware in the given Docker tag
check_malware() {
    local tag=$1
    echo -n "Checking $image:$tag "
    result=$(docker run -it --rm -v "$(pwd):/data:ro" --network=none "$image:$tag" scan)
    if echo "$result" | grep -q "$PATTERN"; then
        echo "Pattern found"
        return 0
    else
        echo "Pattern NOT found"
        return 1
    fi
}

# Function to run the binary search
binary_search() {
    local tags=("$@")
    local low=0
    local high=$(( ${#tags[@]} - 1 ))
    local mid

    # Determine the initial states for low and high
    check_malware "${tags[$low]}"
    local low_result=$?
    
    check_malware "${tags[$high]}"
    local high_result=$?

    # If the results for low and high are the same, there is no switch point in the range
    if [ $low_result -eq $high_result ]; then
        echo "No change in detection within the tag range."
        return
    fi

    # Binary search to find the exact switching point
    while [ $((low + 1)) -lt $high ]; do
        mid=$(((low + high) / 2))
        check_malware "${tags[$mid]}"
        local mid_result=$?

        if [ $mid_result -eq $low_result ]; then
            low=$mid
        else
            high=$mid
        fi
    done

    # Output the tag of the first image that finds the malware switch
    echo "The detection changes between tags: ${tags[$low]} - ${tags[$high]}"
}


# Main script execution

# Try to retrieve all tags without authentication
echo "Retrieving all tags for $image from Docker Hub..."
all_tags=($(get_tags))
echo $all_tags

# If tags retrieval failed, prompt for credentials and retry
if [ ${#all_tags[@]} -eq 0 ]; then
    echo "Failed to retrieve tags without authentication. Trying with credentials..."
    get_token
    all_tags=($(get_tags))
    if [ ${#all_tags[@]} -eq 0 ]; then
        echo "Failed to retrieve tags even with authentication."
        exit 1
    fi
fi

# Filter tags to include only those within the date range
echo "Filtering tags from $start_tag to $end_tag..."
filtered_tags=($(filter_tags_by_date "${all_tags[@]}"))

# Run the binary search on the filtered tags
echo "Running binary search on the filtered tags..."
binary_search "${filtered_tags[@]}"