ssd_health/files/scripts/store_drive.sh

#!/bin/bash
# script for handling adding and updating local drive database

# Function to display usage
usage() {
    echo "Usage: $0 [-i] [-v] [-x] -d /path/to/drive_records.db [-a 'serial,model,flavor,capacity,TBW,smart' OR -u 'serial,TBW,smart'] "
    echo "Options - choose only one of a, u, or i, v and x are optional and not exclusive, and always provide the d"
    echo "  -d /path/to/drive_records.db                   Specify path to DB, required"
    echo "  -a 'serial,model,flavor,capacity,TBW,smart'    Add new drive to sqlite db"
    echo "  -u 'serial,TBW,smart'                          Update drive data in sqlite db"
    echo "  -i                                             Initialize database if not present"
    echo "  -v                                             Output verbose information"
    echo "  -x                                             Output debug information"
    exit 1
}

# init_db subroutine
init_db() {
    if [ "$BE_VERBOSE" == "true" ] ; then
        echo "initializing database"
    fi
    # Check if the file does not exist
    if [ ! -e "$DB_FILE" ]; then
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "No database file, initializing at $DB_FILE"
        fi
        sqlite3 "$DB_FILE" "$CREATE_TABLE"
        chmod 666 $DB_FILE
    else
        if [ "$BE_VERBOSE" == "true" ] ; then
            ls -lah $DB_FILE
            echo "Database file exists, checking tables."
        fi
        DB_PRESENT=$(sqlite3 $DB_FILE .tables | grep drive_records)
        if [ -z "$DB_PRESENT" ]; then
            sqlite3 "$DB_FILE" "$CREATE_TABLE"
            chmod 666 $DB_FILE
        fi
    fi
}

# Define variables
DB_FILE="drive_records.db"
NEEDS_ARGS=false
INIT_ONLY=false
ADD_DRIVE=false
UPDATE_DRIVE=false
DB_PROVIDED=false
VALID_FLAGS=false
BE_VERBOSE=false
OUTPUT_DEBUG=false
CREATE_TABLE="CREATE TABLE drive_records (
    id INTEGER PRIMARY KEY, 
    serial TEXT NOT NULL,
    model TEXT NOT NULL, 
    flavor TEXT NOT NULL, 
    capacity TEXT NOT NULL, 
    TBW TEXT NOT NULL, 
    smart TEXT NOT NULL
);"

# Parse command line options
while getopts ":d:a:u:ivx" opt; do
  case ${opt} in
    v ) # process option v
        echo "Be Verbose"
        BE_VERBOSE=true
        ;;
    x ) # process option x
        echo "Debug Info Enabled"
        OUTPUT_DEBUG=true
        ;;
    d ) # process option d
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "database path provided"
        fi
        DB_PROVIDED=true
        DB_FILE=$OPTARG
        ;;
    a ) # process option a
        if [ "$BE_VERBOSE" == "true" ] ; then   
            echo "add new drive"
        fi
        ADD_DRIVE=true
        NEEDS_ARGS=true
        VALID_FLAGS=true
        DRIVE_DATA=$OPTARG
        ;;
    u ) # process option u
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "update existing drive"
        fi
        UPDATE_DRIVE=true
        NEEDS_ARGS=true
        VALID_FLAGS=true
        DRIVE_DATA=$OPTARG
        ;;
    i ) # process option i
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "initialize database"
        fi
        VALID_FLAGS=true
        init_db
        ;;
    \? ) usage
        ;;
  esac
done
shift $((OPTIND -1))

# Check if all required options are provided
if [ "$DB_PROVIDED" == "false" ] ; then
    if [ "$BE_VERBOSE" == "true" ] ; then
        echo "Database Not Provided"
    fi
    usage
fi
if [ "$VALID_FLAGS" == "false" ] ; then
    if [ "$BE_VERBOSE" == "true" ] ; then
        echo "Invalid Flags"
    fi
    usage
fi
if [ "$NEEDS_ARGS" == "true" ] ; then
    if [ "$BE_VERBOSE" == "true" ] ; then
        echo "NEEDS_ARGS: $NEEDS_ARGS"
    fi
    if [ -z "$DRIVE_DATA" ] ; then
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "Missing Arguments"
        fi
        usage
    fi
fi

# add new drive
if [ "$ADD_DRIVE" == "true" ]; then
    if [ "$BE_VERBOSE" == "true" ] ; then
        echo "Adding new drive"
    fi
    if [ "$OUTPUT_DEBUG" == "true" ] ; then   
        echo "DRIVE_DATA for new drive:"
        echo "$DRIVE_DATA"
    fi
    # Extract the values from the argument string
    IFS=',' read -ra data <<< "$DRIVE_DATA"
    # Check if we have exactly 6 arguments
    if [ ${#data[@]} -ne 6 ]; then
        echo "Exactly 6 arguments are required."
        exit 1
    fi
    # Check if the file does not exist, this should never fail
    if [ ! -e "$DB_FILE" ]; then
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "No database file, exiting"
        exit 2
        fi
    fi
    DRIVE_EXISTS=$(curl -s http://0.0.0.0:5000/check?serial_lookup=${data[0]} | jq .serial_number_exists)
    if [ "$DRIVE_EXISTS" == "false" ]; then
    # Insert the values into the database
        sqlite3 "$DB_FILE" <<EOF
INSERT INTO drive_records (serial, model, flavor, capacity, TBW, smart) 
VALUES ('${data[0]}', '${data[1]}', '${data[2]}', '${data[3]}', '${data[4]}', '${data[5]}');
EOF
    else
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "Drive already exists, skipping"
        fi
    fi
fi

# update existing drive
if [ "$UPDATE_DRIVE" == "true" ]; then
    if [ "$BE_VERBOSE" == "true" ] ; then
        echo "Update existing drive"
    fi
    if [ "$OUTPUT_DEBUG" == "true" ] ; then   
        echo "DRIVE_DATA for new drive:"
        echo "$DRIVE_DATA"
    fi
    # Extract the values from the argument string
    IFS=',' read -ra data <<< "$DRIVE_DATA"
    # Check if we have exactly 3 arguments
    if [ ${#data[@]} -ne 3 ]; then
        echo "Exactly 3 arguments are required."
        exit 1
    fi
    # Check if the file does not exist, this should never fail
    if [ ! -e "$DB_FILE" ]; then
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "No database file, exiting"
        exit 3
        fi
    fi
    DRIVE_EXISTS=$(curl -s http://0.0.0.0:5000/check?serial_lookup=${data[0]} | jq .serial_number_exists)
    if [ "$DRIVE_EXISTS" == "true" ]; then
        # Update the values in the database
        sqlite3 "$DB_FILE" <<EOF
UPDATE drive_records
SET TBW = '${data[1]}', smart = '${data[2]}'
WHERE serial = '${data[0]}';
EOF
    else
        if [ "$BE_VERBOSE" == "true" ] ; then
            echo "Drive does not exist, skipping"
        fi
    fi
fi