import subprocess
import requests

# Function to get all available drives
def get_all_drives():
    try:
        # Run lsblk to get all available drives
        result = subprocess.run(['lsblk', '-d', '-n', '-o', 'NAME'], capture_output=True, text=True)
        drives = [f"/dev/{line.strip()}" for line in result.stdout.splitlines() if line.strip()]
        return drives
    except Exception as e:
        print(f"Error getting drives: {e}")
        return []

# Function to check SMART status of a drive
def check_smart_status(drive):
    try:
        # Run smartctl command for the specified drive
        result = subprocess.run(['sudo', 'smartctl', drive, '-a'], capture_output=True, text=True)
        output = result.stdout
        
        # Look for the "SMART overall-health self-assessment test result" line
        health_line = None
        for line in output.splitlines():
            if "SMART overall-health self-assessment test result" in line:
                health_line = line
                break
        
        # Check if the health test result is "PASSED"
        if health_line and "PASSED" in health_line:
            return True, health_line
        else:
            return False, health_line if health_line else "SMART health result not found"
    
    except Exception as e:
        return False, f"Error checking SMART status: {e}"

# Function to send a notification using ntfy
def send_ntfy_notification(title, body, priority="default", tags=None):
    url = "https://ntfy.hrasci.eu/rpi5"  # Replace with your ntfy topic URL
    headers = {
        "Title": title,
        "Priority": priority,
        "Tags": ",".join(tags) if tags else "",
    }
    
    try:
        response = requests.post(url, headers=headers, data=body)
        if response.status_code == 200:
            print("Notification sent successfully.")
        else:
            print(f"Failed to send notification: {response.status_code}")
    except Exception as e:
        print(f"Error sending notification: {e}")

def main():
    # Get all available drives
    drives = get_all_drives()
    
    if not drives:
        print("No drives found.")
        return

    unhealthy_drives = []
    healthy_drives = []

    for drive in drives:
        # Check the SMART status for each drive
        smart_ok, smart_output = check_smart_status(drive)
        
        if smart_ok:
            healthy_drives.append(f"{drive}: PASSED")
        else:
            unhealthy_drives.append(f"{drive}: {smart_output}")

    # Send notifications based on the drive status
    if unhealthy_drives:
        # If there are unhealthy drives, send urgent notifications
        for drive_info in unhealthy_drives:
            title = f"SMART Alert for {drive_info.split(':')[0]}"
            tags = ["warning", "skull"]
            send_ntfy_notification(title, drive_info, priority="urgent", tags=tags)
    else:
        # If all drives are healthy, send a low priority notification
        healthy_message = "\n".join(healthy_drives)
        send_ntfy_notification("All Drives Healthy", healthy_message, priority="low", tags=["ok", "check"])

if __name__ == "__main__":
    main()