#!/bin/bash

# Define variables
DISK=$1  # Replace with your actual disk identifier (e.g., /dev/sda)
MOUNT_POINT="/mnt/recovery"  # Mount point for data recovery

# Create mount point directory if it doesn't exist
if [ ! -d "$MOUNT_POINT" ]; then
    sudo mkdir -p "$MOUNT_POINT"
fi

echo """

▗▄▄▖ ▗▄▄▄▖▗▄▄▖  ▗▄▖ ▗▄▄▄▖▗▄▄▖ ▗▄▄▄ ▗▄▄▄▖ ▗▄▄▖▗▖ ▗▖
▐▌ ▐▌▐▌   ▐▌ ▐▌▐▌ ▐▌  █  ▐▌ ▐▌▐▌  █  █  ▐▌   ▐▌▗▞▘
▐▛▀▚▖▐▛▀▀▘▐▛▀▘ ▐▛▀▜▌  █  ▐▛▀▚▖▐▌  █  █   ▝▀▚▖▐▛▚▖ 
▐▌ ▐▌▐▙▄▄▖▐▌   ▐▌ ▐▌▗▄█▄▖▐▌ ▐▌▐▙▄▄▀▗▄█▄▖▗▄▄▞▘▐▌ ▐▌
                                                  
                                                  
                                                  

"""

show_help() {
    echo "Usage: $0 <DISK_DEVICE>"
    echo ""
    echo "This script performs a series of checks and repairs on a disk."
    echo "It requires root privileges to run certain commands (use sudo)."
    echo ""
    echo "Arguments:"
    echo "  DISK_DEVICE     The device to check and repair (e.g., /dev/sda)."
    echo ""
    echo "Examples:"
    echo "  $0 /dev/sda   # Check and repair /dev/sda"
    echo ""
    exit 0
}

# Check if help is requested
if [[ "$1" == "--help" ]]; then
    show_help
fi

LOG_DIR="/var/log/disk_health_check"
BADBLOCKS_LOG="${LOG_DIR}/badblocks.log"
SMART_LOG="${LOG_DIR}/smartctl.log"
FSCK_LOG="${LOG_DIR}/fsck.log"
DDRESCUE_LOG="${LOG_DIR}/ddrescue.log"
CLAMAV_LOG="${LOG_DIR}/clamav.log"
REPORT="${LOG_DIR}/disk_health_report.txt"

# Create log directory if it doesn't exist
mkdir -p "$LOG_DIR"

# Check if the disk is unmounted, then unmount if necessary
if mount | grep -q "$DISK"; then
  echo "Unmounting $DISK for checking..."
  sudo umount "$DISK"
fi

echo "Starting disk health and repair script..."

# 1. Check for bad sectors
echo "Running badblocks on $DISK..."
sudo badblocks -v -o "$BADBLOCKS_LOG" "$DISK" || echo "badblocks check failed."

# 2. File System Check (fsck)
echo "Running fsck on $DISK..."
sudo fsck -f -y "$DISK" &> "$FSCK_LOG"

# 3. e2fsck with bad block checking (only for ext filesystems)
echo "Running e2fsck on $DISK..."
sudo e2fsck -c -y "$DISK" &>> "$FSCK_LOG"

# 4. SMART Test
echo "Running SMART test on $DISK..."
sudo smartctl -t long "$DISK"
sleep 10m  # Wait for the SMART long test to complete, adjust timing as needed
sudo smartctl -a "$DISK" &> "$SMART_LOG"

# 5. Attempt to mount the disk
echo "Attempting to mount $DISK to $MOUNT_POINT..."
if sudo mount "$DISK" "$MOUNT_POINT"; then
    echo "Mount successful. Proceeding with data recovery and malware scan."

    # 6. Data Recovery with ddrescue
    echo "Running ddrescue on $DISK..."
    sudo ddrescue -f -n "$DISK" "$MOUNT_POINT/disk_image.img" "$DDRESCUE_LOG"

    # 7. Malware scan with ClamAV
    echo "Running ClamAV scan on $MOUNT_POINT..."
    sudo clamscan --remove --infected --recursive "$MOUNT_POINT" --log="$CLAMAV_LOG"

    # 8. Malware scan with Maldet
    echo "Running maldet scan on $MOUNT_POINT..."
    sudo /usr/local/sbin/maldet -a "$MOUNT_POINT"

    # Unmount after operations
    echo "Unmounting $DISK..."
    sudo umount "$MOUNT_POINT"
else
    echo "Failed to mount $DISK. Skipping mount point operations and continuing with block device checks."
fi

# Generate the report
echo "Generating report..."
{
  echo """


  ┳┓      •  ┳┓• ┓   ┳┓         
  ┣┫┏┓┏┓┏┓┓┏┓┃┃┓┏┃┏  ┣┫┏┓┏┓┏┓┏┓╋
  ┛┗┗ ┣┛┗┻┗┛ ┻┛┗┛┛┗  ┛┗┗ ┣┛┗┛┛ ┗
      ┛                  ┛      

  """
  echo "=== Disk Health Check Report ==="
  echo "Timestamp: $(date)"
  echo ""
  echo "1. Badblocks Test:"
  cat "$BADBLOCKS_LOG"
  echo ""
  echo "2. Filesystem Check (fsck):"
  cat "$FSCK_LOG"
  echo ""
  echo "3. SMART Test Results:"
  cat "$SMART_LOG"
  echo ""
  echo "4. Data Recovery (ddrescue):"
  cat "$DDRESCUE_LOG"
  echo ""
  echo "5. Malware Scan (ClamAV):"
  cat "$CLAMAV_LOG"
  echo ""
  echo "=== End of Report ==="
} > "$REPORT"

echo "Disk health check and repair completed. Report saved to $REPORT."